backend/main.py

"""FastAPI wrapper for the MOFScreen-Agent Python pipeline.

Run locally:
    uvicorn backend.main:app --host 0.0.0.0 --port 8000
"""
from __future__ import annotations

import asyncio
import math
import os
import sys
import tempfile
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from agent.conversation import run_conversational_screening, stream_llm_result_response
from tools.six_step_screening import build_result_row


app = FastAPI(title="MOFScreen-Agent API", version="1.0.0")

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


class RunJobRequest(BaseModel):
    llm_api_key: str | None = None


def _supabase_admin():
    url = os.getenv("SUPABASE_URL") or os.getenv("NEXT_PUBLIC_SUPABASE_URL")
    key = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
    if not url or not key:
        raise RuntimeError("Supabase is not configured. Set SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY.")

    from supabase import create_client

    return create_client(url, key)


def _clean_json(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(k): _clean_json(v) for k, v in value.items()}
    if isinstance(value, list | tuple):
        return [_clean_json(v) for v in value]
    if isinstance(value, np.ndarray):
        return _clean_json(value.tolist())
    if isinstance(value, np.generic):
        return _clean_json(value.item())
    if isinstance(value, float) and (math.isnan(value) or math.isinf(value)):
        return None
    if isinstance(value, Path):
        return str(value)
    return value


def _tox_display(value: Any, confidence: str | None) -> str:
    if value is None:
        return "N/A"
    suffix = ""
    if confidence == "low":
        suffix = " low"
    elif confidence == "medium":
        suffix = " med"
    try:
        return f"{float(value):.2f}{suffix}"
    except Exception:
        return str(value)


def _result_row(result: dict[str, Any], mof_id: str | None = None) -> dict[str, Any]:
    if isinstance(result.get("row"), dict):
        row = dict(result["row"])
        if mof_id:
            row["MOF ID"] = row.get("MOF ID") or mof_id
        return row
    if isinstance(result.get("six_step"), dict):
        row = build_result_row(result)
        if mof_id:
            row["MOF ID"] = row.get("MOF ID") or mof_id
        return row

    selective = result.get("results") or {}
    adsorption = selective.get("adsorption") or result.get("adsorption") or {}
    toxicity = selective.get("toxicity") or result.get("toxicity") or {}
    safety = selective.get("safety") or result.get("safety") or {}
    linker = selective.get("linker") or result.get("linker") or {}

    benzene = adsorption.get("benzene", {}).get("uptake_mg_g") if isinstance(adsorption.get("benzene"), dict) else None
    toluene = adsorption.get("toluene", {}).get("uptake_mg_g") if isinstance(adsorption.get("toluene"), dict) else None
    benzene = benzene if benzene is not None else adsorption.get("benzene_uptake_mg_g")
    toluene = toluene if toluene is not None else adsorption.get("toluene_uptake_mg_g")

    return {
        "MOF ID": mof_id or result.get("mof_id", "unknown"),
        "Metal": ", ".join(linker.get("metals", [])) if linker else "N/A",
        "Linker": linker.get("linker_name") or linker.get("linker_formula") or "N/A",
        "Linker SMILES": linker.get("linker_smiles") or "N/A",
        "Benzene (mg/g)": benzene,
        "Toluene (mg/g)": toluene,
        "LC50 fish (-log mol/L)": _tox_display(
            toxicity.get("LC50_Pimephales"),
            toxicity.get("LC50_Pimephales_confidence"),
        ),
        "LC50 daphnia (-log mol/L)": _tox_display(
            toxicity.get("LC50_Daphnia"),
            toxicity.get("LC50_Daphnia_confidence"),
        ),
        "IGC50 tetrahymena (-log mol/L)": _tox_display(
            toxicity.get("IGC50_Tetrahymena"),
            toxicity.get("IGC50_Tetrahymena_confidence"),
        ),
        "IBC50 vibrio (-log mol/L)": _tox_display(
            toxicity.get("IBC50_Vibrio"),
            toxicity.get("IBC50_Vibrio_confidence"),
        ),
        "Safety": safety.get("metal_tier", "N/A") if safety else "N/A",
        "Score": result.get("final_score"),
        "Decision Class": (result.get("decision_record") or {}).get("decision_class", "N/A"),
        "Recommendation": result.get("recommendation", "N/A"),
    }


def _frontend_result(result: dict[str, Any]) -> dict[str, Any]:
    """Return traceable but compact records for the web UI."""
    compact = dict(result)
    compact.pop("llm_api_key", None)
    scm = compact.get("scm_meta")
    if isinstance(scm, dict):
        compact["scm_meta"] = {
            key: value
            for key, value in scm.items()
            if key not in {"benzene_eigenvalues", "toluene_eigenvalues"}
        }
    return compact


