import json
import re
from typing import Dict, List, Optional

from .ollama_client import chat_vision
from .prompts import build_closedset_prompt


def _extract_json_obj(text: str):
    """尽量从返回文本中提取出合法 JSON（支持代码块/前后说明）。"""
    if not isinstance(text, str):
        return None
    s = text.strip()
    # 去掉 Markdown 代码块包裹
    if s.startswith("```"):
        s = re.sub(r"^```[a-zA-Z]*\n|\n```$", "", s).strip()
    # 直接解析
    try:
        obj = json.loads(s)
        if isinstance(obj, dict):
            return obj
    except Exception:
        pass
    # 提取第一个花括号对象
    m = re.search(r"\{[\s\S]*?\}", s)
    if m:
        frag = m.group(0)
        try:
            obj = json.loads(frag)
            if isinstance(obj, dict):
                return obj
        except Exception:
            pass
    return None


def classify_image(
    image_path: str,
    labels: List[str],
    fewshot_hints: Optional[Dict[str, str]] = None,
    ingredient_only: bool = False,
) -> Dict:
    if not labels:
        print("[VLM] Empty labels provided, return Unknown")
        return {"label": "Unknown", "confidence": 0.0}

    print(
        f"[VLM] Start classify image={image_path} labels={len(labels)} "
        f"ingredient_only={ingredient_only} fewshot={bool(fewshot_hints)}"
    )
    prompt = build_closedset_prompt(labels, fewshot_hints, ingredient_only)
    _preview_prompt = (prompt[:120]).replace("\n", " ")
    print(f"[VLM] Prompt length={len(prompt)} preview={_preview_prompt}...")

    print("[VLM] Calling chat_vision...")
    text = chat_vision(prompt, [image_path], temperature=0.1)
    _preview_text = (str(text)[:200]).replace("\n", " ")
    print(f"[VLM] Received text length={len(str(text))} preview={_preview_text}...")

    # 解析 JSON（更健壮）
    obj = _extract_json_obj(str(text))
    if obj and "label" in obj:
        label = obj.get("label", "")
        # 若输出不在封闭集，尝试包含匹配；否则 Unknown
        if label not in labels:
            for lb in labels:
                if lb in str(label):
                    label = lb
                    break
            else:
                label = "Unknown"
        conf = obj.get("confidence", 0.0)
        try:
            conf = float(conf)
        except Exception:
            conf = 0.0
        result = {"label": label, "confidence": conf}
        print(f"[VLM] Parsed JSON result={result}")
        return result

    print("[VLM] JSON parse failed or missing 'label', enter fallback")

    # 回退：字符串包含匹配（无法确定置信度时返回 0.0）
    for lb in labels:
        if lb in str(text):
            print(f"[VLM] Fallback match label={lb}")
            return {"label": lb, "confidence": 0.0}

    print("[VLM] No match, return Unknown")
    return {"label": "Unknown", "confidence": 0.0}