"""FastAPI backend for the Needle 3 structured data lab.

Runs Cactus Compute's Needle 3 model locally via the `cactus` package
(pip: cactus-compute). The engine and the 35 MB needle3.cact weights are
cached by the package; this server warms the cache in a background thread
on startup, then exposes:

  GET  /api/status   -> {"state": "loading"|"ready"|"error", ...}
  GET  /api/tools    -> the built-in demo tools (for the frontend)
  POST /api/extract  -> one-shot structured extraction (cactus.extract)
  POST /api/agent    -> tool-calling run over built-in tools (Cactus.run)
  POST /api/embed    -> text embedding vector (3072 dims)

Serves the static frontend from ./static at /.

Notes:
- cactus.extract only accepts Pydantic models as schemas, so /api/extract
  builds one dynamically from the request's field list (all fields are
  Optional: ungrounded values come back as null instead of failing).
- One Cactus instance == one conversation. /api/agent creates a fresh
  instance per request so requests don't leak context into each other;
  the warm instance is reused only for embeddings.
"""

import asyncio
import threading
import time
import typing
from pathlib import Path

import cactus
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field, create_model

ROOT = Path(__file__).resolve().parent
STATIC_DIR = ROOT / "static"
WEIGHTS = ROOT / "needle3.cact"

app = FastAPI(title="Needle 3 Structured Data Lab")

# --- Engine state -----------------------------------------------------------

_agent = None
_state = {"state": "loading", "message": "Engine not loaded yet.", "loaded_at": None, "load_seconds": None}
_state_lock = threading.Lock()


def _warm() -> None:
    global _agent
    t0 = time.monotonic()
    try:
        _agent = cactus.Cactus(weights=str(WEIGHTS))
        with _state_lock:
            _state.update(
                state="ready",
                message="Engine ready.",
                loaded_at=time.strftime("%Y-%m-%d %H:%M:%S"),
                load_seconds=round(time.monotonic() - t0, 1),
            )
    except Exception as exc:  # noqa: BLE001 - surface any load failure in the UI
        with _state_lock:
            _state.update(state="error", message=f"Engine load failed: {exc}")


@app.on_event("startup")
async def _startup() -> None:
    asyncio.get_event_loop().run_in_executor(None, _warm)


def _require_agent():
    if _agent is None:
        with _state_lock:
            state, message = _state["state"], _state["message"]
        if state == "loading":
            raise HTTPException(409, "Engine is still loading — check /api/status")
        raise HTTPException(503, message)
    return _agent


# --- Built-in demo tools (the model card / docs examples) -------------------


@cactus.tool
def get_weather(city: str):
    "Get the current weather for a city."
    return {"city": city, "temp_c": 27, "sky": "clear"}


@cactus.tool
def set_thermostat(temperature: int, mode: typing.Literal["heat", "cool", "auto"] = "auto"):
    """Set the thermostat.

    Args:
        temperature: target temperature in Celsius
        mode: heating strategy to use
    """
    return {"temperature": temperature, "mode": mode}


@cactus.tool
def control_device(device: str, action: typing.Literal["on", "off", "toggle"]):
    "Switch or toggle any named smart-home device."
    return {"device": device, "action": action}


@cactus.tool
def send_money(
    amount: typing.Annotated[float, cactus.Field(gt=0, le=10000, description="USD, up to 10,000")],
    to: typing.Annotated[str, cactus.Field(pattern=r"^@[a-z0-9_]+$", description="recipient handle")],
    memo: str = "",
):
    "Send money to a handle."
    return {"sent": amount, "to": to, "memo": memo}


BUILTIN_TOOLS = {
    "get_weather": get_weather,
    "set_thermostat": set_thermostat,
    "control_device": control_device,
    "send_money": send_money,
}

# --- Structured extraction ---------------------------------------------------

FIELD_TYPES = {"string", "integer", "number", "boolean", "enum"}
_PY_TYPES = {"string": str, "integer": int, "number": float, "boolean": bool}
MAX_FIELDS = 20


def _build_schema(fields: dict) -> type[BaseModel]:
    """Build a Pydantic model from {name: {type, description?, values?}}.

    All fields are Optional: the engine omits values it cannot ground, and a
    required field would make the whole extraction return None.
    """
    if not fields:
        raise HTTPException(422, "fields must define at least one field")
    if len(fields) > MAX_FIELDS:
        raise HTTPException(422, f"at most {MAX_FIELDS} fields are allowed")
    sig = {}
    for name, spec in fields.items():
        if not name.isidentifier() or name.startswith("_") or name in sig:
            raise HTTPException(422, f"field {name!r}: must be a unique valid identifier (letters, digits, underscores)")
        spec = spec or {}
        ftype = spec.get("type", "string")
        if ftype == "enum":
            values = spec.get("values") or []
            if not values or len(values) > 20:
                raise HTTPException(422, f"field {name!r}: enum needs 1-20 values")
            if not all(isinstance(v, str) and v.strip() for v in values):
                raise HTTPException(422, f"field {name!r}: enum values must be nonempty strings")
            ftype = typing.Literal[tuple(values)]
        elif ftype in _PY_TYPES:
            ftype = _PY_TYPES[ftype]
        else:
            raise HTTPException(422, f"field {name!r}: type must be one of {sorted(FIELD_TYPES)}")
        desc = spec.get("description")
        sig[name] = (
            typing.Optional[ftype],
            Field(default=None, description=desc) if desc else None,
        )
    return create_model("ExtractionSchema", **sig)


class ExtractRequest(BaseModel):
    text: str = Field(..., min_length=1, description="Source text to extract from")
    fields: dict = Field(..., description='e.g. {"city": {"type": "string", "description": "shipping city"}, "status": {"type": "enum", "values": ["confirmed", "pending"]}}')
    system: str | None = Field(None, description="Optional system facts (date, locale, ...)")
    strict: bool = Field(True, description="Reject ungrounded / invented values instead of returning them")


@app.post("/api/extract")
def api_extract(req: ExtractRequest) -> dict:
    _require_agent()
    schema = _build_schema(req.fields)
    t0 = time.monotonic()
    try:
        result = cactus.extract(req.text, schema, system=req.system, weights=str(WEIGHTS), strict=req.strict)
    except cactus.ExtractionValidationError as exc:
        raise HTTPException(422, f"Extraction rejected (strict): {exc}") from exc
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(502, f"Extraction failed: {exc}") from exc
    if result is None:
        return {"extraction": {}, "matched": False, "missing": list(schema.model_fields), "latency_ms": round((time.monotonic() - t0) * 1000, 1)}
    data = result.model_dump()
    return {
        "extraction": data,
        "matched": any(v is not None for v in data.values()),
        "missing": [k for k, v in data.items() if v is None],
        "latency_ms": round((time.monotonic() - t0) * 1000, 1),
    }


# --- Tool calling -------------------------------------------------------------


class AgentRequest(BaseModel):
    query: str = Field(..., min_length=1, description="User query, e.g. 'make it 21 and cool the room'")
    tools: list[str] = Field(default_factory=lambda: list(BUILTIN_TOOLS), description="Names of built-in tools to enable")
    system: str | None = Field(None, description="Optional system facts (date, locale, device, ...)")
    max_steps: int = Field(8, ge=1, le=16, description="Max tool-call rounds")
    strict: bool = Field(True, description="Refuse ungrounded arguments")


@app.get("/api/tools")
def api_tools() -> dict:
    return {
        "tools": [
            {"name": name, "description": (fn.__doc__ or "").strip().splitlines()[0] if fn.__doc__ else ""}
            for name, fn in BUILTIN_TOOLS.items()
        ]
    }


@app.post("/api/agent")
def api_agent(req: AgentRequest) -> dict:
    _require_agent()
    unknown = [t for t in req.tools if t not in BUILTIN_TOOLS]
    if unknown:
        raise HTTPException(422, f"unknown tools: {unknown} (available: {list(BUILTIN_TOOLS)})")
    if not req.tools:
        raise HTTPException(422, "enable at least one tool")
    t0 = time.monotonic()
    try:
        # Fresh instance per request: one toolset and one conversation each.
        kwargs = {"tools": [BUILTIN_TOOLS[t] for t in req.tools], "weights": str(WEIGHTS)}
        if req.system:
            kwargs["system"] = req.system
        agent = cactus.Cactus(**kwargs)
        result = agent.run(req.query, max_steps=req.max_steps, strict=req.strict)
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(502, f"Agent run failed: {exc}") from exc
    return {"result": result, "latency_ms": round((time.monotonic() - t0) * 1000, 1)}


# --- Embeddings ---------------------------------------------------------------


class EmbedRequest(BaseModel):
    text: str = Field(..., min_length=1, description="Text to embed")


@app.post("/api/embed")
def api_embed(req: EmbedRequest) -> dict:
    agent = _require_agent()
    t0 = time.monotonic()
    try:
        vec = agent.embed(req.text)
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(502, f"Embedding failed: {exc}") from exc
    embedding = [round(float(x), 6) for x in vec]
    return {"embedding": embedding, "dim": len(embedding), "latency_ms": round((time.monotonic() - t0) * 1000, 1)}


# Static frontend — mounted last so /api routes take precedence.
app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="static")