def _sse(event: str, data: Any) -> str:
    return f"event: {event}\ndata: {json.dumps(_clean_json(data), ensure_ascii=False)}\n\n"


def _tool_summary(name: str, payload: dict[str, Any]) -> dict[str, Any]:
    return {
        "tool": name,
        "status": payload.get("status"),
        "step": payload.get("step"),
        "summary": payload.get("summary")
        or payload.get("reason")
        or payload.get("price_status")
        or payload.get("mof")
        or payload.get("linker_smiles")
        or payload.get("status"),
        "result": payload,
    }


def _run_files(
    file_specs: list[tuple[str, bytes]],
    user_request: str,
    mode: str,
    llm_provider: str,
    llm_api_key: str | None,
) -> dict[str, Any]:
    results: list[dict[str, Any]] = []
    rows: list[dict[str, Any]] = []

    for filename, content in file_specs:
        safe_name = Path(filename or "candidate.cif").name
        if not safe_name.lower().endswith(".cif"):
            safe_name = f"{Path(safe_name).stem or 'candidate'}.cif"

        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = str(Path(tmp_dir) / safe_name)
            Path(tmp_path).write_bytes(content)
            try:
                result = run_conversational_screening(
                    user_request=user_request,
                    cif_path=tmp_path,
                    llm_provider=llm_provider,
                    llm_api_key=llm_api_key,
                )
                mof_id = result.get("mof_id") or Path(filename or tmp_path).stem
                result["uploaded_filename"] = filename
                result["display_mof_id"] = mof_id
                rows.append(_result_row(result, mof_id=mof_id))
                results.append(_clean_json(_frontend_result(result)))
            except Exception as exc:
                failed = {
                    "uploaded_filename": filename,
                    "display_mof_id": Path(filename or "unknown").stem,
                    "errors": [str(exc)],
                    "warnings": [],
                }
                results.append(failed)
                rows.append({
                    "MOF ID": failed["display_mof_id"],
                    "Benzene (mg/g)": None,
                    "Toluene (mg/g)": None,
                    "Metal status": "N/A",
                    "Metals": "N/A",
                    "Linker": "N/A",
                    "Linker SMILES": "N/A",
                    "SA score": None,
                    "Mean toxicity": None,
                    "Price USD/g": None,
                    "Price USD/mmol": None,
                    "Gate status": "error",
                    "Recommendation": "error",
                    "Score": None,
                })

    return _clean_json({"rows": rows, "results": results, "count": len(results)})


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok", "service": "mofscreen-agent-api"}


@app.post("/screen")
async def screen(
    files: list[UploadFile] = File(...),
    user_request: str = Form("Run full safe-by-design MOF screening."),
    mode: str = Form("full_screening"),
    llm_provider: str = Form("rule_based"),
    llm_api_key: str | None = Form(None),
) -> dict[str, Any]:
    results: list[dict[str, Any]] = []
    rows: list[dict[str, Any]] = []

    for uploaded in files:
        safe_name = Path(uploaded.filename or "candidate.cif").name
        if not safe_name.lower().endswith(".cif"):
            safe_name = f"{Path(safe_name).stem or 'candidate'}.cif"

        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = str(Path(tmp_dir) / safe_name)
            Path(tmp_path).write_bytes(await uploaded.read())
            try:
                result = run_conversational_screening(
                    user_request=user_request,
                    cif_path=tmp_path,
                    llm_provider=llm_provider,
                    llm_api_key=llm_api_key,
                )
                mof_id = result.get("mof_id") or Path(uploaded.filename or tmp_path).stem
                result["uploaded_filename"] = uploaded.filename
                result["display_mof_id"] = mof_id
                rows.append(_result_row(result, mof_id=mof_id))
                results.append(_clean_json(_frontend_result(result)))
            except Exception as exc:
                failed = {
                    "uploaded_filename": uploaded.filename,
                    "display_mof_id": Path(uploaded.filename or "unknown").stem,
                    "errors": [str(exc)],
                    "warnings": [],
                }
                results.append(failed)
                rows.append({
                    "MOF ID": failed["display_mof_id"],
                    "Benzene (mg/g)": None,
                    "Toluene (mg/g)": None,
                    "Metal status": "N/A",
                    "Metals": "N/A",
                    "Linker": "N/A",
                    "Linker SMILES": "N/A",
                    "SA score": None,
                    "Mean toxicity": None,
                    "Price USD/g": None,
                    "Price USD/mmol": None,
                    "Gate status": "error",
                    "Recommendation": "error",
                    "Score": None,
                })

    return _clean_json({
        "rows": rows,
        "results": results,
        "count": len(results),
    })


