"""
serve_rag.py — STEP 5: RAG-enabled /draft endpoint.

Same /draft contract as before, so your NestJS previous-nc module and
Next.js buttons do NOT change. The difference is INSIDE: it retrieves
similar real findings, then asks a generation model to write a new one.

The generation model is pluggable:
  - For LEARNING now: uses a free Google Gemini API (set GEMINI_API_KEY).
  - For PRIVATE production later: replace generate_with_model() with a call
    to your self-hosted open model (e.g. Ollama). Nothing else changes.

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

import os
import re
from typing import Optional
from fastapi import FastAPI
from pydantic import BaseModel
import chromadb
from sentence_transformers import SentenceTransformer

DB_PATH = "./nc_vectors"
COLLECTION = "findings"

# ---- load retrieval pieces once at startup ----
embedder = SentenceTransformer("all-MiniLM-L6-v2")
client = chromadb.PersistentClient(path=DB_PATH)
col = client.get_collection(COLLECTION)

app = FastAPI()
VALID_TYPES = {"minor": "Minor", "major": "Major", "observation": "Observation"}


class DraftRequest(BaseModel):
    nc_type: Optional[str] = None
    note: str = ""
    n_examples: int = 3


def guess_nc_type(note: str) -> str:
    n = (note or "").lower()
    major = ["no system", "not established", "fire", "safety", "no records",
             "not implemented", "expired", "absence of", "critical"]
    obs = ["improve", "recommend", "suggest", "consider", "opportunity"]
    if any(w in n for w in major):
        return "Major"
    if any(w in n for w in obs):
        return "Observation"
    return "Minor"


def retrieve_examples(note: str, k: int):
    """RAG core: find the k most similar real findings."""
    qvec = embedder.encode(note).tolist()
    hits = col.query(query_embeddings=[qvec], n_results=k)
    return hits["documents"][0]


def build_prompt(nc_type: str, note: str, examples: list) -> str:
    """Ground the model in real examples from your own data."""
    ex_text = "\n\n".join(f"Example {i+1}:\n{e}" for i, e in enumerate(examples))
    return (
        "You are an ISO audit assistant. Write a single, professional "
        "Non-Conformity (NC) finding based on the auditor's note. "
        "Match the STYLE and STRUCTURE of the real examples below.\n\n"
        "Respond in EXACTLY this format, and ALL fields are REQUIRED:\n"
        "Type: <the NC type>\n"
        "Statement: <one clear professional sentence>\n"
        "Clause: <the most relevant ISO standard and clause>\n"
        "Corrective Action: <a concrete, professionally-worded corrective action "
        "written in formal ISO audit language. State what must be done to correct "
        "the non-conformity and prevent recurrence. NEVER leave blank and NEVER write "
        "a placeholder like to be provided by the organization>\n\n"
        "Keep each field to one concise sentence.\n\n"
        f"NC type: {nc_type}\n\n"
        f"Real examples from our records:\n{ex_text}\n\n"
        f"Auditor's note: {note}\n\n"
        "Write the new finding now:"
    )


def generate_with_model(prompt: str) -> str:
    """
    PLUGGABLE generation step.
    Now: free Google Gemini (for learning). Data leaves the server, so use
    non-sensitive notes while testing.
    Later: replace this whole function body with a call to your self-hosted
    open model (e.g. POST to a local Ollama server) for full privacy.

    Retries automatically on temporary errors (429 rate limit, 5xx overload).
    """
    import urllib.request
    import urllib.error
    import json
    import time

    api_key = os.environ.get("GEMINI_API_KEY", "")
    if not api_key:
        return "[No GEMINI_API_KEY set — this is where the model output goes.]"

    url = (
        "https://generativelanguage.googleapis.com/v1beta/models/"
        "gemini-2.5-flash:generateContent?key=" + api_key
    )
    payload = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {"temperature": 0.4, "maxOutputTokens": 1024},
    }
    data_bytes = json.dumps(payload).encode()

    last_err = None
    for attempt in range(4):  # up to 4 tries
        try:
            req = urllib.request.Request(
                url, data=data_bytes,
                headers={"Content-Type": "application/json"},
            )
            with urllib.request.urlopen(req, timeout=30) as r:
                data = json.loads(r.read())
            return data["candidates"][0]["content"]["parts"][0]["text"]
        except urllib.error.HTTPError as e:
            last_err = e
            # Retry only on temporary errors: 429 (rate limit), 5xx (overload)
            if e.code in (429, 500, 502, 503, 504) and attempt < 3:
                time.sleep(1 * (attempt + 1))  # wait 1s, 2s, 3s
                continue
            raise
    raise last_err


def _clean(s: str) -> str:
    """Remove markdown bold/asterisks and tidy whitespace."""
    s = s.replace("**", "").replace("*", "")
    s = " ".join(s.split())
    return s.strip()


def parse_draft(text: str) -> dict:
    text = text.replace("**", "")  # strip markdown bold up front
    statement, clause, corrective = text.strip(), "", ""
    m = re.search(r"(?:NCR statement|Statement)[:\-]\s*(.*?)(?:\n\s*(?:Clause|Criteria)|\Z)", text, re.S | re.I)
    if m:
        statement = m.group(1).strip()
    m = re.search(r"(?:Clause|Criteria)[^:]*:\s*(.*?)(?:\n\s*Corrective|\Z)", text, re.S | re.I)
    if m:
        clause = " ".join(m.group(1).split())
    m = re.search(r"Corrective[^:]*:\s*(.*)\Z", text, re.S | re.I)
    if m:
        corrective = m.group(1).strip()
    return {
        "statement": _clean(statement),
        "clause": _clean(clause),
        "corrective_action": _clean(corrective),
    }


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


@app.post("/draft")
def draft(req: DraftRequest):
    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"

    examples = retrieve_examples(req.note, req.n_examples)   # RAG retrieval
    prompt = build_prompt(nc_type, req.note, examples)       # RAG prompt
    raw = generate_with_model(prompt)                        # generation
    fields = parse_draft(raw)

    ca = fields["corrective_action"].strip()
    placeholder = (not ca) or ("to be provided" in ca.lower()) or ("organization to" in ca.lower())
    if placeholder:
        fields["corrective_action"] = (
            "Conduct a root cause analysis of the identified non-conformity, implement "
            "appropriate corrective measures to bring the affected process into compliance "
            "with the applicable requirements, and verify the effectiveness of the actions "
            "taken within an agreed timeframe to prevent recurrence."
        )

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