#!/usr/bin/env python3
"""FastAPI research API for the MSC Aging Agent Corpus (Hugging Face Space)."""

from __future__ import annotations

import json
import os
from contextlib import asynccontextmanager
from typing import Any

import pandas as pd
from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel, Field

from msc_corpus import connect

DEFAULT_REVISION = os.environ.get("CORPUS_REVISION", "corpus-v2026.06.6")
CACHE_DIR = os.environ.get("CORPUS_CACHE_DIR", "/tmp/msc_corpus_cache")

corpus = None


@asynccontextmanager
async def lifespan(_app: FastAPI):
    global corpus
    corpus = connect(remote_only=True, revision=DEFAULT_REVISION, cache_dir=CACHE_DIR)
    yield


app = FastAPI(
    title="MSC Aging Corpus Research API",
    description=(
        "Query the Syndicate Laboratories MSC Aging Agent Corpus: biomarker search, "
        "sample-level expression, OLS statistical reruns, and APA citations. "
        "Data source: Hugging Face Dataset S4MPL3BI4S/msc-aging-agent-corpus."
    ),
    version="1.0.0",
    lifespan=lifespan,
)


class ResearchRequest(BaseModel):
    gene: str = Field(..., description="Gene symbol, e.g. NDRG1")
    dataset_id: str | None = Field(None, description="Optional GEO accession, e.g. GSE39540")
    species: str | None = Field(None, description="Optional species filter, e.g. Homo sapiens")
    fdr_only: bool = False
    actions: list[str] = Field(
        default=["search", "rerun", "cite"],
        description="Steps: search, expression, rerun, mechanism, cite",
    )


def _require_corpus():
    if corpus is None:
        raise HTTPException(status_code=503, detail="Corpus not loaded yet.")
    return corpus


def _records(frame: pd.DataFrame, limit: int = 500) -> dict[str, Any]:
    if frame.empty:
        return {"count": 0, "rows": []}
    trimmed = frame.head(limit)
    return {
        "count": int(len(frame)),
        "truncated": len(frame) > limit,
        "rows": json.loads(trimmed.to_json(orient="records")),
    }


@app.get("/")
def root() -> dict[str, str]:
    return {
        "service": "MSC Aging Corpus Research API",
        "curator": "James Utley, PhD · Syndicate Laboratories",
        "dataset": "S4MPL3BI4S/msc-aging-agent-corpus",
        "revision": DEFAULT_REVISION,
        "openapi_docs": "/docs",
        "research_endpoint": "POST /research",
    }


@app.get("/health")
def health() -> dict[str, str]:
    _require_corpus()
    return {"status": "ok", "revision": DEFAULT_REVISION}


@app.get("/connect")
def connection_manifest() -> dict[str, Any]:
    return _require_corpus().connection_info()


@app.get("/manifest")
def manifest(
    species: str | None = None,
    limit: int = Query(100, ge=1, le=500),
) -> dict[str, Any]:
    frame = _require_corpus().manifest()
    if species:
        frame = frame.loc[frame["species"] == species]
    return _records(frame, limit=limit)


@app.get("/search")
def search_gene(
    gene: str,
    species: str | None = None,
    dataset_id: str | None = None,
    fdr_only: bool = False,
    limit: int = Query(200, ge=1, le=1000),
) -> dict[str, Any]:
    frame = _require_corpus().search_gene(
        gene,
        species=species,
        dataset_id=dataset_id,
        fdr_only=fdr_only,
    )
    return _records(frame, limit=limit)


@app.get("/expression")
def expression(
    gene: str,
    dataset_id: str,
    limit: int = Query(500, ge=1, le=5000),
) -> dict[str, Any]:
    try:
        frame = _require_corpus().expression(gene, dataset_id)
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    return _records(frame, limit=limit)


@app.get("/rerun")
def rerun_model(gene: str, dataset_id: str) -> dict[str, Any]:
    try:
        result = _require_corpus().rerun_model(gene, dataset_id)
    except ValueError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    result.pop("summary", None)
    return result


@app.get("/cite")
def cite(gene: str, datasets: str = Query(..., description="Comma-separated GSE IDs")) -> dict[str, str]:
    used = [item.strip() for item in datasets.split(",") if item.strip()]
    return _require_corpus().cite_gene(gene, datasets_used=used)


@app.post("/research")
def research(body: ResearchRequest) -> dict[str, Any]:
    try:
        return _require_corpus().research(
            body.gene,
            dataset_id=body.dataset_id,
            species=body.species,
            actions=body.actions,
            fdr_only=body.fdr_only,
        )
    except ValueError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc