"""
ClinIQ — gr.Server backend.

Key feature: proactive Drug Safety Check runs automatically after upload
— no question needed. The 3B model acts as an on-device pharmacist.

Hackathon quests:
  🎨 Off-Brand        — gr.Server + full custom HTML/CSS/JS
  📡 Sharing is Caring — agent traces pushed to HF Hub dataset
  🦙 Llama Champion   — inference via llama.cpp on Modal
  🎯 Well-Tuned       — Qwen2.5-3B-Instruct from Hub
  🐜 Tiny Titan       — ≤4B model
"""

from __future__ import annotations

import base64
import datetime
import json
import os
import re
import tempfile
from pathlib import Path
from typing import Iterator, List, Tuple

import gradio as gr
from fastapi.responses import HTMLResponse, Response

from agent import build_graph, run_query, stream_query
from retriever import HybridRetriever
from safety_checker import run_safety_check, report_to_dict

# ── Config ─────────────────────────────────────────────────────────────────────

SAMPLE_DIR      = Path(__file__).parent / "sample_docs"
FRONTEND        = Path(__file__).parent / "frontend"
HF_TOKEN        = os.getenv("HF_TOKEN", "")
HF_DATASET_REPO = os.getenv("HF_DATASET_REPO", "karthikmulugu08/cliniq-traces")
MODAL_ENDPOINT  = os.getenv("MODAL_ENDPOINT", "")

# ── Global state ───────────────────────────────────────────────────────────────

_retriever   = HybridRetriever()
_graph       = build_graph(_retriever)
_raw_docs: List[Tuple[str, str]] = []   # (original_text, filename) for safety checker


def _auto_load_samples() -> None:
    """Pre-load sample docs at startup so the app is ready immediately after any restart."""
    if not SAMPLE_DIR.exists():
        return
    docs = [
        (p.read_text(encoding="utf-8", errors="ignore"), p.name)
        for p in sorted(SAMPLE_DIR.glob("*.txt"))
    ]
    if not docs:
        return
    _retriever.add_documents(docs)
    _raw_docs.clear()
    _raw_docs.extend(docs)
    print(f"[startup] Auto-loaded {len(docs)} sample docs → {len(_retriever.chunks)} chunks ready")

_auto_load_samples()


# ── Model call helper (shared by agent + safety checker) ──────────────────────

def _call_model(prompt: str, max_tokens: int = 600, json_mode: bool = False) -> str:
    import httpx
    if MODAL_ENDPOINT:
        resp = httpx.post(
            MODAL_ENDPOINT,
            json={"prompt": prompt, "max_tokens": max_tokens, "json_mode": json_mode},
            timeout=300,
        )
        resp.raise_for_status()
        return resp.json()["text"].strip()
    # Local CPU fallback
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline as hf_pipeline
    MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-3B-Instruct")
    if not hasattr(_call_model, "_pipe"):
        tok   = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
        _call_model._pipe = hf_pipeline("text-generation", model=model, tokenizer=tok)
    out = _call_model._pipe(prompt, max_new_tokens=max_tokens, do_sample=False,
                            temperature=None, top_p=None)
    return out[0]["generated_text"][len(prompt):].strip()


# ── Document helpers ───────────────────────────────────────────────────────────

def _read_file(path: str) -> str:
    """
    Extract text from TXT, PDF (native or scanned), or image files.
    - TXT/MD         → plain read
    - PDF with text  → pypdf fast extract
    - PDF scanned    → pdf2image + tesseract OCR (fallback)
    - JPG/PNG/TIFF   → tesseract OCR directly
    """
    p = Path(path)
    suffix = p.suffix.lower()

    # ── Plain text ────────────────────────────────────────────────────────────
    if suffix in (".txt", ".md", ".csv"):
        return p.read_text(encoding="utf-8", errors="ignore")

    # ── Native PDF (selectable text) ──────────────────────────────────────────
    if suffix == ".pdf":
        from pypdf import PdfReader
        pages = PdfReader(path).pages
        text = "\n".join(page.extract_text() or "" for page in pages).strip()
        if len(text) > 100:          # has real text → done
            return text
        # Scanned PDF — fall through to OCR below
        return _ocr_pdf(path)

    # ── Image files ───────────────────────────────────────────────────────────
    if suffix in (".jpg", ".jpeg", ".png", ".tiff", ".tif", ".bmp", ".webp"):
        return _ocr_image(path)

    # ── Unknown fallback ──────────────────────────────────────────────────────
    try:
        return p.read_text(encoding="utf-8", errors="ignore")
    except Exception:
        return ""


