"""Resolve a free-text segmentation target (e.g. "liver") to a canonical
TotalSegmentator label name + integer id, using this repo's own
dataset/TotalSegmentator_v2.json (CT) / dataset/TotalSegmentator_MRI.json
(MRI) as the source of truth for the name<->id mapping.

Resolution order: exact match -> normalized match -> unique substring match
-> LLM disambiguation against the full label list.
"""

import json
import re

from . import config
from .query_parser import extract_json_from_output, get_oai_client_for_model


def _labels_path_for_modality(modality):
    return config.CT_LABELS_PATH if modality.upper().startswith("CT") else config.MRI_LABELS_PATH


def _load_labels(modality):
    path = _labels_path_for_modality(modality)
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return data["labels"]  # {name: int_id}, includes "background": 0


def _normalize(name):
    return re.sub(r"[\s\-]+", "_", name.strip().lower())


def _llm_disambiguate(target, labels, model):
    names = [n for n in labels.keys() if n != "background"]
    prompt = (
        "You are a medical imaging expert. Given the segmentation target "
        f'"{target}", pick the single closest matching name from this list '
        f"of available structures: {json.dumps(names, ensure_ascii=False)}\n"
        "Respond ONLY in this exact format:\n"
        'Final Answer: {"label": "chosen_name"}'
    )
    client = get_oai_client_for_model(model)
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
    )
    text = response.choices[0].message.content or ""
    parsed = extract_json_from_output(text, expected_type="dict")
    label_name = parsed.get("label") if parsed else None
    if not label_name or label_name not in labels:
        raise ValueError(
            f"「{target}」に該当する構造が見つかりませんでした（LLM応答: {text!r}）。"
        )
    return label_name


def resolve_target(target, modality, model):
    """Returns {"name": str, "id": int}. Raises ValueError if unresolved."""
    labels = _load_labels(modality)
    norm_target = _normalize(target)

    if norm_target in labels:
        return {"name": norm_target, "id": labels[norm_target]}

    candidates = [name for name in labels if norm_target in name or name in norm_target]
    candidates = [name for name in candidates if name != "background"]
    if len(candidates) == 1:
        return {"name": candidates[0], "id": labels[candidates[0]]}

    label_name = _llm_disambiguate(target, labels, model)
    return {"name": label_name, "id": labels[label_name]}