main.py

"""
main.py — FastAPI backend for Radiology AI Report Generation
Vision is ENABLED by default (BLIP-2 reads the actual image).
"""

import io
import json
import logging
import os
import sys
import uuid
from datetime import datetime
from pathlib import Path
from typing import Optional

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, Response

# Database
import importlib.util as _ilu
_db_spec = _ilu.spec_from_file_location("backend.database",
    __file__.replace("main.py","database.py"))
_db_mod = _ilu.module_from_spec(_db_spec)
_db_spec.loader.exec_module(_db_mod)
db = _db_mod

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(name)s — %(message)s",
)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="Radiology AI Report Generation System",
    version="2.0.0",
    description="AI-powered radiology report generation with BLIP-2 vision + Phi-2 LLM",
)

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

_report_store: dict = {}
_models_loaded   = False
_slm_model       = None
_slm_tokenizer   = None
_vision_model    = None
_vision_processor = None

# Vision is ON by default — set ENABLE_VISION_CAPTION=false to disable
ENABLE_VISION_CAPTION = False


# ═══════════════════════════════════════════════════════════════════════════
# Module loader helper (avoids import path issues on Windows)
# ═══════════════════════════════════════════════════════════════════════════

def _load_module(name, filepath):
    import importlib.util
    if name in sys.modules:
        return sys.modules[name]
    spec = importlib.util.spec_from_file_location(name, filepath)
    mod  = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod


_BACKEND = Path(__file__).resolve().parent


def _ml():
    return _load_module("backend.model_loader",    _BACKEND / "model_loader.py")
def _pp():
    return _load_module("backend.preprocessing",   _BACKEND / "preprocessing.py")
def _rg():
    return _load_module("backend.report_generator",_BACKEND / "report_generator.py")
def _vm():
    return _load_module("backend.validation_module",_BACKEND / "validation_module.py")
def _mc():
    return _load_module("backend.medical_coding",  _BACKEND / "medical_coding.py")


# ═══════════════════════════════════════════════════════════════════════════
# Model loading
# ═══════════════════════════════════════════════════════════════════════════

def _ensure_models_loaded():
    global _models_loaded, _slm_model, _slm_tokenizer, _vision_model, _vision_processor
    if _models_loaded:
        return

    logger.info("Loading AI models...")
    ml = _ml()

    # Always load language model
    _slm_model, _slm_tokenizer = ml.load_slm()

    # Load vision model (BLIP-2) if enabled
    if ENABLE_VISION_CAPTION:
        logger.info("Loading BLIP-2 vision model (this takes ~30s on first run)...")
        try:
            _vision_model, _vision_processor = ml.load_vision_encoder()
            logger.info("BLIP-2 vision model loaded — image content will be analysed.")
        except Exception as exc:
            logger.warning(f"BLIP-2 failed to load: {exc}. Reports will use form-based captions only.")

    _models_loaded = True
    logger.info("All models loaded.")


@app.on_event("startup")
async def startup_event():
    try:
        _ensure_models_loaded()
    except Exception as exc:
        logger.error(f"Startup model loading failed: {exc}")


# ═══════════════════════════════════════════════════════════════════════════
# Health & Info
# ═══════════════════════════════════════════════════════════════════════════

@app.get("/api/health", tags=["System"])
def health_check():
    import torch
    return {
        "status": "healthy",
        "models_loaded": _models_loaded,
        "vision_enabled": ENABLE_VISION_CAPTION,
        "vision_model_loaded": _vision_model is not None,
        "device": "cuda" if torch.cuda.is_available() else "cpu",
        "timestamp": datetime.utcnow().isoformat(),
    }


@app.get("/api/models", tags=["System"])
def list_models():
    import torch
    return {
        "slm_model": os.getenv("SLM_MODEL_ID", "microsoft/phi-2"),
        "vision_model": os.getenv("VISION_MODEL_ID", "Salesforce/blip2-opt-2.7b"),
        "vision_enabled": ENABLE_VISION_CAPTION,
        "vision_loaded": _vision_model is not None,
        "device": "cuda" if torch.cuda.is_available() else "cpu",
        "models_loaded": _models_loaded,
    }


# ═══════════════════════════════════════════════════════════════════════════
# Main report generation endpoint
# ═══════════════════════════════════════════════════════════════════════════

@app.post("/api/generate-report", tags=["Report"])
async def generate_report(
    image: UploadFile = File(...),
    patient_name: str = Form(""),
    patient_age: int = Form(...),
    patient_gender: str = Form(...),
    symptoms: str = Form(...),
    clinical_history: str = Form(...),
    modality: str = Form(...),
    body_region: str = Form(...),
    view: Optional[str] = Form(""),
):
    # 1. Read image bytes
    file_bytes = await image.read()
    filename   = image.filename or "upload.jpg"
    logger.info(f"Image received: {filename} ({len(file_bytes)/1024:.1f} KB)")

    # 2. Preprocess image
    try:
        pp = _pp()
        preprocessor = pp.RadiologyPreprocessor(
            modality=modality.lower(),
            body_region=body_region.lower(),
        )
        pil_image = preprocessor.preprocess_pil(file_bytes, filename)
    except Exception as exc:
        logger.error(f"Preprocessing error: {exc}")
        raise HTTPException(status_code=422, detail=f"Image preprocessing failed: {exc}")

    # 3. Load models
    try:
        _ensure_models_loaded()
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"Model loading failed: {exc}")

    # 4. Generate image caption using BLIP-2 (reads actual image pixels)
    image_caption = (
        f"{modality} image of the {body_region}. "
        f"Patient presents with {symptoms}."
    )
    if ENABLE_VISION_CAPTION and _vision_model is not None:
        try:
            ml = _ml()
            _, blip_caption = ml.extract_visual_features(
                pil_image, _vision_model, _vision_processor
            )
            if blip_caption and len(blip_caption) > 20:
                image_caption = blip_caption
                logger.info(f"BLIP-2 caption used: {image_caption[:120]}")
        except Exception as exc:
            logger.warning(f"Vision captioning failed: {exc} — using fallback caption")

    # 5. Generate report using Phi-2
    try:
        rg        = _rg()
        patient   = rg.PatientInfo(
            age=patient_age, gender=patient_gender,
            symptoms=symptoms, clinical_history=clinical_history,
        )
        imaging   = rg.ImagingInfo(
            modality=modality, body_region=body_region, view=view or ""
        )
        generator = rg.ReportGenerator(
            slm_model=_slm_model, slm_tokenizer=_slm_tokenizer
        )
        report    = generator.generate(
            patient=patient, imaging=imaging, image_caption=image_caption
        )
        logger.info(f"Report generated in {report.generation_time_seconds:.1f}s")
    except Exception as exc:
        logger.error(f"Report generation error: {exc}")
        raise HTTPException(status_code=500, detail=f"Report generation failed: {exc}")

    # 6. Validate
    vm         = _vm()
    validation = vm.validate_report(report)
    report.confidence_score = validation.confidence_score

    # 7. Medical coding (UMLS / ICD-10 / SNOMED CT)
    mc            = _mc()
    coding_result = mc.extract_medical_codes(
        report_text=report.raw_text,
        sections={
            "findings":        report.findings,
            "impression":      report.impression,
            "recommendations": report.recommendations,
        },
    )

    # 8. Build and store response
    report_id = str(uuid.uuid4())
    # Clean patient name
    patient_name_clean = patient_name.strip() if patient_name.strip() else "Unknown Patient"

    payload   = {
        "report_id":    report_id,
        "timestamp":    datetime.utcnow().isoformat(),
        "image_caption": image_caption,
        "patient_name": patient_name_clean,
        "patient": {
            "name":             patient_name_clean,
            "age":              patient_age,
            "gender":           patient_gender,
            "symptoms":         symptoms,
            "clinical_history": clinical_history,
        },
        "imaging": {
            "modality":    modality,
            "body_region": body_region,
            "view":        view,
        },
        "report":       report.to_dict(),
        "report_text":  report.to_text(),
        "validation":   validation.to_dict(),
        "medical_coding": coding_result.to_dict(),
    }
    _report_store[report_id] = payload

    # Save to SQLite database
    try:
        db.save_report(payload, patient_name_clean)
    except Exception as db_exc:
        logger.warning(f"DB save failed (non-critical): {db_exc}")

    logger.info(
        f"Report stored: {report_id[:8]} | "
        f"confidence: {report.confidence_score:.2f} | "
        f"codes: {len(coding_result.codes)} | "
        f"diagnosis: {coding_result.primary_diagnosis_label}"
    )
    return JSONResponse(content=payload)


