"""HF Space demo for construction-code-cite (v3 · Llama 4 Scout 17B-16E).

The v3 model is 109B total params (17B active MoE) and does not fit in-Space.
Inference goes to Together AI's hosted endpoint; the Space runs the RAG
pipeline + verifier + Gradio UI only. If TOGETHER_API_KEY is missing or the
endpoint returns an error, we fall back to the v2 (Llama 3.2 3B) in-Space
adapter so the demo never goes dark.

Set the following secrets in the Space:
- TOGETHER_API_KEY   (required for v3)
- V3_MODEL_ID        (default: rigidhat/llama-4-scout-17b-construction-codecite-v3)
- FALLBACK_ADAPTER   (default: rigidhat/llama-3.2-3b-construction-codecite-v2)
"""
from __future__ import annotations

import json
import os
import re
import time
from pathlib import Path

import gradio as gr

V3_MODEL_ID = os.environ.get(
    "V3_MODEL_ID", "rigidhat/llama-4-scout-17b-construction-codecite-v3"
)
TOGETHER_API_KEY = os.environ.get("TOGETHER_API_KEY", "")
FALLBACK_BASE = os.environ.get("FALLBACK_BASE", "meta-llama/Llama-3.2-3B-Instruct")
FALLBACK_ADAPTER = os.environ.get(
    "FALLBACK_ADAPTER", "rigidhat/llama-3.2-3b-construction-codecite-v2"
)
DATASET_REPO = os.environ.get("DATASET_REPO", "rigidhat/construction-code-corpus-v1")
MAX_NEW_TOKENS = 384
RAG_K = 5

CORPUS_PATH = Path(__file__).parent / "osha_1926_corpus.jsonl"


def ensure_corpus() -> Path:
    if CORPUS_PATH.exists():
        return CORPUS_PATH
    from huggingface_hub import hf_hub_download

    downloaded = hf_hub_download(
        repo_id=DATASET_REPO,
        filename="osha_1926_corpus.jsonl",
        repo_type="dataset",
    )
    Path(downloaded).replace(CORPUS_PATH)
    return CORPUS_PATH


_STANDARD_RE = re.compile(r"1926(?:\.\d+[A-Za-z]?)(?:\([a-zA-Z0-9ivxIVX]+\))*")


def build_verifier(corpus_path: Path):
    sections: dict[str, dict] = {}
    with corpus_path.open("r", encoding="utf-8") as fh:
        for line in fh:
            rec = json.loads(line)
            cite = (rec.get("citation") or "").strip()
            match = _STANDARD_RE.search(cite)
            if match:
                section = match.group(0).split("(")[0]
                sections[section] = rec

    def verify(raw: str) -> tuple[bool, str]:
        match = _STANDARD_RE.search(raw or "")
        if not match:
            return False, ""
        section = match.group(0).split("(")[0]
        rec = sections.get(section)
        if not rec:
            return False, ""
        return True, rec.get("heading") or ""

    return verify


def build_bm25(corpus_path: Path):
    from rank_bm25 import BM25Okapi

    token_re = re.compile(r"[a-zA-Z][a-zA-Z\-]+|\d+")
    stopwords = frozenset(
        "the a an and or but of in on for to with at by from as is are be been being "
        "this that these those it its which who whom whose what when where why how".split()
    )

    def tok(text: str) -> list[str]:
        return [t.lower() for t in token_re.findall(text or "") if t.lower() not in stopwords]

    records: list[dict] = []
    tokens: list[list[str]] = []
    with corpus_path.open("r", encoding="utf-8") as fh:
        for line in fh:
            rec = json.loads(line)
            records.append(rec)
            tokens.append(tok(f"{rec.get('heading', '')}\n{rec.get('text', '')}"))
    bm25 = BM25Okapi(tokens)

    def search(query: str, k: int = RAG_K):
        query_tokens = tok(query)
        if not query_tokens:
            return []
        scores = bm25.get_scores(query_tokens)
        import numpy as np

        top = np.argpartition(scores, -k)[-k:]
        order = sorted(top, key=lambda i: scores[i], reverse=True)
        hits = []
        for i in order[:k]:
            rec = records[i]
            hits.append({
                "section": rec.get("citation") or "",
                "heading": rec.get("heading") or "",
                "bm25": float(scores[i]),
            })
        return hits

    return search