def _ocr_image(path: str) -> str:
    """Run tesseract OCR on a single image file."""
    try:
        import pytesseract
        from PIL import Image
        img = Image.open(path)
        # Upscale small images for better OCR accuracy
        w, h = img.size
        if w < 1000:
            scale = 1000 / w
            img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
        return pytesseract.image_to_string(img, lang="eng", config="--psm 6")
    except Exception as e:
        return f"[OCR failed: {e}]"


def _ocr_pdf(path: str) -> str:
    """Convert each PDF page to image then OCR (for scanned PDFs)."""
    try:
        import pytesseract
        from pdf2image import convert_from_path
        pages = convert_from_path(path, dpi=250)
        texts = []
        for page_img in pages:
            texts.append(pytesseract.image_to_string(page_img, lang="eng", config="--psm 6"))
        return "\n\n".join(texts)
    except Exception as e:
        return f"[OCR failed: {e}]"


def _sanitize(text: str) -> str:
    """Replace non-ASCII chars that break JSON serialization in llama-server."""
    import re
    return re.sub(r'[^\x00-\x7F]', '-', text)


def _build_full_context() -> str:
    """
    Build safety context from the original (un-chunked) document text.
    Extracts ALLERGIES and MEDICATIONS sections directly, preserving newlines
    so the regex in safety_checker can find them reliably.
    """
    if not _raw_docs:
        return ""

    parts = []
    for text, name in _raw_docs:
        text = _sanitize(text)
        # Flexible patterns covering many real-world clinical document formats
        am = re.search(
            r'((?:drug\s+|known\s+|food\s+)?(?:allergies?|allergy|sensitivities?|adverse\s+reactions?)'
            r'[^\n]*\n(?:.*\n?){1,10})',
            text, re.IGNORECASE
        )
        mm = re.search(
            r'((?:current\s+|active\s+|home\s+|discharge\s+)?(?:medications?|meds?|drugs?|prescriptions?|rx)'
            r'[^\n]*\n(?:.*\n?){1,15})',
            text, re.IGNORECASE
        )
        section = f"[{name}]"
        if am:
            section += "\n" + am.group(0)[:500].strip()
        if mm:
            section += "\n" + mm.group(0)[:500].strip()
        parts.append(section)

    return "\n\n---\n\n".join(parts)


# ── HF Hub trace sharing ───────────────────────────────────────────────────────

def _push_trace(question: str, answer: str, trace: list, structured_data=None) -> str:
    if not HF_TOKEN:
        return ""
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=HF_TOKEN)
        api.create_repo(HF_DATASET_REPO, repo_type="dataset", exist_ok=True)
        record = {
            "timestamp":      datetime.datetime.utcnow().isoformat() + "Z",
            "question":       question,
            "answer":         answer,
            "structured_data": json.dumps(structured_data) if structured_data else None,
            "trace":          json.dumps(trace),
            "model":          "Qwen2.5-3B-Instruct (llama.cpp Q4_K_M)",
        }
        fname = f"traces/trace_{datetime.datetime.utcnow().strftime('%Y%m%d_%H%M%S_%f')}.json"
        api.upload_file(
            path_or_fileobj=json.dumps(record, indent=2).encode(),
            path_in_repo=fname,
            repo_id=HF_DATASET_REPO,
            repo_type="dataset",
            commit_message=f"ClinIQ agent trace: {question[:60]}",
        )
        return f"https://huggingface.co/datasets/{HF_DATASET_REPO}/blob/main/{fname}"
    except Exception as e:
        return f"(trace share error: {e})"


