"""
ChemGraph Loop — a guarded, natural-language API around the real ChemGraph agent.

A visitor asks a chemistry question in plain language ("what's the IR spectrum of
water?"). An LLM intent router identifies the target molecule (by any name /
synonym / formula) + the task; the real ChemGraph LLM agent then plans the
workflow and tool calls for that intent, runs (with its deterministic workflow
guard), and we return the full agentic trace + a 3D structure + the task-specific
result so the site can render the loop.

Two tiers, by physics cost on a free CPU:
  • energy / dipole  → single-point, ~10-25s → run LIVE.
  • vibrations·IR / thermochemistry → finite-difference Hessian, ~1-3 min →
    served from PRECOMPUTED real agent runs (flagged cached:true).

Public-safe guards: molecule + task allow-list, per-call timeout, hourly rate
limit. OPENAI_API_KEY comes from a Space secret.
"""
import asyncio
import glob
import json
import os
import time

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from cg_extract import extract

# --- molecule allow-list + synonyms/formulae (only a fallback; the LLM router
#     below is the primary, agentic intent path) ------------------------------
MOLECULES = {
    "water": ["water", "h2o", "dihydrogen monoxide", "dihydrogen oxide", "oxidane", "aqua", "hydrogen oxide"],
    "methane": ["methane", "ch4", "marsh gas", "carbane", "natural gas", "methyl hydride"],
    "ammonia": ["ammonia", "nh3", "azane", "hydrogen nitride", "spirit of hartshorn"],
    "methanol": ["methanol", "ch3oh", "ch4o", "methyl alcohol", "wood alcohol", "carbinol", "meoh"],
    "ethanol": ["ethanol", "c2h5oh", "ch3ch2oh", "c2h6o", "ethyl alcohol", "grain alcohol", "drinking alcohol", "etoh"],
    "carbon dioxide": ["carbon dioxide", "co2", "carbonic anhydride", "carbonic acid gas", "dry ice"],
    "benzene": ["benzene", "c6h6", "benzol", "cyclohexatriene"],
    "acetic acid": ["acetic acid", "ch3cooh", "ch3co2h", "c2h4o2", "ethanoic acid", "glacial acetic acid", "vinegar acid"],
    "formaldehyde": ["formaldehyde", "ch2o", "hcho", "methanal", "formalin", "methylene oxide"],
    "hydrogen peroxide": ["hydrogen peroxide", "h2o2", "dihydrogen dioxide", "hydrogen dioxide", "peroxide"],
}
CALC_LABEL = {"emt": "EMT", "tblite": "TBLite"}

# unicode sub/superscripts → ascii digits, so "CO₂" / "H²O" resolve like "CO2"
_DIGITS = {"₀": "0", "₁": "1", "₂": "2", "₃": "3", "₄": "4", "₅": "5", "₆": "6", "₇": "7", "₈": "8", "₉": "9",
           "⁰": "0", "¹": "1", "²": "2", "³": "3", "⁴": "4", "⁵": "5", "⁶": "6", "⁷": "7", "⁸": "8", "⁹": "9"}


def normalize(text: str) -> str:
    return "".join(_DIGITS.get(c, c) for c in (text or "")).lower()

TASKS = ("energy", "dipole", "ir", "thermo")
LIVE_TASKS = {"energy", "dipole"}
CACHED_TASKS = {"ir", "thermo"}

MODEL = os.environ.get("CHEMGRAPH_MODEL", "gpt-4o-mini")
INTENT_MODEL = os.environ.get("CHEMGRAPH_INTENT_MODEL", MODEL)
RUN_TIMEOUT = int(os.environ.get("CHEMGRAPH_TIMEOUT", "95"))
RATE_MAX = int(os.environ.get("CHEMGRAPH_RATE_MAX", "40"))

# --- load precomputed heavy runs (real agent, cached) -------------------------
_PRECOMPUTED = {}
_HERE = os.path.dirname(os.path.abspath(__file__))
for fp in glob.glob(os.path.join(_HERE, "precomputed", "*.json")):
    try:
        base = os.path.basename(fp)[:-5]          # e.g. carbon_dioxide__ir
        slug, task = base.rsplit("__", 1)
        with open(fp, "r", encoding="utf-8") as f:
            _PRECOMPUTED[(slug, task)] = json.load(f)
    except Exception:
        pass