PROMPT = """You are an OSHA-trained construction-safety classifier.

Output STRICT JSON only with this shape (no prose, no markdown):
{{"hazards":[{{"code_event":{{"id":"<OIICS event id>","title":"<short>"}},
"code_source":{{"id":"<OIICS source id>","title":"<short>"}},
"code_nature":{{"id":"<OIICS nature id>","title":"<short>"}},
"code_body":{{"id":"<OIICS body id>","title":"<short>"}},
"severity":"low|moderate|high"}}],
"citations":[{{"standard":"<1926.X>","section_heading":"<heading>"}}]}}

OIICS code IDs are short numeric strings (1-4 digits). Use "OTHER" only when
no specific code applies. Cite 0-3 OSHA 1926 sections from the candidate list
below if any apply.

Retrieved OSHA 1926 sections (BM25-ranked candidates):
{candidates}

Incident narrative:
\"\"\"{narrative}\"\"\"

JSON:"""


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


def parse_json(raw: str) -> dict:
    if not raw:
        return {}
    match = _JSON_RE.search(raw)
    if not match:
        return {}
    snippet = match.group(0)
    try:
        return json.loads(snippet)
    except json.JSONDecodeError:
        last = snippet.rfind("}")
        if last != -1:
            try:
                return json.loads(snippet[: last + 1])
            except json.JSONDecodeError:
                pass
    return {}


_STATE = {"search": None, "verify": None, "fallback": None}


def get_search_verify():
    if _STATE["search"] is not None:
        return _STATE["search"], _STATE["verify"]
    corpus = ensure_corpus()
    _STATE["search"] = build_bm25(corpus)
    _STATE["verify"] = build_verifier(corpus)
    return _STATE["search"], _STATE["verify"]


def get_fallback():
    """Lazy-load v2 in-Space adapter as the fallback path."""
    if _STATE["fallback"] is not None:
        return _STATE["fallback"]
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    print(f"Loading fallback base: {FALLBACK_BASE}")
    tokenizer = AutoTokenizer.from_pretrained(FALLBACK_BASE, trust_remote_code=True)
    base = AutoModelForCausalLM.from_pretrained(
        FALLBACK_BASE, torch_dtype=torch.float32, trust_remote_code=True
    )
    print(f"Loading fallback adapter: {FALLBACK_ADAPTER}")
    model = PeftModel.from_pretrained(base, FALLBACK_ADAPTER)
    _STATE["fallback"] = {"tokenizer": tokenizer, "model": model, "torch": torch}
    return _STATE["fallback"]


def generate_together(prompt: str) -> tuple[str, str]:
    """Call Together AI hosted endpoint. Returns (text, path_label)."""
    if not TOGETHER_API_KEY:
        raise RuntimeError("TOGETHER_API_KEY not set")
    try:
        from together import Together
    except ImportError as e:
        raise RuntimeError(f"together package not installed: {e}")

    client = Together(api_key=TOGETHER_API_KEY)
    response = client.chat.completions.create(
        model=V3_MODEL_ID,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=MAX_NEW_TOKENS,
        temperature=0.0,
    )
    return response.choices[0].message.content, "v3 · Llama 4 Scout 17B-16E (Together)"


def generate_fallback(prompt: str) -> tuple[str, str]:
    """Fall back to v2 in-Space."""
    pipe = get_fallback()
    messages = [{"role": "user", "content": prompt}]
    chat = pipe["tokenizer"].apply_chat_template(
        messages, add_generation_prompt=True, tokenize=False
    )
    inputs = pipe["tokenizer"](chat, return_tensors="pt")
    with pipe["torch"].no_grad():
        out = pipe["model"].generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            pad_token_id=pipe["tokenizer"].eos_token_id,
        )
    generated = out[0][inputs["input_ids"].shape[1]:]
    raw = pipe["tokenizer"].decode(generated, skip_special_tokens=True)
    return raw, "v2 · Llama 3.2 3B (in-Space fallback)"