# ═══════════════════════════════════════════════════════════════════════════
# Report retrieval & export
# ═══════════════════════════════════════════════════════════════════════════

@app.get("/api/report/{report_id}", tags=["Report"])
def get_report(report_id: str):
    if report_id not in _report_store:
        raise HTTPException(status_code=404, detail="Report not found.")
    return JSONResponse(content=_report_store[report_id])


@app.get("/api/export/json/{report_id}", tags=["Export"])
def export_json(report_id: str):
    if report_id not in _report_store:
        raise HTTPException(status_code=404, detail="Report not found.")
    return Response(
        content=json.dumps(_report_store[report_id], indent=2).encode(),
        media_type="application/json",
        headers={"Content-Disposition": f'attachment; filename="report_{report_id[:8]}.json"'},
    )


@app.get("/api/export/pdf/{report_id}", tags=["Export"])
def export_pdf(report_id: str):
    if report_id not in _report_store:
        raise HTTPException(status_code=404, detail="Report not found.")
    try:
        pdf_bytes = _generate_pdf(_report_store[report_id])
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"PDF generation failed: {exc}")
    return Response(
        content=pdf_bytes,
        media_type="application/pdf",
        headers={"Content-Disposition": f'attachment; filename="report_{report_id[:8]}.pdf"'},
    )


# ═══════════════════════════════════════════════════════════════════════════
# PDF generator
# ═══════════════════════════════════════════════════════════════════════════



# ═══════════════════════════════════════════════════════════════════════════
# Database endpoints
# ═══════════════════════════════════════════════════════════════════════════

@app.get("/api/reports", tags=["Database"])
def list_all_reports(limit: int = 50, offset: int = 0):
    """List all stored reports from database."""
    try:
        reports = db.get_all_reports(limit=limit, offset=offset)
        return JSONResponse(content={"reports": reports, "count": len(reports)})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/reports/patient/{patient_name}", tags=["Database"])
def get_patient_reports(patient_name: str):
    """Get all reports for a specific patient."""
    try:
        reports = db.get_reports_by_patient(patient_name)
        return JSONResponse(content={"patient": patient_name,
                                      "reports": reports,
                                      "count": len(reports)})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/reports/search", tags=["Database"])
def search_patients(q: str = ""):
    """Search patients by name."""
    try:
        results = db.search_patients(q)
        return JSONResponse(content={"results": results})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/stats", tags=["Database"])
def get_stats():
    """Get database statistics."""
    try:
        return JSONResponse(content=db.get_db_stats())
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/api/reports/{report_id}", tags=["Database"])
def delete_report(report_id: str):
    """Delete a report by ID."""
    try:
        success = db.delete_report(report_id)
        if not success:
            raise HTTPException(status_code=404, detail="Report not found.")
        return JSONResponse(content={"deleted": report_id})
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