@app.post("/chat/stream")
async def chat_stream(
    files: list[UploadFile] | None = File(None),
    user_request: str = Form("Run full MOF screening."),
    mode: str = Form("full_screening"),
    llm_provider: str = Form("rule_based"),
    llm_api_key: str | None = Form(None),
):
    file_specs: list[tuple[str, bytes]] = []
    for uploaded in files or []:
        file_specs.append((uploaded.filename or "candidate.cif", await uploaded.read()))

    async def events():
        if not file_specs:
            result = run_conversational_screening(
                user_request=user_request,
                cif_path=None,
                llm_provider=llm_provider,
                llm_api_key=llm_api_key,
                final_llm=False,
            )
            for chunk in stream_llm_result_response(user_request, result, llm_provider, llm_api_key):
                yield _sse("assistant_delta", chunk)
            yield _sse("final_result", {"rows": [], "results": [_frontend_result(result)], "count": 1})
            return

        rows: list[dict[str, Any]] = []
        results: list[dict[str, Any]] = []
        for filename, content in file_specs:
            safe_name = Path(filename or "candidate.cif").name
            if not safe_name.lower().endswith(".cif"):
                safe_name = f"{Path(safe_name).stem or 'candidate'}.cif"
            with tempfile.TemporaryDirectory() as tmp_dir:
                tmp_path = str(Path(tmp_dir) / safe_name)
                Path(tmp_path).write_bytes(content)
                yield _sse("workflow_start", {"filename": filename, "status": "running"})
                screening_task = asyncio.create_task(asyncio.to_thread(
                    run_conversational_screening,
                    user_request=user_request,
                    cif_path=tmp_path,
                    llm_provider=llm_provider,
                    llm_api_key=llm_api_key,
                    final_llm=False,
                ))
                while True:
                    try:
                        result = await asyncio.wait_for(asyncio.shield(screening_task), timeout=8)
                        break
                    except asyncio.TimeoutError:
                        yield _sse("workflow_heartbeat", {"filename": filename, "status": "running"})
                mof_id = result.get("mof_id") or Path(filename or tmp_path).stem
                result["uploaded_filename"] = filename
                result["display_mof_id"] = mof_id
                for name, payload in (result.get("six_step") or result.get("results") or {}).items():
                    if isinstance(payload, dict):
                        yield _sse("tool_start", {"tool": name})
                        yield _sse("tool_result", _tool_summary(name, payload))
                text_result = dict(result)
                text_result["assistant_message"] = None
                for chunk in stream_llm_result_response(user_request, text_result, llm_provider, llm_api_key):
                    yield _sse("assistant_delta", chunk)
                rows.append(_result_row(result, mof_id=mof_id))
                results.append(_clean_json(_frontend_result(result)))
        yield _sse("final_result", {"rows": rows, "results": results, "count": len(results)})

    return StreamingResponse(events(), media_type="text/event-stream")


@app.post("/jobs/{job_id}/run")
async def run_job(job_id: str, payload: RunJobRequest | None = None) -> dict[str, Any]:
    supabase = _supabase_admin()

    job_response = supabase.table("screening_jobs").select("*").eq("id", job_id).single().execute()
    job = job_response.data
    if not job:
        return {"error": f"Job not found: {job_id}"}

    supabase.table("screening_jobs").update({"status": "running", "error": None}).eq("id", job_id).execute()

    bucket = job.get("storage_bucket") or "mof-cifs"
    file_paths = job.get("file_paths") or []
    uploaded = job.get("uploaded_filenames") or []

    try:
        file_specs: list[tuple[str, bytes]] = []
        for index, storage_path in enumerate(file_paths):
            data = supabase.storage.from_(bucket).download(storage_path)
            filename = uploaded[index] if index < len(uploaded) else Path(storage_path).name
            file_specs.append((filename, data))

        output = _run_files(
            file_specs=file_specs,
            user_request=job.get("user_request") or "Run full safe-by-design MOF screening.",
            mode=job.get("mode") or "full_screening",
            llm_provider=job.get("llm_provider") or "rule_based",
            llm_api_key=payload.llm_api_key if payload else None,
        )

        update = {
            "status": "completed",
            "rows": output["rows"],
            "results": output["results"],
            "error": None,
            "completed_at": datetime.now(timezone.utc).isoformat(),
        }
        supabase.table("screening_jobs").update(update).eq("id", job_id).execute()
        return {"job_id": job_id, **output}
    except Exception as exc:
        message = str(exc)
        supabase.table("screening_jobs").update({"status": "failed", "error": message}).eq("id", job_id).execute()
        return {"job_id": job_id, "error": message, "rows": [], "results": []}