main.py

import os
import io
import uuid
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
from PIL import Image
from dotenv import load_dotenv

load_dotenv()

RESULTS_DIR = os.environ.get("RESULTS_DIR", "/tmp/results")
os.makedirs(RESULTS_DIR, exist_ok=True)

_hf_host = os.environ.get("SPACE_HOST", "")
API_BASE_URL = f"https://{_hf_host}" if _hf_host else "http://localhost:7860"

# Track startup state so /health can report honestly
_startup_ok = False
_startup_error = ""

app = FastAPI(title="RadGuard AI Engine", version="1.0")

app.mount("/results", StaticFiles(directory=RESULTS_DIR), name="results")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
async def load_model():
    global _startup_ok, _startup_error
    try:
        from huggingface_hub import hf_hub_download

        chexbert_ckpt = os.environ.get("CHEXBERT_CKPT", "/app/CheXbert/src/chexbert.pth")
        if not os.path.exists(chexbert_ckpt):
            print("📥 Downloading chexbert.pth from HuggingFace Hub...")
            os.makedirs(os.path.dirname(chexbert_ckpt), exist_ok=True)
            hf_hub_download(
                repo_id="alyrraza/radguard-v11",
                filename="chexbert.pth",
                local_dir=os.path.dirname(chexbert_ckpt),
            )
            print("✅ chexbert.pth downloaded")
        else:
            print("✅ chexbert.pth already present")

        from inference.model import get_model, get_tokenizer
        get_model()
        get_tokenizer()
        _startup_ok = True
        print("✅ Model ready!")
    except Exception as e:
        import traceback
        _startup_error = traceback.format_exc()
        print(f"❌ Startup failed: {e}\n{_startup_error}")


def generate_heatmap(image: Image.Image, attn_map: np.ndarray,
                     condition_name: str, request_id: str,
                     server_url: str) -> str:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    try:
        from scipy.ndimage import zoom as nd_zoom, gaussian_filter
        am  = attn_map.reshape(14, 14)
        aup = nd_zoom(am, 448/14, order=3)
        aup = gaussian_filter(aup, sigma=8)
        aup = np.clip(aup, 0, None)
    except Exception:
        aup = np.array(
            Image.fromarray(attn_map.reshape(14, 14).astype(np.float32))
            .resize((448, 448), resample=Image.BICUBIC),
            dtype=np.float32)

    if aup.max() > aup.min():
        aup = (aup - aup.min()) / (aup.max() - aup.min())

    img448 = np.array(image.resize((448, 448)))

    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
    fig.patch.set_facecolor('#0d1117')
    ax.imshow(img448)
    ax.imshow(aup, cmap='jet', alpha=0.5)
    ax.set_title(condition_name.replace('_', ' '),
                 color='white', fontsize=12, fontweight='bold')
    ax.axis('off')
    plt.tight_layout()

    filename = f"{request_id}_{condition_name}.png"
    filepath = os.path.join(RESULTS_DIR, filename)
    fig.savefig(filepath, dpi=100, bbox_inches='tight', facecolor='#0d1117')
    plt.close(fig)

    return f"{server_url}/results/{filename}"


@app.post("/analyze")
async def analyze(
    file: UploadFile = File(...),
    ai_report: str = Form(""),
):
    report_text = ai_report.strip()

    if not report_text:
        return JSONResponse(
            status_code=400,
            content={"error": "ai_report field is required and cannot be empty"})

    try:
        # Image loading is now inside try/except so errors surface properly
        image_bytes = await file.read()
        if not image_bytes:
            return JSONResponse(
                status_code=400,
                content={"error": "Uploaded file is empty"})

        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')

        from inference.pipeline import run_full_pipeline, CONDITIONS
        from inference.model import device

        result = run_full_pipeline(image, report_text)

        server_url = API_BASE_URL
        request_id = str(uuid.uuid4())[:8]

        all_attn     = result.pop('all_attn')
        sentences    = result.pop('sentences')
        all_chexbert = result.pop('all_chexbert')

        heatmaps = {}
        active_names = [c['name'] for c in result['conditions']][:4]

        # Build lookup: condition → sentence index that drove its verdict
        # Pipeline stores _best_si per condition; fall back to 0 if missing.
        best_si_lookup = {c['name']: c.get('_best_si', 0)
                          for c in result['conditions']}

        if all_attn and len(all_attn) > 0:
            for cond in active_names:
                ci       = CONDITIONS.index(cond)
                best_si  = best_si_lookup.get(cond, 0)
                # Clamp in case sentence count changed unexpectedly
                best_si  = min(best_si, len(all_attn) - 1)
                attn_map = all_attn[best_si][ci]
                url      = generate_heatmap(
                    image, attn_map, cond, request_id, server_url)
                heatmaps[cond] = url

        task2 = {}
        for cond in result['conditions']:
            task2[cond['name']] = {
                'xray_present': cond['xray_present'],
                'confidence':   cond['confidence'],
            }

        return JSONResponse(content={
            "task1_elrrs":         result['elrrs'],
            "task1_conditions":    result['conditions'],
            "task2_xray_findings": task2,
            "task3_heatmaps":      heatmaps,
            "not_mentioned":       result['not_mentioned'],
            "sentences_analyzed":  len(sentences),
            "request_id":          request_id,
        })

    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        print(f"❌ /analyze error: {e}\n{tb}")
        return JSONResponse(
            status_code=500,
            content={"error": str(e), "traceback": tb})


@app.get("/health")
def health():
    if not _startup_ok:
        return JSONResponse(
            status_code=503,
            content={
                "status": "starting" if not _startup_error else "error",
                "model":  "RadGuard V11",
                "ready":  False,
                "startup_error": _startup_error or "Still initializing..."
            })
    return {"status": "ok", "model": "RadGuard V11", "ready": True}


@app.get("/debug")
def debug():
    """Diagnostic endpoint — shows exactly what is loaded and what paths exist."""
    import sys
    chexbert_ckpt = os.environ.get("CHEXBERT_CKPT", "/app/CheXbert/src/chexbert.pth")
    chexbert_dir  = os.environ.get("CHEXBERT_DIR", "/app/CheXbert")

    try:
        from inference.model import _model, _tokenizer, device, MODEL_PATH
        model_loaded     = _model is not None
        tokenizer_loaded = _tokenizer is not None
        model_path_used  = MODEL_PATH
    except Exception as e:
        model_loaded = tokenizer_loaded = False
        model_path_used = str(e)
        device = "unknown"

    return {
        "startup_ok":        _startup_ok,
        "startup_error":     _startup_error or None,
        "model_loaded":      model_loaded,
        "tokenizer_loaded":  tokenizer_loaded,
        "device":            str(device),
        "model_path":        model_path_used,
        "chexbert_ckpt_exists": os.path.exists(chexbert_ckpt),
        "chexbert_dir_exists":  os.path.exists(chexbert_dir),
        "chexbert_ckpt_path":   chexbert_ckpt,
        "results_dir":       RESULTS_DIR,
        "api_base_url":      API_BASE_URL,
        "python":            sys.version,
    }


@app.get("/")
def root():
    return {"message": "RadGuard AI Engine — use /analyze endpoint"}