def _avail(task: str) -> list:
    return sorted({slug.replace("_", " ") for (slug, t) in _PRECOMPUTED if t == task})


SPECTRA_MOLECULES = sorted({slug.replace("_", " ") for (slug, _t) in _PRECOMPUTED})
TASK_PHRASE = {"energy": "energy", "dipole": "dipole moment",
               "ir": "IR spectrum", "thermo": "thermochemistry"}

app = FastAPI(title="ChemGraph Loop")
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://sciencesloop.com", "https://www.sciencesloop.com",
        "http://localhost:4321", "http://127.0.0.1:4321",
    ],
    allow_origin_regex=r"https://[a-z0-9-]+\.vercel\.app",
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)

_hits: dict[str, list[float]] = {}


def _rate_ok(ip: str) -> bool:
    now = time.time()
    w = _hits.setdefault(ip, [])
    w[:] = [t for t in w if now - t < 3600]
    if len(w) >= RATE_MAX:
        return False
    w.append(now)
    return True


# --- PRIMARY: agentic LLM intent router --------------------------------------
# An LLM reads the plain-language question and identifies the target molecule
# (by any name/synonym/formula) + the task. The ChemGraph agent then plans the
# actual workflow and tool calls for that intent. This is intentionally not a
# rule/keyword engine — the model handles the open-ended naming space.
_client = None


def _openai():
    global _client
    if _client is None:
        from openai import OpenAI
        _client = OpenAI()
    return _client


_INTENT_SYS = (
    "You are the intent router for a computational-chemistry agent that is limited "
    "to a fixed set of small molecules and four tasks. Read the user's natural-language "
    "question, identify the target molecule and the task, and reply with ONLY a compact "
    "JSON object. Do not add commentary."
)


def llm_parse_query(text: str) -> dict:
    mols = ", ".join(sorted(MOLECULES.keys()))
    user = (
        f"Allowed molecules (canonical names): {mols}.\n"
        "Map the user's molecule to one of these ONLY if it is the SAME compound, by any "
        "name — common name, trivial/trade name, IUPAC name, or chemical formula. Examples: "
        "'H2O' / 'dihydrogen monoxide' / 'aqua' -> water; 'CO2' / 'dry ice' -> carbon dioxide; "
        "'EtOH' / 'ethyl alcohol' / 'grain alcohol' -> ethanol; 'NH3' / 'azane' -> ammonia; "
        "'benzol' -> benzene.\n"
        "CRITICAL: a substituted or derivative molecule is a DIFFERENT compound — return "
        "null, do NOT map it to the parent. e.g. dimethoxybenzene / toluene / nitrobenzene / "
        "phenol / aniline are NOT benzene; acetaldehyde is NOT formaldehyde; propanol is NOT "
        "ethanol. Never match just because the name contains an allowed molecule's name. If "
        "the molecule is not EXACTLY one of the allowed ones, use null.\n"
        "Tasks: 'energy' = single-point energy; 'dipole' = dipole moment; 'ir' = "
        "vibrational frequencies / IR spectrum; 'thermo' = thermochemistry (enthalpy, "
        "entropy, Gibbs free energy). Pick the closest task; default to 'energy' if none "
        "is implied.\n"
        "Calculator: 'emt' only if the user explicitly asks for EMT and the task is "
        "energy; otherwise 'tblite'.\n"
        "When molecule is null, also give: 'nearest' = the single closest allowed molecule "
        "if the user's is a close relative/derivative (e.g. dimethoxybenzene -> benzene, "
        "1-propanol -> ethanol, acetone -> null if nothing close), else null; and 'note' = "
        "ONE short, friendly sentence naming the user's molecule and why it's outside this "
        "small-molecule demo (e.g. \"Caffeine is too large for this small-molecule demo.\").\n"
        f'User question: "{text}"\n'
        'Reply as JSON: {"molecule": <canonical name or null>, "task": '
        '"energy|dipole|ir|thermo", "calculator": "emt|tblite", '
        '"nearest": <allowed name or null>, "note": <string or null>}'
    )
    resp = _openai().chat.completions.create(
        model=INTENT_MODEL,
        temperature=0,
        response_format={"type": "json_object"},
        messages=[{"role": "system", "content": _INTENT_SYS},
                  {"role": "user", "content": user}],
    )
    data = json.loads(resp.choices[0].message.content or "{}")
    mol = data.get("molecule")
    mol = mol.strip().lower() if isinstance(mol, str) else ""
    task = data.get("task") or "energy"
    task = task.strip().lower() if isinstance(task, str) else "energy"
    if mol not in MOLECULES:
        nearest = data.get("nearest")
        nearest = nearest.strip().lower() if isinstance(nearest, str) else None
        return {"error": "no_molecule",
                "note": data.get("note") if isinstance(data.get("note"), str) else None,
                "nearest": nearest if nearest in MOLECULES else None,
                "task": task if task in TASKS else "energy"}
    if task not in TASKS:
        task = "energy"
    # calculator isn't "intent": TBLite (real QM) by default; EMT only for an
    # energy query where the user explicitly named it. Deterministic, not model-guessed.
    calc = "emt" if (task == "energy" and "emt" in normalize(text)) else "tblite"
    return {"molecule": mol, "task": task, "calculator": calc}