def format_candidates(hits) -> str:
    if not hits:
        return "(no high-confidence candidates)"
    return "\n".join(f"- {h['section']}: {h['heading'][:80]}" for h in hits)


def predict(narrative: str):
    if not (narrative or "").strip():
        return "{}", "(paste an incident narrative first)", "—", "—"
    t0 = time.time()
    search, verify = get_search_verify()
    hits = search(narrative, k=RAG_K)
    prompt = PROMPT.format(candidates=format_candidates(hits), narrative=narrative[:1800])

    path_label = ""
    raw = ""
    try:
        raw, path_label = generate_together(prompt)
    except Exception as e:
        print(f"Together path failed: {e}. Falling back to v2 in-Space.")
        raw, path_label = generate_fallback(prompt)

    parsed = parse_json(raw)
    if parsed and "citations" in parsed:
        for c in parsed["citations"]:
            is_valid, heading = verify(c.get("standard", ""))
            c["verified"] = is_valid
            if heading and not c.get("section_heading"):
                c["section_heading"] = heading

    rag_view = "\n".join(f"- {h['section']}: {h['heading'][:80]}" for h in hits)
    return json.dumps(parsed, indent=2), rag_view, f"{time.time() - t0:.2f}s", path_label


EXAMPLES = [
    "Worker fell from second-story scaffold platform while installing siding, sustained multiple fractures.",
    "Employee's hand was caught between two pieces of trench shoring equipment causing partial amputation of two fingers.",
    "Electrician contacted overhead power line while operating boom lift on a commercial roofing project.",
]


with gr.Blocks(title="Construction Code-Citation") as demo:
    gr.Markdown("# Construction Code-Citation Model (v3 · Llama 4 Scout 17B-16E · AutoScientist)")
    gr.Markdown(
        "Llama 4 Scout 17B-16E (MoE) fine-tuned by **AutoScientist** on OSHA Severe "
        "Injury Reports for the [Adaption Labs AutoScientist Challenge](https://adaptionlabs.ai/auto-scientist) "
        "\"All Other Domains\" category. Given a construction-site incident narrative, "
        "returns strict JSON with OIICS hazard codes plus OSHA 29 CFR 1926 citations, "
        "verifier-grounded against the corpus. **Inference via Together AI hosted "
        "endpoint** — v2 (Llama 3.2 3B) auto-falls-back if the endpoint is unavailable."
    )
    with gr.Row():
        with gr.Column():
            narrative = gr.Textbox(
                label="Incident narrative",
                lines=5,
                placeholder="Describe the construction-site incident...",
            )
            submit = gr.Button("Classify", variant="primary")
            gr.Examples(EXAMPLES, inputs=narrative)
        with gr.Column():
            output_json = gr.Code(label="Hazards + Citations (JSON)", language="json")
            rag_view = gr.Textbox(label="OSHA 1926 RAG candidates (BM25)", lines=6)
            with gr.Row():
                elapsed = gr.Textbox(label="Latency", lines=1)
                path = gr.Textbox(label="Model path", lines=1)

    submit.click(predict, inputs=[narrative], outputs=[output_json, rag_view, elapsed, path])

    gr.Markdown(
        "**Artifacts:** "
        "[Dataset](https://huggingface.co/datasets/rigidhat/construction-code-corpus-v1) · "
        "[v3 Model (Llama 4 Scout 17B-16E · AutoScientist)](https://huggingface.co/rigidhat/llama-4-scout-17b-construction-codecite-v3) · "
        "[v2 Model (Llama 3.2 3B · AutoScientist)](https://huggingface.co/rigidhat/llama-3.2-3b-construction-codecite-v2) · "
        "[v1 Baseline (Qwen 2.5 1.5B)](https://huggingface.co/rigidhat/qwen-2.5-construction-codecite-v1) · "
        "[Source](https://github.com/snakezilla/construction-code-llm)"
    )


if __name__ == "__main__":
    demo.launch()