app.py

import os
import threading
import time
import uuid
from pathlib import Path
from typing import Optional

import uvicorn
from chainlit.utils import mount_chainlit
from fastapi import Cookie, FastAPI, File, Form, HTTPException, Response, UploadFile
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from werkzeug.utils import secure_filename

from agents.agent_decision_multi_agent import process_query, resume_review
from config import Config

config = Config()
os.environ.setdefault("CHAINLIT_AUTH_SECRET", config.auth.auth_secret)

app = FastAPI(title="SkinVL Dermoscopy Assistant Demo", version="1.0")

UPLOAD_FOLDER = Path("uploads/backend")

UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)

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

ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg"}


class QueryRequest(BaseModel):
    query: str


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


def _has_valid_image_signature(filename: str, content: bytes) -> bool:
    extension = filename.rsplit(".", 1)[-1].lower()
    if extension == "png":
        return content.startswith(b"\x89PNG\r\n\x1a\n")
    return content.startswith(b"\xff\xd8\xff")


async def _read_upload_with_limit(image: UploadFile, limit_bytes: int) -> bytes:
    return await image.read(limit_bytes + 1)


def _default_backend_metadata() -> dict:
    demo_mode = bool(config.medical_cv.demo_mode)
    return {
        "mode": "demo" if demo_mode else "runtime",
        "scenario": config.medical_cv.demo_scenario,
        "image_router": "mock_image_router" if demo_mode else "image_router",
        "skinvl": "mock_skinvl" if demo_mode else "transformers_skinvl",
        "rag": "demo_rag" if demo_mode else "qdrant_rag",
        "web_search": "enabled" if config.web_search.enabled else "disabled",
        "simulated_stages": ["image_routing", "skinvl", "rag"]
        if demo_mode
        else [],
    }


def cleanup_expired_uploads(now: Optional[float] = None) -> int:
    """Delete backend uploads only after the configured review retention window."""
    cutoff = (now if now is not None else time.time()) - (
        config.api.upload_retention_minutes * 60
    )
    deleted = 0
    for file_path in UPLOAD_FOLDER.iterdir():
        try:
            if file_path.is_file() and file_path.stat().st_mtime < cutoff:
                file_path.unlink()
                deleted += 1
        except OSError as exc:  # pragma: no cover - cleanup best effort
            print(f"Failed to remove expired upload {file_path}: {exc}")
    return deleted


def cleanup_old_uploads() -> None:
    """Periodically expire uploads while preserving files needed by HITL review."""
    interval_seconds = min(
        300,
        max(60, config.api.upload_retention_minutes * 30),
    )
    while True:
        try:
            cleanup_expired_uploads()
        except Exception as exc:  # pragma: no cover - background maintenance
            print(f"Error during upload cleanup: {exc}")
        time.sleep(interval_seconds)


upload_cleanup_thread = threading.Thread(target=cleanup_old_uploads, daemon=True)
upload_cleanup_thread.start()


def _build_http_result(result: dict, session_id: str) -> dict:
    payload = {
        "status": result["status"],
        "response": result["response"],
        "agent_outputs": result.get("agent_outputs", []),
        "confidence": result.get("confidence", 1.0),
        "needs_validation": result.get("needs_validation", False),
        "review_payload": result.get("review_payload"),
        "attachments": result.get("attachments", []),
        "session_id": session_id,
        "simulation": result.get("simulation", config.medical_cv.demo_mode),
        "backend_metadata": result.get(
            "backend_metadata", _default_backend_metadata()
        ),
        "image_routing": result.get("image_routing"),
        "image_findings": result.get("image_findings"),
        "retrieval_bundle": result.get("retrieval_bundle"),
        "sources": result.get("sources", []),
        "review_state": result.get("review_state"),
    }

    image_attachments = [
        attachment
        for attachment in result.get("attachments", [])
        if attachment.get("type") == "image" and attachment.get("url")
    ]
    if image_attachments:
        payload["result_image"] = image_attachments[0]["url"]

    return payload


@app.get("/")
async def index():
    return RedirectResponse(url=f"{config.ui.chainlit_path}/")


@app.get("/health")
def health_check():
    return {
        "status": "healthy",
        "service": "SkinVL Dermoscopy Assistant Demo",
        "simulation": bool(config.medical_cv.demo_mode),
        "demo_mode": bool(config.medical_cv.demo_mode),
        "demo_scenario": config.medical_cv.demo_scenario,
        "web_search_enabled": bool(config.web_search.enabled),
        "backend_metadata": _default_backend_metadata(),
        "upload_retention_minutes": config.api.upload_retention_minutes,
    }


@app.post("/chat")
def chat(
    request: QueryRequest,
    response: Response,
    session_id: Optional[str] = Cookie(None),
):
    thread_id = session_id or str(uuid.uuid4())
    response.set_cookie(key="session_id", value=thread_id)

    try:
        result = process_query(request.query, thread_id=thread_id)
        return _build_http_result(result, thread_id)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(
            status_code=500,
            detail="SkinVL demo could not complete the request.",
        ) from exc


@app.post("/upload")
async def upload_image(
    response: Response,
    image: UploadFile = File(...),
    text: str = Form(""),
    session_id: Optional[str] = Cookie(None),
):
    if not image.filename or not allowed_file(image.filename):
        return JSONResponse(
            status_code=400,
            content={
                "status": "error",
                "response": "Unsupported file type. Allowed formats: PNG, JPG, JPEG",
            },
        )

    limit_bytes = config.api.max_image_upload_size * 1024 * 1024
    file_content = await _read_upload_with_limit(image, limit_bytes)
    if len(file_content) > limit_bytes:
        return JSONResponse(
            status_code=413,
            content={
                "status": "error",
                "response": (
                    f"File too large. Maximum size allowed: {config.api.max_image_upload_size}MB"
                ),
            },
        )

    if not _has_valid_image_signature(image.filename, file_content):
        return JSONResponse(
            status_code=400,
            content={
                "status": "error",
                "response": "The uploaded file content is not a valid PNG or JPEG image.",
            },
        )

    thread_id = session_id or str(uuid.uuid4())
    response.set_cookie(key="session_id", value=thread_id)

    filename = secure_filename(f"{uuid.uuid4()}_{image.filename}")
    file_path = UPLOAD_FOLDER / filename
    file_path.write_bytes(file_content)

    try:
        result = process_query({"text": text, "image": str(file_path)}, thread_id=thread_id)
        return _build_http_result(result, thread_id)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(
            status_code=500,
            detail="SkinVL demo could not process the uploaded image.",
        ) from exc


@app.post("/validate")
def validate_medical_output(
    response: Response,
    validation_result: str = Form(...),
    comments: Optional[str] = Form(None),
    edited_response: Optional[str] = Form(None),
    session_id: Optional[str] = Cookie(None),
):
    thread_id = session_id or str(uuid.uuid4())
    response.set_cookie(key="session_id", value=thread_id)

    decision_type = validation_result.strip().lower()
    if decision_type == "yes":
        decision_type = "approve"
    elif decision_type == "no":
        decision_type = "reject"

    decision = {"type": decision_type}
    if comments:
        decision["comment"] = comments
    if edited_response:
        decision["edited_response"] = edited_response

    try:
        result = resume_review(thread_id=thread_id, decision=decision)
        return _build_http_result(result, thread_id)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(
            status_code=500,
            detail="The review could not be resumed. The draft remains unapproved.",
        ) from exc


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


mount_chainlit(app=app, target="chainlit_app.py", path=config.ui.chainlit_path)


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