Remodel/packages/api/src/remodel_api/routers/scenarios.py

"""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())