"""Natural-language query parsing: free text -> {"modality": ..., "target": ...}.

NOTE on why this doesn't import from agents/ or utils/:
`utils/chat_utils.py` does `from agents import *`, which (via agents/__init__.py)
eagerly imports all ~10 agent pattern modules and the full pyautogen stack just
to reuse a couple of small, self-contained pieces. To keep this webapp
decoupled from that research-code import chain (today it's harmless, but it's
incidental, not structural -- see the implementation plan), the prompt text
and JSON-extraction helper are copied here verbatim from
agents/c2f_tree_search.py (lines ~35-91 at time of writing), and the OpenAI
client lookup is reimplemented directly instead of importing
utils.chat_utils.get_OAI_client_for_model.
"""

import json
import logging
import re

from autogen import config_list_from_json
from openai import OpenAI

from . import config

# Copied verbatim from agents/c2f_tree_search.py (SYSTEM_PROMPT_QUERY_PARSER).
SYSTEM_PROMPT_QUERY_PARSER = """
# Role:
You are an AI assistant specialized in parsing queries related to medical image segmentation. Your task is to split the query into its imaging modality and segmentation target with as much specificity as possible.

## Skills:
1. Query Parsing:
- Identify and extract the imaging modality, including specific subtypes (e.g., CT contrast, CT non-contrast, MR ADC, MR FLAIR, etc.) from the query.
- If specific subtypes are not mentioned, extract only the main modality (e.g., CT, MR).
- Identify and extract the segmentation target (e.g., brain, liver, tumor, etc.).

## Constraints:
1. Response Format
- MUST use the following JSON format for the Final Answer:
Final Answer: {"modality": "modality_name", "target": "segmentation_target"}
Here are some examples:
{"modality": "CT", "target": "liver"}
{"modality": "MR", "target": "spleen"}
{"modality": "MR ADC", "target": "brain tumor"}
"""


def extract_json_from_output(llm_output, expected_type="dict"):
    """Adapted from agents/c2f_tree_search.py's extract_json_from_output.

    One deliberate change from the original: the default expected_type is
    "dict" here (the original repo function defaults to "list", which is a
    footgun for this exact prompt -- SYSTEM_PROMPT_QUERY_PARSER above always
    answers with a JSON *object*, so relying on the original default would
    silently fail to match every single response).
    """
    if expected_type == "list":
        pattern = r"Final Answer:\s*(\[[\s\S]*?\])"
    elif expected_type == "dict":
        pattern = r"Final Answer:\s*(\{[\s\S]*?\})"
    else:
        raise ValueError("expected_type must be 'list' or 'dict'")

    match = re.search(pattern, llm_output)
    if not match:
        logging.error("No 'Final Answer' found in LLM output: %r", llm_output)
        return None

    json_str = match.group(1)
    json_str = re.sub(r"//.*", "", json_str)
    json_str = re.sub(r"/\*[\s\S]*?\*/", "", json_str)
    try:
        return json.loads(json_str)
    except json.JSONDecodeError:
        logging.info("Initial JSON parse failed, retrying with quote fix...")
        json_str_fixed = json_str.replace("'", '"')
        try:
            return json.loads(json_str_fixed)
        except json.JSONDecodeError:
            logging.error("Failed to parse JSON from LLM output: %r", json_str)
            return None


def get_oai_client_for_model(model, config_path=None):
    """Equivalent to utils/chat_utils.get_OAI_client_for_model, reimplemented
    locally to avoid importing utils.chat_utils (see module docstring)."""
    path = str(config_path or config.OAI_CONFIG_PATH)
    cfg = config_list_from_json(env_or_file=path, filter_dict={"model": [model, "default"]})[0]
    return OpenAI(api_key=cfg["api_key"], base_url=cfg["base_url"])


def parse_query(query, model=None):
    """Parse a free-text segmentation request into {"modality": str, "target": str}.

    Raises ValueError if the LLM response couldn't be parsed or is missing
    required fields.
    """
    model = model or config.MODEL_NAME
    client = get_oai_client_for_model(model)

    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT_QUERY_PARSER},
            {"role": "user", "content": query},
        ],
    )
    text = response.choices[0].message.content or ""
    parsed = extract_json_from_output(text, expected_type="dict")

    if not parsed:
        raise ValueError(f"クエリを解析できませんでした（LLM応答: {text!r}）")

    modality = parsed.get("modality")
    target = parsed.get("target") or parsed.get("targets")
    if not modality or not target:
        raise ValueError(f"クエリからmodality/targetを抽出できませんでした: {parsed!r}")

    return {"modality": str(modality), "target": str(target)}