"""
Glaucoma CDSS
=============
Gradio UI  : /
REST API   : /predict  (JSON)  |  /predict/pdf  (PDF file)  |  /health
"""

import os, sys, tempfile, shutil
from pathlib import Path

import cv2
import numpy as np
import gradio as gr
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse, FileResponse

sys.path.insert(0, os.path.dirname(__file__))
from phase3pipeline import Phase3Pipeline

# ── config ────────────────────────────────────────────────────────────────────
REPO_ID = "Nj-1111/EyeeSEE"
EPOCH   = None
TOKEN   = os.getenv("HF_TOKEN_2") or os.getenv("HF_TOKEN")
# ─────────────────────────────────────────────────────────────────────────────

print(f"Loading model  repo={REPO_ID}  epoch={'latest' if EPOCH is None else EPOCH}")
pipeline = Phase3Pipeline(
    repo_id=REPO_ID, epoch=EPOCH, mc_passes=1,
    uncertainty_threshold=0.05, token=TOKEN,
)
print("Model ready.")


# ── shared helpers ────────────────────────────────────────────────────────────

def _load_image(path: str) -> np.ndarray:
    img = cv2.imread(path)
    if img is None:
        raw = np.fromfile(path, dtype=np.uint8)
        img = cv2.imdecode(raw, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"Cannot decode image: {path}")
    return img


def _run_pipeline(path: str) -> dict:
    img    = _load_image(path)
    result = pipeline.run(img)
    return result["report"]


def _build_text(r: dict) -> str:
    isnt = r["isnt"]
    lines = [
        "=" * 52,
        "  GLAUCOMA CDSS — CLINICAL REPORT",
        "=" * 52,
        f"  vCDR              : {r['vcdr']:.4f}",
        f"  Risk Level        : {r['risk_level']}",
        f"  Uncertainty       : {r['uncertainty']:.6f}",
        f"  Sanity Check      : {'PASSED' if r['sanity_passed'] else 'CORRECTED (auto)'}",
        f"  ISNT Rule         : {'Satisfied' if isnt['rule_satisfied'] else 'Violated'}",
        "",
        "  ISNT Rim Thickness",
        f"    Inferior  : {isnt['inferior']:.2f}",
        f"    Superior  : {isnt['superior']:.2f}",
        f"    Nasal     : {isnt['nasal']:.2f}",
        f"    Temporal  : {isnt['temporal']:.2f}",
        "",
        "  Structural",
        f"    Disc Area : {r['disc_area_px']:,} px",
        f"    Cup Area  : {r['cup_area_px']:,} px",
        f"    Cup/Disc  : {r['cup_area_px'] / max(r['disc_area_px'], 1) * 100:.1f}%",
        f"    MC Passes : {pipeline.mc_passes}",
    ]
    if r.get("warnings"):
        lines += ["", "  Warnings"]
        for w in r["warnings"]:
            lines.append(f"    ! {w}")
    lines += [
        "", "=" * 52,
        "  DISCLAIMER: Research prototype. Not a certified medical device.",
        "  Validate with a qualified ophthalmologist.",
        "=" * 52,
    ]
    return "\n".join(lines)


def _build_pdf(text: str) -> str:
    from fpdf import FPDF
    safe = text.replace("—", "-").replace("•", "-").replace("→", "->")
    pdf  = FPDF()
    pdf.add_page()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.set_font("Courier", size=10)
    for line in safe.splitlines():
        pdf.multi_cell(0, 6, txt=line)
    out = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
    pdf.output(out.name)
    return out.name


def _save_upload(upload: UploadFile) -> str:
    suffix = Path(upload.filename).suffix or ".jpg"
    tmp    = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    shutil.copyfileobj(upload.file, tmp)
    tmp.flush()
    return tmp.name


# ── FastAPI app + endpoints ───────────────────────────────────────────────────

app = FastAPI(title="Glaucoma CDSS API", version="1.0")


