from __future__ import annotations
import hashlib, os, time, uuid
from pathlib import Path
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
import torch
from model import OncoCT1B, OncoCTConfig

app = FastAPI(title="OncoCT-1B Research Inference", version="0.1.0")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL = OncoCT1B(OncoCTConfig()).to(DEVICE).eval()
CHECKPOINT = os.getenv("ONCOCT_CHECKPOINT")
if CHECKPOINT and Path(CHECKPOINT).exists():
    state = torch.load(CHECKPOINT, map_location=DEVICE, weights_only=False)
    MODEL.load_state_dict(state.get("model", state), strict=False)

class InferenceRequest(BaseModel):
    volume_tokens_path: str = Field(..., description="Path to a preprocessed research tensor; do not expose raw filesystem paths in production")
    study_time: float = 0.0
    request_id: str | None = None

@app.get("/health")
def health(): return {"status": "ok", "device": str(DEVICE), "model_loaded": CHECKPOINT is not None}

@app.post("/infer")
def infer(req: InferenceRequest):
    rid = req.request_id or str(uuid.uuid4())
    started = time.time()
    p = Path(req.volume_tokens_path)
    if not p.exists(): raise HTTPException(404, "tensor input not found")
    item = torch.load(p, map_location="cpu", weights_only=True)
    tokens = item["volume_tokens"].float().unsqueeze(0).to(DEVICE)
    t = torch.tensor([[req.study_time]], device=DEVICE)
    with torch.inference_mode(): out = MODEL(tokens, t)
    abstain = float(torch.sigmoid(out["abstain_logit"])[0].cpu())
    response_prob = torch.softmax(out["response_logits"], -1)[0].cpu().tolist()
    audit_hash = hashlib.sha256(f"{rid}|{p}|{time.time_ns()}".encode()).hexdigest()
    return {"request_id": rid, "status": "abstain" if abstain >= 0.5 else "review_required", "abstention_score": abstain, "response_probabilities": response_prob, "evidence": {"volume_log1p": float(out["volume_log1p"][0].cpu()), "long_axis_mm": float(out["long_axis_mm"][0].cpu())}, "latency_ms": (time.time()-started)*1000, "audit_hash": audit_hash, "clinical_notice": "Research output only. Requires qualified radiologist/oncologist review; does not diagnose or recommend treatment."}