def _generate_pdf(payload: dict) -> bytes:
    from reportlab.lib import colors
    from reportlab.lib.enums import TA_CENTER, TA_JUSTIFY
    from reportlab.lib.pagesizes import A4
    from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
    from reportlab.lib.units import cm
    from reportlab.platypus import (HRFlowable, Paragraph, SimpleDocTemplate,
                                     Spacer, Table, TableStyle)

    buf    = io.BytesIO()
    doc    = SimpleDocTemplate(buf, pagesize=A4,
                               rightMargin=2*cm, leftMargin=2*cm,
                               topMargin=2.5*cm, bottomMargin=2*cm)
    styles = getSampleStyleSheet()
    BLUE   = colors.HexColor("#1a4a8a")
    LIGHT  = colors.HexColor("#e8f0fb")

    story = []
    story.append(Paragraph("RADIOLOGY REPORT",
        ParagraphStyle("T", parent=styles["Title"], textColor=BLUE,
                       fontSize=20, spaceAfter=4, fontName="Helvetica-Bold")))

    report    = payload.get("report", {})
    patient   = payload.get("patient", {})
    imaging   = payload.get("imaging", {})
    validation = payload.get("validation", {})
    report_id = payload.get("report_id", "N/A")
    timestamp = payload.get("timestamp", "")

    story.append(Paragraph(
        f"AI-Generated Report  |  {timestamp[:10]}  |  ID: {report_id[:8]}",
        ParagraphStyle("S", parent=styles["Normal"], textColor=colors.HexColor("#555"),
                       fontSize=10, spaceAfter=12, alignment=TA_CENTER)))
    story.append(HRFlowable(width="100%", thickness=2, color=BLUE))
    story.append(Spacer(1, 0.4*cm))

    cell_style = ParagraphStyle("cell", parent=styles["Normal"], fontSize=9,
                                leading=12)
    lbl_style  = ParagraphStyle("lbl", parent=styles["Normal"], fontSize=9,
                                leading=12, fontName="Helvetica-Bold")
    symptoms_text = str(patient.get("symptoms", "—"))

    tbl = Table([
        [Paragraph("Patient Name", lbl_style), Paragraph(str(payload.get("patient_name","—")), cell_style),
         Paragraph("Report ID",   lbl_style), Paragraph(str(report_id[:8]), cell_style)],
        [Paragraph("Patient Age", lbl_style), Paragraph(str(patient.get("age","—")), cell_style),
         Paragraph("Modality",    lbl_style), Paragraph(str(imaging.get("modality","—")), cell_style)],
        [Paragraph("Gender",      lbl_style), Paragraph(str(patient.get("gender","—")), cell_style),
         Paragraph("Body Region", lbl_style), Paragraph(str(imaging.get("body_region","—")), cell_style)],
        [Paragraph("View",        lbl_style), Paragraph(str(imaging.get("view","—") or "—"), cell_style),
         Paragraph("",            cell_style), Paragraph("", cell_style)],
        [Paragraph("Symptoms",    lbl_style), Paragraph(symptoms_text, cell_style),
         Paragraph("",            cell_style), Paragraph("", cell_style)],
    ], colWidths=[3*cm, 7*cm, 3*cm, 4*cm])
    tbl.setStyle(TableStyle([
        ("FONTSIZE",      (0,0),(-1,-1), 9),
        ("BACKGROUND",    (0,0),(-1,0),  LIGHT),
        ("BACKGROUND",    (0,2),(-1,2),  LIGHT),
        ("GRID",          (0,0),(-1,-1), 0.5, colors.HexColor("#ccc")),
        ("VALIGN",        (0,0),(-1,-1), "TOP"),
        ("TOPPADDING",    (0,0),(-1,-1), 5),
        ("BOTTOMPADDING", (0,0),(-1,-1), 5),
        ("SPAN",          (1,2),(3,2)),
        ("SPAN",          (1,3),(3,3)),
    ]))
    story.append(tbl)
    story.append(Spacer(1, 0.5*cm))
    story.append(HRFlowable(width="100%", thickness=0.5, color=colors.HexColor("#aaa")))

    hdr  = ParagraphStyle("H", parent=styles["Heading2"], textColor=BLUE,
                          fontSize=11, fontName="Helvetica-Bold", spaceBefore=12, spaceAfter=4)
    body = ParagraphStyle("B", parent=styles["Normal"], fontSize=10,
                          leading=15, spaceAfter=6, alignment=TA_JUSTIFY)

    for header, key in [
        ("Examination Details", "examination_details"),
        ("Clinical Information", "clinical_information"),
        ("Findings",            "findings"),
        ("Impression",          "impression"),
        ("Recommendations",     "recommendations"),
    ]:
        story.append(Paragraph(header.upper(), hdr))
        story.append(Paragraph(str(report.get(key,"—") or "—"), body))

    story.append(Spacer(1, 0.5*cm))
    story.append(HRFlowable(width="100%", thickness=1, color=BLUE))
    confidence = validation.get("confidence_score", 0)
    story.append(Paragraph(
        f"CONFIDENCE SCORE: {confidence:.2f} / 1.00",
        ParagraphStyle("C", parent=styles["Normal"], fontSize=12,
                       fontName="Helvetica-Bold", textColor=BLUE,
                       alignment=TA_CENTER, spaceBefore=10)))

    story.append(Spacer(1, 1*cm))
    story.append(HRFlowable(width="100%", thickness=0.5, color=colors.HexColor("#aaa")))
    story.append(Paragraph(
        "This report was generated by an AI system. "
        "Clinical correlation by a licensed radiologist is mandatory.",
        ParagraphStyle("D", parent=styles["Normal"], fontSize=8,
                       textColor=colors.HexColor("#888"),
                       alignment=TA_CENTER, spaceBefore=6)))
    doc.build(story)
    return buf.getvalue()