app.py

import os
import re
import uuid
from pathlib import Path
from typing import Optional

import requests
import uvicorn
from fastapi import Cookie, FastAPI, File, Form, HTTPException, Response, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from agents.agent_decision import (
    load_report_from_disk,
    medical_agent_instance,
    reports_db,
)
from config import Config


config = Config()

app = FastAPI(
    title="AI-MD Integrated Multi-Agent Medical API",
    version="5.0",
    description="One API for medical chat, anatomical-layer triage, image analysis, RAG, voice transcription, and speech.",
)

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

UPLOAD_DIR = Path("uploads/backend")
REPORTS_DIR = Path("uploads/backend/reports")
SKIN_LESION_OUTPUT_DIR = Path("uploads/skin_lesion_output")
SPEECH_DIR = Path("uploads/speech")

for directory in [UPLOAD_DIR, REPORTS_DIR, SKIN_LESION_OUTPUT_DIR, SPEECH_DIR]:
    directory.mkdir(parents=True, exist_ok=True)

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

ALLOWED_IMAGE_EXTENSIONS = {"png", "jpg", "jpeg", "webp"}
ALLOWED_AUDIO_EXTENSIONS = {"webm", "wav", "mp3", "m4a", "ogg"}
_image_agent = None


class StartRequest(BaseModel):
    body_part: str
    chief_complaint: Optional[str] = None


class DiagnosticRequest(BaseModel):
    session_id: str
    body_part: str
    message: str


class SpeechRequest(BaseModel):
    text: str
    voice_id: Optional[str] = None


def secure_filename(filename: str) -> str:
    filename = os.path.basename(filename or "upload")
    filename = re.sub(r"[^A-Za-z0-9_.-]+", "_", filename).strip("._")
    return filename or "upload"


def allowed_file(filename: str, allowed: set[str]) -> bool:
    return "." in filename and filename.rsplit(".", 1)[1].lower() in allowed


def get_image_agent():
    global _image_agent
    if _image_agent is None:
        from agents.image_analysis_agent import ImageAnalysisAgent

        _image_agent = ImageAnalysisAgent(config)
    return _image_agent


def save_upload(upload: UploadFile, folder: Path, allowed: set[str], max_mb: int) -> Path:
    if not upload.filename or not allowed_file(upload.filename, allowed):
        raise HTTPException(status_code=400, detail="Unsupported file type.")

    content = upload.file.read()
    if not content:
        raise HTTPException(status_code=400, detail="Empty uploaded file.")

    if len(content) > max_mb * 1024 * 1024:
        raise HTTPException(status_code=413, detail=f"File too large. Max: {max_mb}MB.")

    filename = f"{uuid.uuid4()}_{secure_filename(upload.filename)}"
    path = folder / filename
    path.write_bytes(content)
    return path


def run_image_analysis(image_path: Path, image_type: str, text: str = "") -> dict:
    agent = get_image_agent()
    normalized_type = (image_type or "auto").strip().lower()

    if normalized_type == "auto":
        try:
            classifier_result = str(agent.analyze_image(str(image_path)))
        except Exception as exc:
            classifier_result = f"auto-classification unavailable: {exc}"

        lower = classifier_result.lower()
        if "xray" in lower or "x-ray" in lower or "chest" in lower:
            normalized_type = "chest_xray"
        elif "skin" in lower or "lesion" in lower:
            normalized_type = "skin_lesion"
        else:
            return {
                "agent": "IMAGE_CLASSIFIER",
                "image_type": "auto",
                "result": classifier_result,
                "medical_note": "Image type was not confidently routed to a specialist image agent.",
            }

    if normalized_type in {"chest", "xray", "x-ray", "chest_xray"}:
        result = agent.classify_chest_xray(str(image_path))
        return {
            "agent": "CHEST_XRAY_AGENT",
            "image_type": "chest_xray",
            "result": result,
            "medical_note": "AI image output must be reviewed by a qualified clinician.",
        }

    if normalized_type in {"skin", "lesion", "skin_lesion"}:
        result = agent.segment_skin_lesion(str(image_path))
        mask_url = None
        mask_path = Path(config.medical_cv.skin_lesion_segmentation_output_path)
        if mask_path.exists():
            mask_url = "/" + str(mask_path).replace("\\", "/")
        return {
            "agent": "SKIN_LESION_AGENT",
            "image_type": "skin_lesion",
            "result": result,
            "mask_url": mask_url,
            "medical_note": "AI image output must be reviewed by a qualified clinician.",
        }

    raise HTTPException(
        status_code=400,
        detail="image_type must be auto, chest_xray, or skin_lesion.",
    )


@app.get("/")
def root():
    return {
        "status": "ok",
        "service": "AI-MD Integrated Multi-Agent Medical API",
        "docs": "/docs",
        "endpoints": [
            "GET /health",
            "POST /api/v1/start",
            "POST /api/v1/chat",
            "POST /api/v1/image/analyze",
            "POST /api/v1/multimodal-chat",
            "POST /api/v1/transcribe",
            "POST /api/v1/generate-speech",
            "GET /api/v1/report/{session_id}",
        ],
    }


@app.get("/health")
def health_check():
    return {"status": "healthy"}