# ── gr.Server ──────────────────────────────────────────────────────────────────

app = gr.Server()


@app.api(name="ingest")
def ingest(files_json: str) -> str:
    items = json.loads(files_json)
    docs: List[Tuple[str, str]] = []
    for item in items:
        raw    = base64.b64decode(item["content_b64"])
        suffix = Path(item["name"]).suffix or ".txt"
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tf:
            tf.write(raw); tf.flush()
            text = _read_file(tf.name)
        docs.append((text, item["name"]))
    _retriever.add_documents(docs)
    _raw_docs.clear()
    _raw_docs.extend(docs)
    return json.dumps({
        "status": "ok",
        "files": len(docs),
        "chunks": len(_retriever.chunks),
        "names": [d[1] for d in docs],
    })


@app.api(name="load_samples")
def load_samples() -> str:
    if not SAMPLE_DIR.exists():
        return json.dumps({"status": "error", "message": "sample_docs/ not found"})
    docs = [(p.read_text(encoding="utf-8", errors="ignore"), p.name)
            for p in sorted(SAMPLE_DIR.glob("*.txt"))]
    if not docs:
        return json.dumps({"status": "error", "message": "No sample files found"})
    _retriever.add_documents(docs)
    _raw_docs.clear()
    _raw_docs.extend(docs)
    return json.dumps({
        "status": "ok",
        "files":  len(docs),
        "chunks": len(_retriever.chunks),
        "names":  [d[1] for d in docs],
    })


@app.api(name="safety_check")
def safety_check() -> str:
    """
    Proactive drug safety check — runs after document ingestion.
    No question needed. Returns colour-coded alerts.
    """
    if not _retriever.ready:
        return json.dumps({"error": "No documents indexed yet."})
    context = _build_full_context()
    report  = run_safety_check(context, _call_model)
    return json.dumps(report_to_dict(report))


@app.api(name="ask_stream")
def ask_stream(question: str) -> Iterator[str]:
    if not question.strip():
        yield json.dumps({"type": "error", "message": "Empty question"}); return
    if not _retriever.ready:
        yield json.dumps({"type": "error", "message": "No documents indexed."}); return

    final_answer, final_trace, structured = "", [], None
    for event in stream_query(_graph, question):
        if event["type"] == "trace_step":
            final_trace.append(event["step"])
        elif event["type"] == "answer":
            final_answer = event["answer"]
            structured   = event.get("structured_data")
        yield json.dumps(event)

    trace_url = _push_trace(question, final_answer, final_trace, structured)
    yield json.dumps({"type": "done", "trace_url": trace_url})


@app.api(name="ask")
def ask(question: str) -> str:
    if not question.strip():
        return json.dumps({"error": "Empty question"})
    if not _retriever.ready:
        return json.dumps({"error": "No documents indexed."})
    result    = run_query(_graph, question)
    trace_url = _push_trace(question, result["answer"], result["trace"],
                            result.get("structured_data"))
    return json.dumps({
        "answer":          result["answer"],
        "structured_data": result.get("structured_data"),
        "query_type":      result["query_type"],
        "trace":           result["trace"],
        "chunks":          [{"source": c.source, "excerpt": c.text[:350]}
                            for c in result["chunks"][:4]],
        "trace_url":       trace_url,
    })


# ── Static frontend ────────────────────────────────────────────────────────────

@app.get("/")
async def homepage():
    return HTMLResponse((FRONTEND / "index.html").read_text(encoding="utf-8"))

@app.get("/style.css")
async def css():
    return Response((FRONTEND / "style.css").read_text(encoding="utf-8"),
                    media_type="text/css")


if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)