api/routes/analysis.py

import uuid
import logging
import json
import asyncio
from datetime import datetime, timezone
from typing import Optional, List

from fastapi import APIRouter, BackgroundTasks, WebSocket, WebSocketDisconnect, HTTPException

from api.schemas import AnalyzeRequest, AnalyzeResponse, StatusResponse
from src.pipeline import run_full_pipeline, PipelineState
from src.storage import save_pipeline_run, get_pipeline_run

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api")

# WebSocket Connection Manager for real-time progress broadcasts
class ConnectionManager:
    def __init__(self):
        self.active_connections: dict[str, List[WebSocket]] = {}
        self._locks: dict[str, asyncio.Lock] = {}

    async def connect(self, run_id: str, websocket: WebSocket):
        await websocket.accept()
        if run_id not in self.active_connections:
            self.active_connections[run_id] = []
            self._locks[run_id] = asyncio.Lock()
        async with self._locks[run_id]:
            self.active_connections[run_id].append(websocket)

    def disconnect(self, run_id: str, websocket: WebSocket):
        if run_id in self.active_connections:
            try:
                self.active_connections[run_id].remove(websocket)
                if not self.active_connections[run_id]:
                    del self.active_connections[run_id]
                    if run_id in self._locks:
                        del self._locks[run_id]
            except ValueError:
                pass  # WebSocket already removed

    async def broadcast_status(self, run_id: str, data: dict):
        if run_id in self.active_connections and run_id in self._locks:
            async with self._locks[run_id]:
                for connection in list(self.active_connections[run_id]):
                    try:
                        await connection.send_json(data)
                    except Exception as e:
                        logger.warning(f"Failed to send websocket broadcast to connection: {e}")
                        try:
                            self.active_connections[run_id].remove(connection)
                        except ValueError:
                            pass

manager = ConnectionManager()

async def update_run_status(
    run_id: str,
    query: str,
    status: str,
    papers_fetched: Optional[int] = None,
    claims_extracted: Optional[int] = None,
    contradictions_found: Optional[int] = None,
    error_message: Optional[str] = None,
    pmids: Optional[list[str]] = None,
    report_json: Optional[str] = None,
    started_at: Optional[str] = None,
    completed_at: Optional[str] = None,
    status_message: Optional[str] = None,
    total_papers: Optional[int] = None,
    papers_extracted: Optional[int] = None,
    nli_pairs_total: Optional[int] = None,
    nli_pairs_scored: Optional[int] = None,
    judge_pairs_total: Optional[int] = None,
    judge_pairs_scored: Optional[int] = None,
    seed_claim: Optional[str] = None,
    date_from: Optional[int] = None,
    date_to: Optional[int] = None,
    journals: Optional[list[str]] = None
):
    """Helper to save the current run status in SQLite and broadcast it to WebSockets."""
    save_pipeline_run(
        run_id=run_id,
        query=query,
        status=status,
        papers_fetched=papers_fetched,
        claims_extracted=claims_extracted,
        contradictions_found=contradictions_found,
        error_message=error_message,
        pmids=pmids,
        report_json=report_json,
        started_at=started_at,
        completed_at=completed_at,
        status_message=status_message,
        total_papers=total_papers,
        papers_extracted=papers_extracted,
        nli_pairs_total=nli_pairs_total,
        nli_pairs_scored=nli_pairs_scored,
        judge_pairs_total=judge_pairs_total,
        judge_pairs_scored=judge_pairs_scored,
        seed_claim=seed_claim,
        date_from=date_from,
        date_to=date_to,
        journals=journals
    )
    
    await manager.broadcast_status(run_id, {
        "run_id": run_id,
        "query": query,
        "status": status,
        "papers_fetched": papers_fetched,
        "claims_extracted": claims_extracted,
        "contradictions_found": contradictions_found,
        "error_message": error_message,
        "completed_at": completed_at,
        "status_message": status_message,
        "total_papers": total_papers,
        "papers_extracted": papers_extracted,
        "nli_pairs_total": nli_pairs_total,
        "nli_pairs_scored": nli_pairs_scored,
        "judge_pairs_total": judge_pairs_total,
        "judge_pairs_scored": judge_pairs_scored,
        "seed_claim": seed_claim,
        "date_from": date_from,
        "date_to": date_to,
        "journals": journals,
        "timestamp": datetime.now(timezone.utc).isoformat()
    })

async def run_analysis_background(
    run_id: str,
    query: str,
    max_papers: int,
    seed_claim: Optional[str] = None,
    date_from: Optional[int] = None,
    date_to: Optional[int] = None,
    journals: Optional[list[str]] = None
):
    """The async background task executing all RSCE stages using the unified run_full_pipeline."""
    async def on_stage_complete(state: PipelineState):
        pmids = [p.pmid for p in state.papers]
        completed_at = datetime.now(timezone.utc).isoformat() if state.status in ("COMPLETED", "FAILED") else None
        await update_run_status(
            run_id=state.run_id,
            query=state.query,
            status=state.status,
            papers_fetched=state.papers_fetched if state.papers_fetched is not None else len(state.papers),
            claims_extracted=state.claims_extracted if state.claims_extracted is not None else len(state.claims),
            contradictions_found=state.contradictions_found if state.contradictions_found is not None else len(state.contradictions),
            error_message=state.error_message,
            pmids=pmids,
            report_json=state.report.model_dump_json() if state.report else None,
            started_at=state.started_at,
            completed_at=completed_at,
            status_message=state.status_message,
            total_papers=state.total_papers,
            papers_extracted=state.papers_extracted,
            nli_pairs_total=state.nli_pairs_total,
            nli_pairs_scored=state.nli_pairs_scored,
            judge_pairs_total=state.judge_pairs_total,
            judge_pairs_scored=state.judge_pairs_scored,
            seed_claim=state.seed_claim,
            date_from=state.date_from,
            date_to=state.date_to,
            journals=state.journals
        )

    try:
        await run_full_pipeline(
            query=query,
            max_papers=max_papers,
            run_id=run_id,
            seed_claim=seed_claim,
            date_from=date_from,
            date_to=date_to,
            journals=journals,
            on_stage_complete=on_stage_complete
        )
    except Exception as e:
        logger.error(f"Background analysis task failed for run {run_id}: {e}")

@router.post("/analyze", response_model=AnalyzeResponse)
async def analyze(request: AnalyzeRequest, background_tasks: BackgroundTasks):
    """Start the research synthesis pipeline for a query in the background."""
    run_id = str(uuid.uuid4())
    # Create the run record immediately so the results page can observe it
    # before the background worker starts updating status.
    save_pipeline_run(
        run_id=run_id,
        query=request.query,
        status="RUNNING",
        started_at=datetime.now(timezone.utc).isoformat(),
        status_message="Queued for processing...",
        seed_claim=request.seed_claim,
        date_from=request.date_from,
        date_to=request.date_to,
        journals=request.journals
    )
    background_tasks.add_task(
        run_analysis_background,
        run_id=run_id,
        query=request.query,
        max_papers=request.max_papers,
        seed_claim=request.seed_claim,
        date_from=request.date_from,
        date_to=request.date_to,
        journals=request.journals
    )
    return AnalyzeResponse(
        run_id=run_id,
        query=request.query,
        status="RUNNING"
    )

@router.get("/status/{run_id}", response_model=StatusResponse)
async def get_status(run_id: str):
    """Retrieve the current processing status of an analysis run."""
    run = get_pipeline_run(run_id)
    if not run:
        raise HTTPException(status_code=404, detail="Analysis run not found.")
        
    journals_list = None
    if run.get("journals"):
        try:
            journals_list = json.loads(run["journals"])
        except Exception:
            pass

    return StatusResponse(
        run_id=run["id"],
        query=run["query"],
        status=run["status"],
        papers_fetched=run["papers_fetched"],
        claims_extracted=run["claims_extracted"],
        contradictions_found=run["contradictions_found"],
        started_at=run["started_at"],
        completed_at=run["completed_at"],
        error_message=run["error_message"],
        status_message=run.get("status_message"),
        total_papers=run.get("total_papers"),
        papers_extracted=run.get("papers_extracted"),
        nli_pairs_total=run.get("nli_pairs_total"),
        nli_pairs_scored=run.get("nli_pairs_scored"),
        judge_pairs_total=run.get("judge_pairs_total"),
        judge_pairs_scored=run.get("judge_pairs_scored"),
        seed_claim=run.get("seed_claim"),
        date_from=run.get("date_from"),
        date_to=run.get("date_to"),
        journals=journals_list
    )

@router.websocket("/ws/{run_id}")
async def websocket_endpoint(websocket: WebSocket, run_id: str):
    """WebSocket endpoint to subscribe to real-time analysis pipeline stage updates."""
    await manager.connect(run_id, websocket)
    # Immediately send current state if it exists
    run = get_pipeline_run(run_id)
    if run:
        journals_list = None
        if run.get("journals"):
            try:
                journals_list = json.loads(run["journals"])
            except Exception:
                pass
        try:
            await websocket.send_json({
                "run_id": run["id"],
                "query": run["query"],
                "status": run["status"],
                "papers_fetched": run["papers_fetched"],
                "claims_extracted": run["claims_extracted"],
                "contradictions_found": run["contradictions_found"],
                "error_message": run["error_message"],
                "completed_at": run["completed_at"],
                "status_message": run.get("status_message"),
                "total_papers": run.get("total_papers"),
                "papers_extracted": run.get("papers_extracted"),
                "nli_pairs_total": run.get("nli_pairs_total"),
                "nli_pairs_scored": run.get("nli_pairs_scored"),
                "judge_pairs_total": run.get("judge_pairs_total"),
                "judge_pairs_scored": run.get("judge_pairs_scored"),
                "seed_claim": run.get("seed_claim"),
                "date_from": run.get("date_from"),
                "date_to": run.get("date_to"),
                "journals": journals_list,
                "timestamp": datetime.now(timezone.utc).isoformat()
            })
        except Exception:
            pass
            
    try:
        while True:
            # We must await receive_text to keep the connection alive and detect disconnects
            await websocket.receive_text()
    except WebSocketDisconnect:
        manager.disconnect(run_id, websocket)