@app.post("/api/v1/start")
def start_diagnostic_session(request: StartRequest):
    try:
        return medical_agent_instance.start_session(
            body_part=request.body_part,
            chief_complaint=request.chief_complaint,
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.post("/api/v1/chat")
def process_diagnostic_chat(request: DiagnosticRequest):
    try:
        return medical_agent_instance.process(
            session_id=request.session_id,
            body_part=request.body_part,
            user_message=request.message,
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.get("/api/v1/report/{session_id}")
def get_final_report(session_id: str):
    entry = reports_db.get(session_id)

    if not entry:
        saved = load_report_from_disk(session_id)
        if saved:
            entry = {"status": "COMPLETED", "report": saved}
            reports_db[session_id] = entry

    if not entry:
        raise HTTPException(status_code=404, detail="Session not found or report is not ready.")

    if entry.get("status") == "COMPLETED" or entry.get("report"):
        medical_agent_instance.cleanup_session(session_id)
        return {
            "status": "REPORT",
            "session_status": "COMPLETED",
            "report": entry.get("report", {}),
        }

    return entry


@app.post("/api/v1/image/analyze")
async def analyze_image(
    image: UploadFile = File(...),
    image_type: str = Form("auto"),
    text: str = Form(""),
    session_id: Optional[str] = Form(None),
    body_part: Optional[str] = Form(None),
):
    path = save_upload(image, UPLOAD_DIR, ALLOWED_IMAGE_EXTENSIONS, config.api.max_image_upload_size)
    try:
        image_result = run_image_analysis(path, image_type=image_type, text=text)
        diagnostic_update = None
        if session_id and body_part:
            message = (
                f"Image analysis result from {image_result['agent']} "
                f"for {image_result['image_type']}: {image_result.get('result')}. "
                f"Patient context: {text}"
            )
            diagnostic_update = medical_agent_instance.process(
                session_id=session_id,
                body_part=body_part,
                user_message=message,
            )
        return {
            "status": "success",
            "session_id": session_id,
            "image": image_result,
            "diagnostic_update": diagnostic_update,
        }
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.post("/api/v1/upload")
async def upload_image_legacy(
    response: Response,
    image: UploadFile = File(...),
    text: str = Form(""),
    image_type: str = Form("auto"),
    session_id: Optional[str] = Cookie(None),
    body_part: Optional[str] = Form(None),
):
    result = await analyze_image(
        image=image,
        image_type=image_type,
        text=text,
        session_id=session_id,
        body_part=body_part,
    )
    if result.get("session_id"):
        response.set_cookie(key="session_id", value=result["session_id"])
    return result


@app.post("/api/v1/multimodal-chat")
async def multimodal_chat(
    session_id: str = Form(...),
    body_part: str = Form(...),
    message: str = Form(""),
    image: Optional[UploadFile] = File(None),
    image_type: str = Form("auto"),
):
    image_result = None
    combined_message = message

    if image is not None:
        path = save_upload(image, UPLOAD_DIR, ALLOWED_IMAGE_EXTENSIONS, config.api.max_image_upload_size)
        image_result = run_image_analysis(path, image_type=image_type, text=message)
        combined_message = (
            f"{message}\n\nAttached image analysis result: "
            f"{image_result['agent']} / {image_result['image_type']} -> {image_result.get('result')}"
        )

    diagnostic_response = medical_agent_instance.process(
        session_id=session_id,
        body_part=body_part,
        user_message=combined_message,
    )

    return {
        "status": "success",
        "session_id": session_id,
        "image": image_result,
        "diagnostic_response": diagnostic_response,
    }


@app.post("/api/v1/transcribe")
async def transcribe_audio(audio: UploadFile = File(...)):
    if not config.speech.eleven_labs_api_key:
        raise HTTPException(status_code=503, detail="ELEVEN_LABS_API_KEY is not configured.")

    path = save_upload(audio, SPEECH_DIR, ALLOWED_AUDIO_EXTENSIONS, max_mb=25)
    try:
        from elevenlabs.client import ElevenLabs

        client = ElevenLabs(api_key=config.speech.eleven_labs_api_key)
        with path.open("rb") as audio_file:
            transcription = client.speech_to_text.convert(
                file=audio_file,
                model_id="scribe_v1",
                tag_audio_events=True,
                language_code="eng",
                diarize=True,
            )
        return {"status": "success", "transcript": transcription.text}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.post("/api/v1/generate-speech")
async def generate_speech(request: SpeechRequest):
    if not config.speech.eleven_labs_api_key:
        raise HTTPException(status_code=503, detail="ELEVEN_LABS_API_KEY is not configured.")

    voice_id = request.voice_id or config.speech.eleven_labs_voice_id
    url = f"https://api.elevenlabs.io/v1/text-to-speech/{voice_id}/stream"
    headers = {
        "Accept": "audio/mpeg",
        "Content-Type": "application/json",
        "xi-api-key": config.speech.eleven_labs_api_key,
    }
    payload = {
        "text": request.text,
        "model_id": "eleven_monolingual_v1",
        "voice_settings": {"stability": 0.5, "similarity_boost": 0.5},
    }

    response = requests.post(url, headers=headers, json=payload, timeout=120)
    if response.status_code != 200:
        raise HTTPException(status_code=500, detail=response.text)

    output_path = SPEECH_DIR / f"{uuid.uuid4()}.mp3"
    output_path.write_bytes(response.content)
    return FileResponse(path=output_path, media_type="audio/mpeg", filename="generated_speech.mp3")


@app.exception_handler(413)
async def request_entity_too_large(request, exc):
    return JSONResponse(
        status_code=413,
        content={"status": "error", "response": f"File too large. Max: {config.api.max_image_upload_size}MB"},
    )


if __name__ == "__main__":
    uvicorn.run(app, host=config.api.host, port=config.api.port)