IPSExtractor/server/app/main.py

"""Servidor de tracking de IPSExtractor.

Persiste el historial de lotes en SQLite y distribuye actualizaciones
en tiempo real vía WebSockets al dashboard web.
"""
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.encoders import jsonable_encoder
from fastapi.responses import FileResponse
from contextlib import asynccontextmanager
from typing import Dict, List, Any, Optional
from pydantic import BaseModel
import uuid
import os
from datetime import datetime

from .schemas.batch import BatchCreate, BatchUpdate, BatchRead, BatchStatus
from .database import init_db, insert_batch, update_batch, get_all_batches, get_batch, insert_log, get_logs


@asynccontextmanager
async def lifespan(app: FastAPI):
    init_db()
    yield

app = FastAPI(title="IPSExtractor Tracking API", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Cache en memoria para lotes activos (para broadcast rápido)
active_batches: Dict[str, Dict[str, Any]] = {}
active_connections: List[WebSocket] = []


async def broadcast(message: dict):
    json_compatible_message = jsonable_encoder(message)
    disconnected = []
    for connection in active_connections:
        try:
            await connection.send_json(json_compatible_message)
        except Exception:
            disconnected.append(connection)
    for conn in disconnected:
        active_connections.remove(conn)


@app.post("/batches", response_model=BatchRead)
async def create_batch(batch: BatchCreate):
    batch_id = str(uuid.uuid4())
    now = datetime.now()
    new_batch = {
        "id": batch_id,
        **batch.model_dump(),
        "status": BatchStatus.PROCESSING,
        "current_radicado_index": 0,
        "success_count": 0,
        "error_count": 0,
        "total_tokens": 0,
        "total_cost_usd": 0.0,
        "last_error": None,
        "created_at": now,
        "updated_at": now,
    }
    active_batches[batch_id] = new_batch
    insert_batch(new_batch)
    await broadcast({"event": "batch_created", "data": new_batch})
    return new_batch


@app.get("/batches", response_model=List[BatchRead])
async def list_batches():
    # Retorna historial completo desde SQLite
    db_batches = get_all_batches()
    # Sobreescribe con datos en memoria si el lote está activo (más fresco)
    result = []
    for b in db_batches:
        if b["id"] in active_batches:
            result.append(active_batches[b["id"]])
        else:
            result.append(b)
    return result


@app.get("/batches/{batch_id}", response_model=BatchRead)
async def get_batch_endpoint(batch_id: str):
    if batch_id in active_batches:
        return active_batches[batch_id]
    batch = get_batch(batch_id)
    if not batch:
        raise HTTPException(status_code=404, detail="Lote no encontrado")
    return batch


@app.patch("/batches/{batch_id}", response_model=BatchRead)
async def update_batch_endpoint(batch_id: str, update: BatchUpdate):
    if batch_id in active_batches:
        batch = active_batches[batch_id]
        update_data = update.model_dump(exclude_unset=True)
        for key, value in update_data.items():
            batch[key] = value
        batch["updated_at"] = datetime.now()

        # Si el lote se completó, sacarlo del cache activo
        if batch.get("status") in (BatchStatus.COMPLETED, BatchStatus.FAILED):
            active_batches.pop(batch_id, None)

        update_batch(batch_id, update_data)
        await broadcast({"event": "batch_updated", "data": batch})
        return batch

    # Lote no está en memoria (completado), actualizar solo en DB
    batch = get_batch(batch_id)
    if not batch:
        raise HTTPException(status_code=404, detail="Lote no encontrado")
    update_data = update.model_dump(exclude_unset=True)
    update_batch(batch_id, update_data)
    batch.update(update_data)
    return batch


class LogEntry(BaseModel):
    message: str
    level: str = "info"
    radicado: Optional[str] = None

class LogRead(BaseModel):
    id: int
    batch_id: str
    timestamp: str
    level: str
    radicado: Optional[str]
    message: str

@app.post("/batches/{batch_id}/logs", status_code=201)
async def add_log(batch_id: str, entry: LogEntry):
    """Añade una entrada de log para un lote."""
    insert_log(batch_id, entry.message, entry.level, entry.radicado)
    # Broadcast log también por WebSocket para actualizaciones en vivo
    await broadcast({"event": "log_added", "batch_id": batch_id, "data": {
        "level": entry.level, "radicado": entry.radicado, "message": entry.message
    }})
    return {"ok": True}


@app.get("/batches/{batch_id}/logs", response_model=List[LogRead])
async def list_logs(batch_id: str, limit: int = 200):
    """Retorna el historial de logs de un lote."""
    return get_logs(batch_id, limit)


@app.get("/download/{batch_id}")
async def download_csv(batch_id: str):
    """Sirve el archivo CSV generado por el lote de extracción."""
    batch = get_batch(batch_id)
    if not batch:
        raise HTTPException(status_code=404, detail="Lote no encontrado")

    csv_path = batch.get("csv_output_path")
    if not csv_path:
        raise HTTPException(status_code=404, detail="Este lote no tiene un archivo CSV asociado")
    if not os.path.isfile(csv_path):
        raise HTTPException(status_code=404, detail=f"Archivo CSV no encontrado en el servidor: {csv_path}")

    filename = os.path.basename(csv_path)
    return FileResponse(
        path=csv_path,
        media_type="text/csv",
        filename=filename,
    )


@app.websocket("/ws/tracking")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    active_connections.append(websocket)
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        if websocket in active_connections:
            active_connections.remove(websocket)