@app.get("/health")
def health():
    """Check that the model is loaded and ready."""
    return {
        "status": "ok",
        "repo":   REPO_ID,
        "epoch":  "latest" if EPOCH is None else EPOCH,
    }


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    """
    Run glaucoma screening on an uploaded fundus image.

    Request  : multipart/form-data  field name = file
    Accepts  : JPEG, PNG
    Response : JSON clinical report

    JSON fields
    -----------
    vcdr              float   vertical cup-to-disc ratio
    risk_level        string  Healthy | Glaucoma Suspect | High Risk
    uncertainty       float   MC-Dropout variance (lower = more confident)
    sanity_passed     bool    True if cup is naturally inside disc
    isnt              object  inferior/superior/nasal/temporal rim thickness
    structural        object  disc_area_px, cup_area_px, cup_disc_pct
    warnings          list    clinical warning strings
    """
    path = _save_upload(file)
    try:
        r = _run_pipeline(path)
    except Exception as e:
        raise HTTPException(status_code=422, detail=str(e))
    finally:
        os.remove(path)

    return JSONResponse({
        "vcdr":          r["vcdr"],
        "risk_level":    r["risk_level"],
        "uncertainty":   r["uncertainty"],
        "sanity_passed": r["sanity_passed"],
        "isnt": {
            "inferior":       r["isnt"]["inferior"],
            "superior":       r["isnt"]["superior"],
            "nasal":          r["isnt"]["nasal"],
            "temporal":       r["isnt"]["temporal"],
            "rule_satisfied": r["isnt"]["rule_satisfied"],
        },
        "structural": {
            "disc_area_px":  r["disc_area_px"],
            "cup_area_px":   r["cup_area_px"],
            "cup_disc_pct":  round(r["cup_area_px"] / max(r["disc_area_px"], 1) * 100, 1),
        },
        "warnings": r.get("warnings", []),
    })


@app.post("/predict/pdf")
async def predict_pdf(file: UploadFile = File(...)):
    """
    Run glaucoma screening and return a downloadable PDF report.

    Request  : multipart/form-data  field name = file
    Accepts  : JPEG, PNG
    Response : application/pdf
    """
    path = _save_upload(file)
    try:
        r    = _run_pipeline(path)
        text = _build_text(r)
        pdf  = _build_pdf(text)
    except Exception as e:
        raise HTTPException(status_code=422, detail=str(e))
    finally:
        os.remove(path)

    return FileResponse(
        pdf,
        media_type="application/pdf",
        filename="glaucoma_report.pdf",
    )


# ── Gradio UI (mounted at /) ──────────────────────────────────────────────────

def analyse(file_path):
    if file_path is None:
        return "No file uploaded.", None, "Please upload a JPEG or PNG fundus image."
    try:
        r    = _run_pipeline(file_path)
        text = _build_text(r)
        pdf  = _build_pdf(text)
        return text, pdf, "Analysis completed."
    except Exception as e:
        return f"Error: {e}", None, f"Analysis failed: {e}"


def set_busy():
    return gr.update(interactive=False), "Running analysis, please wait..."


def set_ready():
    return gr.update(interactive=True)


with gr.Blocks(title="Glaucoma CDSS") as gradio_ui:

    gr.Markdown("## Glaucoma CDSS\nUpload a retinal fundus image to receive a clinical screening report.")

    with gr.Row():
        file_in = gr.File(
            label="Fundus Image (JPEG / PNG)",
            file_types=[".jpg", ".jpeg", ".png"],
            type="filepath",
        )
        run_btn = gr.Button("Analyse", variant="primary")

    status_box  = gr.Textbox(label="Status", value="Awaiting upload.", interactive=False)
    report_box  = gr.Textbox(label="Clinical Report", lines=28, interactive=False)
    pdf_out     = gr.File(label="Download PDF Report")

    gr.Markdown("""
---
### How to interpret results

| vCDR | Indication |
|------|------------|
| 0.30 – 0.50 | Usually healthy |
| 0.50 – 0.65 | Borderline |
| 0.65 – 0.80 | Glaucoma suspect |
| > 0.80 | High risk |

**ISNT Rule** — healthy nerves follow Inferior > Superior > Nasal > Temporal.
Violation suggests neuro-retinal rim thinning.

**Uncertainty** — values above 0.05 indicate low model confidence (check image quality).

---
*Research prototype — NOT a medical device. All results must be reviewed by a qualified ophthalmologist.*
""")

    run_btn.click(
        fn=set_busy, inputs=None, outputs=[run_btn, status_box], queue=False,
    ).then(
        fn=analyse, inputs=[file_in], outputs=[report_box, pdf_out, status_box],
    ).then(
        fn=set_ready, inputs=None, outputs=[run_btn], queue=False,
    )


app = gr.mount_gradio_app(app, gradio_ui, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)