"""Scenarios router: CRUD + run + results endpoints.""" from __future__ import annotations import json from collections.abc import AsyncGenerator from datetime import UTC, datetime from typing import Annotated, Any import arq from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import Response from pydantic import BaseModel from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sse_starlette.sse import EventSourceResponse from remodel_api.config import settings from remodel_api.db.models import Scenario from remodel_api.db.session import get_session router = APIRouter() SessionDep = Annotated[AsyncSession, Depends(get_session)] class ScenarioCreate(BaseModel): name: str inputs: dict[str, Any] | None = None class ScenarioRead(BaseModel): id: str name: str status: str kpis_json: str | None created_at: datetime runtime_s: float | None = None model_config = {"from_attributes": True} class ScenarioDetail(ScenarioRead): inputs_json: str | None = None statements_json: str | None = None debt_schedule_json: str | None = None error_message: str | None = None timeseries_path: str | None = None @router.post("/scenarios", response_model=ScenarioRead, status_code=201) async def create_scenario(body: ScenarioCreate, db: SessionDep) -> Scenario: inputs_json = json.dumps(body.inputs or {}) scenario = Scenario(name=body.name, status="queued", inputs_json=inputs_json) db.add(scenario) await db.commit() await db.refresh(scenario) pool = await arq.create_pool(arq.connections.RedisSettings.from_dsn(settings.redis_url)) await pool.enqueue_job("run_scenario_task", scenario.id) await pool.aclose() return scenario @router.get("/scenarios", response_model=list[ScenarioRead]) async def list_scenarios( db: SessionDep, archived: bool = False, ) -> list[Scenario]: q = select(Scenario) if not archived: q = q.where(Scenario.archived_at.is_(None)) result = await db.execute(q.order_by(Scenario.created_at.desc())) return list(result.scalars().all()) @router.get("/scenarios/{scenario_id}", response_model=ScenarioDetail) async def get_scenario(scenario_id: str, db: SessionDep) -> Scenario: scenario = await db.get(Scenario, scenario_id) if scenario is None: raise HTTPException(status_code=404, detail="Scenario not found") return scenario class ScenarioInputsUpdate(BaseModel): inputs: dict[str, Any] @router.patch("/scenarios/{scenario_id}/inputs", response_model=ScenarioRead) async def update_scenario_inputs( scenario_id: str, body: ScenarioInputsUpdate, db: SessionDep ) -> Scenario: """Update inputs and re-queue the scenario for execution.""" scenario = await db.get(Scenario, scenario_id) if scenario is None: raise HTTPException(status_code=404, detail="Scenario not found") scenario.inputs_json = json.dumps(body.inputs) scenario.status = "queued" scenario.kpis_json = None scenario.statements_json = None scenario.debt_schedule_json = None scenario.error_message = None scenario.runtime_s = None await db.commit() await db.refresh(scenario) pool = await arq.create_pool(arq.connections.RedisSettings.from_dsn(settings.redis_url)) await pool.enqueue_job("run_scenario_task", scenario.id) await pool.aclose() return scenario @router.delete("/scenarios/{scenario_id}", status_code=200) async def archive_scenario(scenario_id: str, db: SessionDep) -> dict[str, str]: scenario = await db.get(Scenario, scenario_id) if scenario is None: raise HTTPException(status_code=404, detail="Scenario not found") scenario.archived_at = datetime.now(UTC) await db.commit() return {"status": "archived"} @router.get("/scenarios/{scenario_id}/kpis") async def get_scenario_kpis(scenario_id: str, db: SessionDep) -> dict[str, Any]: scenario = await db.get(Scenario, scenario_id) if scenario is None: raise HTTPException(status_code=404, detail="Scenario not found") if scenario.status != "success": raise HTTPException( status_code=409, detail=f"Scenario is not complete (status={scenario.status})", ) if not scenario.kpis_json: return {} return json.loads(scenario.kpis_json) # type: ignore[no-any-return] @router.get("/scenarios/{scenario_id}/statements") async def get_scenario_statements(scenario_id: str, db: SessionDep) -> dict[str, Any]: scenario = await db.get(Scenario, scenario_id) if scenario is None: raise HTTPException(status_code=404, detail="Scenario not found") if scenario.status != "success": raise HTTPException( status_code=409, detail=f"Scenario is not complete (status={scenario.status})", ) if not scenario.statements_json: return {"pnl": [], "cfs": [], "bs": []} return json.loads(scenario.statements_json) # type: ignore[no-any-return] @router.get("/scenarios/{scenario_id}/export/excel") async def export_scenario_excel(scenario_id: str, db: SessionDep) -> Response: """Export full scenario results to .xlsx.""" scenario = await db.get(Scenario, scenario_id) if scenario is None: raise HTTPException(status_code=404, detail="Scenario not found") if scenario.status != "success": raise HTTPException( status_code=409, detail=f"Scenario is not complete (status={scenario.status})", ) if not scenario.inputs_json: raise HTTPException(status_code=409, detail="No inputs stored for this scenario") import threading from remodel_engine.io.excel_export import export_to_bytes from remodel_engine.schemas.scenario import ScenarioResult if not scenario.kpis_json: raise HTTPException(status_code=409, detail="No KPIs stored for this scenario") result_json = { "inputs": json.loads(scenario.inputs_json or "{}"), "status": scenario.status, "solved_tariff": json.loads(scenario.kpis_json or "{}").get("solved_tariff_inr_per_kwh"), "kpis": json.loads(scenario.kpis_json or "{}"), "financials": json.loads(scenario.statements_json or "{}") or None, "debt_schedule": json.loads(scenario.debt_schedule_json or "[]"), "irr_metrics": {}, } result = ScenarioResult.model_validate(result_json) buf: list[bytes] = [] exc: list[Exception] = [] def _export() -> None: try: buf.append(export_to_bytes(result)) except Exception as e: exc.append(e) t = threading.Thread(target=_export) t.start() t.join() if exc: raise HTTPException(status_code=500, detail=str(exc[0])) filename = f"scenario_{scenario_id[:8]}.xlsx" return Response( content=buf[0], media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", headers={"Content-Disposition": f'attachment; filename="{filename}"'}, ) @router.get("/scenarios/{scenario_id}/events") async def scenario_events(scenario_id: str) -> EventSourceResponse: # pragma: no cover import redis.asyncio as aioredis async def generator() -> AsyncGenerator[dict[str, Any], None]: r = aioredis.from_url(settings.redis_url) # type: ignore[no-untyped-call] channel = f"scenario:{scenario_id}:events" pubsub = r.pubsub() await pubsub.subscribe(channel) try: async for message in pubsub.listen(): if message["type"] == "message": yield {"data": message["data"].decode()} finally: await pubsub.unsubscribe(channel) await r.aclose() return EventSourceResponse(generator())