cool/api/main.py

"""FastAPI application: REST + WebSocket API and HTMX-rendered frontend."""
import asyncio
from io import StringIO
from pathlib import Path

from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.responses import HTMLResponse, PlainTextResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio import SeqIO

from core import SPECIES_PROFILES, DnaAnalyzer, species_choices, mechanism_for
from core.codons import CODON_USAGE_TABLES
from .schemas import DesignRequest, AnalyzeRequest, GrnaRequest, CAS_TO_PAM
from .service import run_design, analyze_sequence

BASE_DIR = Path(__file__).resolve().parent.parent
templates = Jinja2Templates(directory=str(BASE_DIR / "web" / "templates"))

app = FastAPI(title="Plant DNA Designer", version="2.0")
app.mount("/static", StaticFiles(directory=str(BASE_DIR / "web" / "static")), name="static")


# --- pages -----------------------------------------------------------------
@app.get("/", response_class=HTMLResponse)
def index(request: Request):
    return templates.TemplateResponse(request, "index.html", {
        "species": species_choices(),
        "profiles": SPECIES_PROFILES,
        "cas_types": list(CAS_TO_PAM.keys()),
        "codon_tables": list(CODON_USAGE_TABLES.keys()),
    })


@app.get("/traits/{species}", response_class=HTMLResponse)
def traits_partial(request: Request, species: str):
    """HTMX swap: trait checkboxes for the selected species."""
    prof = SPECIES_PROFILES.get(species)
    if not prof:
        raise HTTPException(404, "Unknown species")
    return templates.TemplateResponse(request, "partials/traits.html", {
        "traits": list(prof['traits'].keys()),
        "gc_range": prof['gc_range'], "codon_table": prof['codon_table'],
    })


@app.get("/api/mechanism/{trait}")
def api_mechanism(trait: str):
    """Trait → Mechanism → Effector protein record (for the UI target layer)."""
    m = mechanism_for(trait)
    if not m:
        raise HTTPException(404, f"No known mechanism for trait: {trait}")
    return JSONResponse(m)


# --- JSON API --------------------------------------------------------------
@app.post("/api/design")
def api_design(req: DesignRequest):
    try:
        return run_design(req)
    except ValueError as e:
        raise HTTPException(400, str(e))


@app.post("/api/analyze")
def api_analyze(req: AnalyzeRequest):
    return analyze_sequence(req.dna.strip().upper(), req.codon_table)


@app.post("/api/grna")
def api_grna(req: GrnaRequest):
    from core import design_grna
    pam = CAS_TO_PAM.get(req.cas_type, "NGG")
    grnas = design_grna(req.dna.strip().upper(), pam_type=pam, grna_num=req.grna_num)
    return [{"grna": g, "position": p, "on_target": round(s["on_target"], 3),
             "off_target": round(s["off_target"], 3)} for g, p, s in grnas]


# --- exports ---------------------------------------------------------------
@app.get("/api/export/{fmt}")
def api_export(fmt: str, dna: str, name: str = "synthetic"):
    dna = dna.strip().upper()
    ana = DnaAnalyzer(dna)
    if fmt == "fasta":
        return PlainTextResponse(f">{name}_{ana.sha()}\n{dna}\n",
                                 headers={"Content-Disposition": f"attachment; filename={ana.sha()}.fasta"})
    if fmt == "sbol":
        return PlainTextResponse(ana.generate_sbol(),
                                 headers={"Content-Disposition": f"attachment; filename={ana.sha()}.sbol"})
    if fmt == "genbank":
        record = SeqRecord(Seq(dna), id=ana.sha(), description=f"Synthetic DNA ({name})")
        record.annotations["molecule_type"] = "DNA"
        buf = StringIO()
        SeqIO.write(record, buf, "genbank")
        return PlainTextResponse(buf.getvalue(),
                                 headers={"Content-Disposition": f"attachment; filename={ana.sha()}.gb"})
    raise HTTPException(400, f"Unknown format: {fmt}")


# --- WebSocket: live design with streaming progress ------------------------
@app.websocket("/ws/design")
async def ws_design(ws: WebSocket):
    await ws.accept()
    loop = asyncio.get_running_loop()
    try:
        payload = await ws.receive_json()
        req = DesignRequest(**payload)
        queue: asyncio.Queue = asyncio.Queue()

        def progress_cb(fraction: float):
            loop.call_soon_threadsafe(queue.put_nowait, fraction)

        async def pump():
            while True:
                frac = await queue.get()
                await ws.send_json({"type": "progress", "value": round(frac, 4)})

        pump_task = asyncio.create_task(pump())
        try:
            result = await loop.run_in_executor(None, lambda: run_design(req, progress_cb))
            await asyncio.sleep(0)  # let the final progress frame flush
            await ws.send_json({"type": "result", "data": result})
        finally:
            pump_task.cancel()
    except WebSocketDisconnect:
        return
    except Exception as e:  # surface validation/runtime errors to the client
        await ws.send_json({"type": "error", "message": str(e)})
    finally:
        await ws.close()


@app.get("/api/health")
def health():
    return JSONResponse({"status": "ok", "species": len(SPECIES_PROFILES)})