"""
serve.py — Model service for the NC drafting feature (cleaned-up version).

Improvements over the basic version:
  1. Strips the auditor's raw note (and their typos) from the output —
     the note only STEERS the model; the returned text is the model's own.
  2. Lower temperature + tighter top_k = less random, fewer weird words.
  3. Trims the statement to end at a complete sentence (no mid-word cutoff).

Run:
    source venv/bin/activate
    uvicorn serve:app --host 0.0.0.0 --port 8000
"""

import re
import torch
from typing import Optional
from fastapi import FastAPI
from pydantic import BaseModel
from model import GPT

# ---- load the trained model ONCE at startup ----
ckpt = torch.load("nc_model.pt", map_location="cpu")
cfg = ckpt["config"]
stoi, itos = ckpt["stoi"], ckpt["itos"]
model = GPT(**cfg)
model.load_state_dict(ckpt["model"])
model.eval()

app = FastAPI()

VALID_TYPES = {"minor": "Minor", "major": "Major", "observation": "Observation"}


class DraftRequest(BaseModel):
    nc_type: Optional[str] = None
    note: str = ""
    max_new_tokens: int = 400
    temperature: float = 0.5     # lower = cleaner, less random
    top_k: int = 20              # tighter = safer word choices


def encode(s: str):
    return [stoi[c] for c in s if c in stoi]


def guess_nc_type(note: str) -> str:
    n = (note or "").lower()
    major_words = [
        "not established", "no system", "completely", "totally", "absence of",
        "never", "no records", "no evidence", "fire", "safety", "injury",
        "accident", "legal", "critical", "serious", "repeated", "systemic",
        "not implemented", "not conducted", "expired", "no procedure",
    ]
    observation_words = [
        "could be improved", "improvement", "recommend", "suggest", "minor",
        "consider", "opportunity", "may benefit", "better to", "advisable",
    ]
    if any(w in n for w in major_words):
        return "Major"
    if any(w in n for w in observation_words):
        return "Observation"
    return "Minor"


def parse_draft(text: str) -> dict:
    """Turn raw model text into statement / clause / corrective_action."""
    if "<|end|>" in text:
        text = text.split("<|end|>")[0]
    text = text.replace("<|nc|>", "").strip()

    statement, clause, corrective = "", "", ""

    m = re.search(r"Statement:\s*(.*?)(?:\n\s*Clause:|\Z)", text, re.S | re.I)
    if m:
        statement = m.group(1).strip()

    m = re.search(r"Clause:\s*(.*?)(?:\n\s*Corrective Action:|\Z)", text, re.S | re.I)
    if m:
        clause = " ".join(m.group(1).split())

    m = re.search(r"Corrective Action:\s*(.*)\Z", text, re.S | re.I)
    if m:
        corrective = m.group(1).strip()

    return {"statement": statement, "clause": clause, "corrective_action": corrective}


def trim_to_sentence(text: str) -> str:
    """Cut at the last full stop / ! / ? so it doesn't end mid-word."""
    text = text.strip()
    for i in range(len(text) - 1, -1, -1):
        if text[i] in ".!?":
            return text[: i + 1].strip()
    return text


def strip_note_prefix(statement: str, note: str) -> str:
    """Remove the auditor's raw note from the front so their typos don't show."""
    note_clean = (note or "").strip()
    s = statement
    if note_clean and s.lower().startswith(note_clean.lower()):
        s = s[len(note_clean):].lstrip(" ,.-—:\n")
    if s:
        s = s[0].upper() + s[1:]
    return s


@app.get("/health")
def health():
    return {"ok": True}


@app.post("/draft")
def draft(req: DraftRequest):
    # 1) NC type: auditor's choice wins; else auto-guess.
    if req.nc_type and req.nc_type.strip():
        nc_type = VALID_TYPES.get(req.nc_type.strip().lower(), req.nc_type.strip())
        nc_type_source = "user"
    else:
        nc_type = guess_nc_type(req.note)
        nc_type_source = "auto"

    # 2) Prompt in the trained format.
    prompt = f"<|nc|>\nType: {nc_type}\nStatement: {req.note}"
    ids = encode(prompt) or [0]
    idx = torch.tensor([ids], dtype=torch.long)

    # 3) Generate.
    out = model.generate(
        idx,
        max_new_tokens=req.max_new_tokens,
        temperature=req.temperature,
        top_k=req.top_k,
    )[0].tolist()
    raw = "".join(itos[i] for i in out)

    # 4) Parse + clean.
    fields = parse_draft(raw)
    stmt = strip_note_prefix(fields["statement"], req.note)
    stmt = trim_to_sentence(stmt)
    fields["statement"] = stmt
    fields["clause"] = trim_to_sentence(fields["clause"]) if fields["clause"] else ""

    return {
        "nc_type": nc_type,
        "nc_type_source": nc_type_source,
        "statement": fields["statement"],
        "clause": fields["clause"],
        "corrective_action": fields["corrective_action"],
        "raw": raw,
    }