from dotenv import load_dotenv
load_dotenv()  # no-op if .env doesn't exist (e.g. on HF Spaces, where env vars are set directly)

import json

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from app.graph.build import build_graph

app = FastAPI(title="VeriScite", version="0.1.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # tighten to your Vercel domain before deploying
    allow_methods=["*"],
    allow_headers=["*"],
)

graph = build_graph()


class AuditRequest(BaseModel):
    text: str  # pasted paper/abstract text


@app.get("/health")
async def health():
    return {"status": "ok"}


@app.post("/audit")
async def audit(req: AuditRequest):
    initial_state = {
        "source_text": req.text,
        "claims": [],
        "audits": [],
        "current_index": 0,
        "report": None,
    }
    final_state = await graph.ainvoke(initial_state)
    return final_state["report"]


@app.post("/audit/stream")
async def audit_stream(req: AuditRequest):
    """Server-Sent Events version of /audit — emits a live event after each
    graph node (and after each individual tool call inside the escalate
    agent loop) completes, rather than only returning at the very end.
    Powers the frontend's live agent-workflow visualization.

    Event shapes:
      {"type": "progress", ...node-specific fields, see nodes.py _writer() calls}
      {"type": "final", "report": {...}}   -- sent once, at the end
    """
    initial_state = {
        "source_text": req.text,
        "claims": [],
        "audits": [],
        "current_index": 0,
        "report": None,
    }

    async def event_generator():
        async for mode, chunk in graph.astream(initial_state, stream_mode=["custom", "values"]):
            if mode == "custom":
                yield f"data: {json.dumps({'type': 'progress', **chunk})}\n\n"
            elif mode == "values" and chunk.get("report") is not None:
                yield f"data: {json.dumps({'type': 'final', 'report': chunk['report']})}\n\n"

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