"""
dt_gemini.py
------------
Gemini AI 用戶端、JSON 解析工具、題目生成、AI 批改回饋。

依賴：無 Dash 元件依賴（可獨立測試）

公開 API：
    _gem_call(txt, model)                    → str
    _gem_json(model, system, schema, payload) → dict
    _ai_generate_similar_quizzes_as_md(...)  → str
    _parse_ai_md_to_quiz_nodes(...)          → List[dict]
    _ai_feedback_fill_gemini(...)            → dict
    _ai_feedback_choice_gemini(...)          → str
    _ai_feedback_multi_gemini(...)           → dict
"""
from __future__ import annotations
import os, json, re, traceback
from pathlib import Path
from dotenv import load_dotenv

# ---------------------------------------------------------------------------
# Gemini client 初始化（單一來源）
# ---------------------------------------------------------------------------
import google.genai as genai
from google.genai import types as genai_types

DOTENV_PATH = Path(__file__).resolve().parents[2] / ".env"  # layout/renderer/ → 上兩層 = 專案根目錄
load_dotenv(dotenv_path=DOTENV_PATH, override=True)

_JSON_RE = re.compile(r"\{.*\}", re.S)

def _get_api_key() -> str:
    raw = (os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY") or "").strip()
    return raw.strip('"').strip("'")

_API_KEY = _get_api_key()
print("[dotenv] path =", DOTENV_PATH, "| exists =", DOTENV_PATH.exists())
print("[dotenv] GEMINI_API_KEY set:", bool(_API_KEY))
_gem_client = genai.Client(api_key=_API_KEY) if _API_KEY else None

# ---------------------------------------------------------------------------
# 基礎工具
# ---------------------------------------------------------------------------

def _norm_model(m: str) -> str:
    m = (m or "").strip()
    return m if m.startswith("models/") else (f"models/{m}" if m else "models/gemini-2.5-flash")


def _s(x) -> str:
    """safe strip：None → ''，其他型別 → str().strip()"""
    if x is None: return ""
    if isinstance(x, str): return x.strip()
    return str(x).strip()


def _first_json(text: str):
    """從文字中抽出第一段 {...} 並解析為 dict；失敗回 None。"""
    m = _JSON_RE.search(text or "")
    if not m: return None
    try:
        return json.loads(m.group(0))
    except Exception:
        return None


def _ai_log_fail(tag: str, exc: Exception, raw_text: str | None = None) -> None:
    print("\n" + "=" * 90, flush=True)
    print(f"[AI FAIL] {tag}", flush=True)
    print(f"[AI FAIL] {type(exc).__name__}: {exc!r}", flush=True)
    if raw_text is not None:
        raw_text = raw_text or ""
        print("[AI FAIL] --- raw response (head 2000 chars) ---", flush=True)
        print(raw_text[:2000], flush=True)
        if len(raw_text) > 2000:
            print("[AI FAIL] --- raw response (tail 400 chars) ---", flush=True)
            print(raw_text[-400:], flush=True)
    traceback.print_exc()
    print("=" * 90 + "\n", flush=True)

# ---------------------------------------------------------------------------
# Gemini 呼叫
# ---------------------------------------------------------------------------

def _gem_call(txt: str, model: str | None = None) -> str:
    """最基本的文字呼叫，無 schema，回傳純文字。"""
    global _gem_client
    if not _gem_client:
        print("[gemini] no API key; skip call")
        return ""
    try:
        r = _gem_client.models.generate_content(
            model=_norm_model(model or "models/gemini-2.5-flash"),
            contents=[{"role": "user", "parts": [{"text": txt}]}],
        )
        return (getattr(r, "text", "") or "").strip()
    except Exception as e:
        print("[gemini] call failed:", e)
        return ""


def _repair_json_newlines(s: str) -> str:
    """修復 JSON 字串值內的真換行，避免 json.loads 爆掉。"""
    if not s: return s
    out, in_str, esc = [], False, False
    for ch in s:
        if not in_str:
            out.append(ch)
            if ch == '"': in_str = True; esc = False
            continue
        if esc:
            out.append(ch); esc = False; continue
        if ch == '\\':
            out.append(ch); esc = True; continue
        if ch == '"':
            out.append(ch); in_str = False; continue
        if ch == '\n': out.append('\\n')
        elif ch == '\r': out.append('\\r')
        elif ch == '\t': out.append('\\t')
        else: out.append(ch)
    return "".join(out)


def _loads_json_robust(text: str):
    """先直接 loads，失敗後修復換行再試。"""
    m = _JSON_RE.search(text or "")
    s = m.group(0) if m else (text or "")
    s = s.strip()
    try:
        return json.loads(s)
    except json.JSONDecodeError:
        pass
    try:
        return json.loads(_repair_json_newlines(s))
    except Exception:
        return None


def _gem_json(
    model: str,
    system: str,
    schema: dict,
    payload: dict,
    temperature: float = 0.2,
    max_tokens: int = 256,
) -> dict:
    """JSON schema 模式呼叫，自動 fallback 至無 schema 模式。"""
    if not _gem_client or not _API_KEY:
        print("[gemini] NO_KEY", flush=True)
        return {"_error": "NO_KEY"}

    DEBUG_AI = str(os.getenv("DEBUG_AI", "0")).lower() in ("1", "true", "yes", "on")
    if DEBUG_AI:
        print("\n" + "=" * 90, flush=True)
        print(f"[gemini][REQ] model = {model}", flush=True)
        print(f"[gemini][REQ] system =\n{system}", flush=True)
        print("=" * 90 + "\n", flush=True)

    cfg_schema = genai_types.GenerateContentConfig(
        system_instruction=system,
        response_mime_type="application/json",
        response_schema=schema,
        temperature=temperature,
        max_output_tokens=max_tokens,
    )

    # --- 第一次嘗試：JSON schema mode ---
    try:
        resp = _gem_client.models.generate_content(
            model=_norm_model(model),
            contents=json.dumps(payload, ensure_ascii=False),
            config=cfg_schema,
        )
        text = getattr(resp, "text", "") or ""
        if not text and getattr(resp, "candidates", None):
            parts = resp.candidates[0].content.parts
            text = getattr(parts[0], "text", "") if parts else ""
        try:
            return json.loads(text or "{}")
        except json.JSONDecodeError as je:
            _ai_log_fail("JSON decode failed (schema mode)", je, text)
            return {"_error": "BAD_JSON", "_why": str(je), "_raw_head": (text or "")[:2000]}
    except Exception as e1:
        _ai_log_fail("Gemini call failed (schema mode)", e1)

    # --- 第二次嘗試：無 schema，強制輸出 JSON ---
    cfg_plain = genai_types.GenerateContentConfig(
        system_instruction=system + " Output ONLY JSON with the requested fields.",
        temperature=temperature,
        max_output_tokens=max_tokens,
    )
    try:
        resp2 = _gem_client.models.generate_content(
            model=_norm_model(model),
            contents=json.dumps(payload, ensure_ascii=False),
            config=cfg_plain,
        )
        text2 = getattr(resp2, "text", "") or ""
        if not text2 and getattr(resp2, "candidates", None):
            parts2 = resp2.candidates[0].content.parts
            text2 = getattr(parts2[0], "text", "") if parts2 else ""
        obj = _loads_json_robust(text2 or "")
        if obj is not None:
            return obj
        _ai_log_fail("JSON decode failed (robust repair also failed)", Exception("BAD_JSON"), text2)
        return {"_error": "BAD_JSON", "_why": "robust repair failed", "_raw_head": (text2 or "")[:2000]}
    except Exception as e2:
        _ai_log_fail("Gemini call failed (no-schema fallback)", e2)
        return {"_error": "CALL_FAIL", "_why": str(e2)}

# ---------------------------------------------------------------------------
# AI 題目生成
# ---------------------------------------------------------------------------

_OPT_LINE_RE = re.compile(r'^\s*([a-dA-D])\s*[\.]\s*(.+?)\s*$')
_RE_AI_CHOICE = re.compile(r'<!--\s*choice(?P<attrs>[^>]*)-->(?P<body>.*?)<!--\s*choice_end\s*-->', re.S | re.I)
_RE_AI_FILL   = re.compile(r'<!--\s*fill(?P<attrs>[^>]*)-->(?P<body>.*?)<!--\s*fill_end\s*-->', re.S | re.I)
_Q_ATTR = re.compile(r'([A-Za-z_][A-Za-z0-9_-]*)\s*=\s*"([^"]*)"'
                     r'|([A-Za-z_][A-Za-z0-9_-]*)\s*=\s*([^\s">]+)')

def _attrs_to_dict_local(s: str) -> dict:
    out: dict = {}
    if not s: return out
    for m in _Q_ATTR.finditer(s):
        k = (m.group(1) or m.group(3)).strip()
        v = (m.group(2) or m.group(4)).strip()
        out[k] = v
    return out


def _parse_options_md(s: str) -> list[dict]:
    out = []
    for line in (s or "").splitlines():
        m = _OPT_LINE_RE.match(line)
        if m:
            out.append({"key": m.group(1).lower(), "text": m.group(2)})
    return out


# DEFAULT_NORMALIZE 由 dt_render_utils 提供，此處需要直接用
DEFAULT_NORMALIZE = ["sym", "numeric", "nospace"]


def _ai_generate_similar_quizzes_as_md(
    prompt_text: str,
    qtype: str,
    model: str = "gemini-2.5-flash",
) -> str:
    """呼叫 Gemini，生成 3 題類似題（回傳 Markdown 格式字串）。"""
    system = (
        "你是數學教材出題助教。請根據『原題』生成 3 題同概念同難度練習題。\n"
        "⚠ 只輸出『題庫註解格式』，不要輸出 JSON、不要 code fence、不要多餘解釋。\n\n"
        "格式（單選）：\n"
        "題目（可含 $...$，允許多行）\n"
        '<!--choice ans="c" card="0" placeholder="請輸入答案"-->\n'
        "a. 選項\nb. 選項\nc. 選項\nd. 選項\n<!--choice_end-->\n\n"
        "格式（填空）：\n"
        "題目（可含 $...$，允許多行）\n"
        '<!--fill ans="答案1|答案2" card="0" placeholder="請輸入答案"-->\n'
        "<!--fill_end-->\n\n"
        "規則：\n"
        "1) 一共輸出 3 題。\n2) 選項請用 a./b./c./d.。\n"
        "3) 內容避免半形雙引號，用「」或單引號。\n"
        "4) 數學式只要一組 $ 字號。\n"
        "5) 除法要用 \\frac{}{}。\n"
        "6) 數學式要用 $ 包起來，禁止用空格。"
    )
    ask = system + "\n\n[原題]\n" + (prompt_text or "") + f"\n\n[題型]\n{qtype}"
    print("\n===== [moreq] ask to gemini BEGIN =====", flush=True)
    print(ask, flush=True)
    print("===== [moreq] ask to gemini END =====\n", flush=True)
    return (_gem_call(ask, model=model) or "").strip()


def _prompt_md_to_runs(md: str) -> list[dict]:
    """把含數學式的 prompt MD 字串轉成 prompt_runs（供 render_inline_runs 用）。"""
    import re as _re
    MATH_FRAG = _re.compile(r'(\\\(.+?\\\)|\\\[.+?\\\]|\$\$.+?\$\$|\$(?!\$).+?(?<!\$)\$)', _re.S)

    def _strip(s: str):
        t = s.strip()
        if t.startswith(r'\(') and t.endswith(r'\)'): return t[2:-2], False
        if t.startswith(r'\[') and t.endswith(r'\]'): return t[2:-2], True
        if t.startswith('$$') and t.endswith('$$'):   return t[2:-2], True
        if t.startswith('$') and t.endswith('$'):     return t[1:-1], False
        return t, False

    md = md or ""
    runs: list[dict] = []

    def push_text(text: str):
        for i, part in enumerate(text.split("\n")):
            r: dict = {"kind": "text", "text": part}
            if i != len(text.split("\n")) - 1:
                r["newline"] = "true"
            runs.append(r)

    pos = 0
    for m in MATH_FRAG.finditer(md):
        if m.start() > pos:
            push_text(md[pos:m.start()])
        latex, is_block = _strip(m.group(0))
        runs.append({"kind": "math_block" if is_block else "math", "latex": latex})
        pos = m.end()
    if pos < len(md):
        push_text(md[pos:])
    return runs


def _parse_ai_md_to_quiz_nodes(
    md_text: str,
    default_qtype: str = "choice",
) -> list[dict]:
    """把 AI 生成的 Markdown 題目字串解析成 quiz node 清單。"""
    md_text = (md_text or "").replace("＄", "$").replace("（", "(").replace("）", ")")

    nodes: list[dict] = []
    blocks = []
    for m in _RE_AI_CHOICE.finditer(md_text):
        blocks.append(("choice", m.start(), m.end(), m))
    for m in _RE_AI_FILL.finditer(md_text):
        blocks.append(("fill", m.start(), m.end(), m))
    blocks.sort(key=lambda x: x[1])

    last_end = 0
    for kind, s, e, m in blocks:
        stem = (md_text[last_end:s] or "").strip()
        stem = re.sub(r"^\s*題目\s*[:：]\s*", "", stem)
        attrs = _attrs_to_dict_local(m.group("attrs") or "")
        body  = (m.group("body") or "").strip()

        if kind == "choice":
            lines = [ln.rstrip() for ln in body.splitlines()]
            prompt_lines, opt_lines, hit_opt = [], [], False
            for ln in lines:
                if _OPT_LINE_RE.match(ln.strip()): hit_opt = True
                (opt_lines if hit_opt else prompt_lines).append(ln)
            prompt_md = ("\n".join(x for x in prompt_lines if x.strip()) or stem).strip().replace("＄", "$")
            opts = []
            for ln in opt_lines:
                mm = _OPT_LINE_RE.match(ln.strip())
                if mm:
                    opts.append({"key": mm.group(1).lower(), "text": mm.group(2).strip()})
            ans = (attrs.get("ans") or "").strip().lower().rstrip(".")
            idxmap = {"1": "a", "2": "b", "3": "c", "4": "d"}
            if ans in idxmap: ans = idxmap[ans]
            nodes.append({
                "type": "quiz", "qtype": "choice", "id": "",
                "prompt_md": prompt_md, "prompt_runs": _prompt_md_to_runs(prompt_md),
                "options": opts, "answer": ans, "card": "1", "ai": "1",
            })
        else:
            ans_raw = attrs.get("ans") or ""
            answers = [a.strip() for a in ans_raw.split("|") if a.strip()]
            prompt_md = (body or stem).strip().replace("＄", "$")
            nodes.append({
                "type": "quiz", "qtype": "fill", "id": "",
                "prompt_md": prompt_md, "prompt_runs": _prompt_md_to_runs(prompt_md),
                "answers": answers, "normalize": DEFAULT_NORMALIZE,
                "card": "1", "ai": "1",
            })
        last_end = e
    return nodes

# ---------------------------------------------------------------------------
# AI 批改回饋
# ---------------------------------------------------------------------------

def _ai_feedback_multi_gemini(
    prompt_text: str,
    selected_keys: list[str],
    options: list[dict],
    answer_keys: list[str],
    rubric: str = "",
    model: str = "gemini-2.5-flash",
    lang: str = "zh-TW",
) -> dict:
    opt_map = {str(o.get("key")).lower(): o.get("text", "") for o in (options or [])}
    ask = (
        "你是數學助教，請用繁體中文回覆。\n"
        "請判斷學生多選題是否正確，並提供具體引導。\n"
        "必要時可用行內 LaTeX（$...$）。\n"
        '只輸出這段 JSON（不得多字）：{"is_correct":true/false,"hint":"...","explain":"..."}\n\n'
        f"題目：{prompt_text}\n"
        f"選項：{json.dumps(opt_map, ensure_ascii=False)}\n"
        f"學生勾選：{json.dumps([str(k).lower() for k in (selected_keys or [])], ensure_ascii=False)}\n"
        f"正確答案：{json.dumps([str(k).lower() for k in (answer_keys or [])], ensure_ascii=False)}\n"
        + (f"評分規則：{rubric}\n" if rubric else "")
    )
    obj = _first_json(_gem_call(ask, model=model)) or {}
    hint = _s(obj.get("hint")); explain = _s(obj.get("explain"))
    if not (hint or explain): return {}
    return {"hint": hint, "explain": explain}


def _ai_feedback_choice_gemini(
    prompt_text: str,
    selected_key: str,
    options: list[dict],
    answer_key: str,
    rubric: str = "",
    model: str = "gemini-2.5-flash",
    lang: str = "zh-TW",
) -> str:
    opt_map = {str(o.get("key")).lower(): o.get("text", "") for o in (options or [])}
    ask = (
        "你是數學助教，請用繁體中文回覆，提供引導式提示，不要直接公布正確答案。\n"
        "必要時可用行內 LaTeX（$...$）。\n"
        '只輸出這段 JSON（不得多字）：{"is_correct":true/false,"hint":"...","explain":"..."}\n\n'
        f"題目：{prompt_text}\n"
        f"選項：{json.dumps(opt_map, ensure_ascii=False)}\n"
        f"學生選的鍵：{selected_key}\n"
        f"學生選的內容：{opt_map.get(selected_key, '')}\n"
        f"正解鍵：{answer_key}\n"
        f"正解內容：{opt_map.get(answer_key, '')}\n"
        + (f"評分規則：{rubric}\n" if rubric else "")
    )
    txt = _gem_call(ask, model=model)
    obj = _first_json(txt)
    print(ask)
    return _s((obj or {}).get("hint"))


def _ai_feedback_fill_gemini(
    prompt_text: str,
    user_input: str,
    answers: list[str],
    rubric: str = "",
    model: str = "gemini-2.5-flash",
    lang: str = "zh-TW",
) -> dict:
    ask = (
        "你是數學助教，請用繁體中文回覆。\n"
        "判斷學生填空題是否正確，提供具體診斷與引導修正。\n"
        "必要時可用行內 LaTeX（$...$）。\n"
        '只輸出這段 JSON（不得多字）：{"is_correct":true/false,"hint":"...","explain":"..."}\n\n'
        f"題目：{prompt_text}\n"
        f"學生作答：{user_input}\n"
        f"可接受答案：{json.dumps(answers, ensure_ascii=False)}\n"
        + (f"評分規則：{rubric}\n" if rubric else "")
    )
    txt = _gem_call(ask, model=model)
    obj = _first_json(txt) or {}
    print(ask)
    return {
        "is_correct": bool((obj or {}).get("is_correct")),
        "hint":    _s((obj or {}).get("hint")),
        "explain": _s((obj or {}).get("explain")),
    }