def keyword_fallback(text: str) -> dict:
    """Non-LLM fallback used ONLY if the intent model call fails. Token-level
    alias match (no regex) so e.g. 'ch4' does not match the token 'ch4o'."""
    t = normalize(text)
    toks = set(t.replace(",", " ").replace("?", " ").replace("(", " ").replace(")", " ").split())
    molecule = None
    for canonical, aliases in MOLECULES.items():
        for a in aliases:
            hit = (a in t) if " " in a else (a in toks)
            if hit:
                molecule = canonical
                break
        if molecule:
            break
    if molecule is None:
        return {"error": "no_molecule"}
    task = "energy"
    for name, kws in (("thermo", ("thermo", "enthalpy", "entropy", "gibbs")),
                      ("ir", ("ir", "infrared", "vibration", "vibrational", "frequency", "frequencies", "spectrum", "spectra")),
                      ("dipole", ("dipole", "polarity"))):
        if toks & set(kws):
            task = name
            break
    calc = "emt" if ("emt" in toks and task == "energy") else "tblite"
    return {"molecule": molecule, "task": task, "calculator": calc}


def canonical_query(molecule: str, task: str, calc: str) -> str:
    label = CALC_LABEL[calc]
    if task == "energy":
        return f"Compute the single point energy of {molecule} using the {label} calculator."
    if task == "dipole":
        return f"Compute the dipole moment of {molecule} using the {label} calculator."
    if task == "ir":
        return f"Compute the IR spectrum of {molecule} using the {label} calculator."
    if task == "thermo":
        return f"Compute the thermochemistry (enthalpy, entropy, Gibbs free energy) of {molecule} using the {label} calculator at 298.15 K."
    return f"Compute the single point energy of {molecule} using the {label} calculator."


@app.get("/")
def health():
    return {
        "ok": True, "service": "chemgraph-loop", "model": MODEL,
        "molecules": sorted(MOLECULES.keys()),
        "tasks": ["energy", "dipole", "ir", "thermo"],
        "live_tasks": sorted(LIVE_TASKS),
        "ir_molecules": _avail("ir"),
        "thermo_molecules": _avail("thermo"),
    }


@app.post("/run")
async def run(req: Request):
    ip = (req.headers.get("x-forwarded-for") or (req.client.host if req.client else "?") or "?").split(",")[0].strip()
    if not _rate_ok(ip):
        return JSONResponse({"error": "Too many requests — try again later."}, status_code=429)

    try:
        body = await req.json()
    except Exception:
        body = {}

    # Accept NL {query}, or legacy {molecule, calculator}.
    if body.get("query"):
        user_text = str(body["query"])
        # PRIMARY: LLM intent router; keyword fallback only if the model call fails.
        try:
            parsed = llm_parse_query(user_text)
        except Exception:
            parsed = keyword_fallback(user_text)
    else:
        mol = str(body.get("molecule", "water")).strip().lower()
        calc = str(body.get("calculator", "emt")).strip().lower()
        parsed = {"molecule": mol if mol in MOLECULES else None,
                  "task": "energy", "calculator": calc if calc in CALC_LABEL else "emt"}
        if parsed["molecule"] is None:
            parsed = {"error": "no_molecule"}
        user_text = None

    # ---- CLARIFICATION NODE: molecule not in the demo's small-molecule set ----
    if parsed.get("error") == "no_molecule":
        note = parsed.get("note") or "That doesn't look like one of the small molecules in this demo."
        nearest = parsed.get("nearest")
        want = parsed.get("task", "energy")
        suggestions = []
        if nearest:
            note += f" Did you mean {nearest}?"
            if want in CACHED_TASKS and nearest in _avail(want):
                suggestions.append({"label": f"{TASK_PHRASE[want]} of {nearest}",
                                    "query": f"{TASK_PHRASE[want]} of {nearest}"})
            suggestions += [{"label": f"energy of {nearest}", "query": f"energy of {nearest}"},
                            {"label": f"dipole of {nearest}", "query": f"dipole of {nearest}"}]
        else:
            suggestions = [{"label": "IR spectrum of water", "query": "IR spectrum of water"},
                           {"label": "dipole of ammonia", "query": "dipole of ammonia"},
                           {"label": "energy of benzene", "query": "energy of benzene"}]
        return JSONResponse({"clarify": True, "message": note,
                             "molecules": sorted(MOLECULES.keys()), "suggestions": suggestions})

    molecule, task, calc = parsed["molecule"], parsed["task"], parsed["calculator"]
    calc_label = CALC_LABEL[calc]
    query = canonical_query(molecule, task, calc)
    display_query = user_text or query

    # ---- cached heavy tasks (vibrations/IR, thermochemistry) ----
    if task in CACHED_TASKS:
        slug = molecule.replace(" ", "_")
        payload = _PRECOMPUTED.get((slug, task))
        if payload is None:
            # CLARIFICATION NODE: molecule is supported, but this heavy task wasn't
            # precomputed for it — offer its live options + where the task IS available.
            avail = _avail(task)
            msg = (f"I don't have a precomputed {TASK_PHRASE[task]} for {molecule} — those need a "
                   f"slow vibrational (Hessian) run, so they're prepared ahead of time"
                   + (f" for {', '.join(avail)}" if avail else "") + ". "
                   f"But I can run {molecule}'s energy or dipole live right now.")
            suggestions = [{"label": f"energy of {molecule}", "query": f"energy of {molecule}"},
                           {"label": f"dipole of {molecule}", "query": f"dipole of {molecule}"}]
            if avail:
                suggestions.append({"label": f"{TASK_PHRASE[task]} of {avail[0]}",
                                    "query": f"{TASK_PHRASE[task]} of {avail[0]}"})
            return JSONResponse({"clarify": True, "message": msg, "suggestions": suggestions})
        out = dict(payload)
        out["display_query"] = display_query
        return JSONResponse(out)

    # ---- live single-point tasks (energy, dipole) ----
    try:
        from chemgraph.agent.llm_agent import ChemGraph
        agent = ChemGraph(model_name=MODEL, return_option="state", recursion_limit=20)
        state = await asyncio.wait_for(agent.run(query), timeout=RUN_TIMEOUT)
    except asyncio.TimeoutError:
        return JSONResponse({"error": "The workflow took too long — please retry."}, status_code=504)
    except Exception as e:  # noqa: BLE001
        return JSONResponse({"error": f"Agent error: {str(e)[:200]}"}, status_code=500)

    out = extract(state, query, molecule, calc_label, task, cached=False)
    out["display_query"] = display_query
    return JSONResponse(out)