"""Per-run event bus + factories + SSE formatting.

Tiny in-memory queue keyed by run_id. The runtime publishes on each
transition; the FastAPI SSE endpoint subscribes and forwards.

Event payloads are constructed via the module-level factory functions below
so the dict shape lives in one place and is grep-able.

Many event types carry an `analysis` field that names the active Analysis
(read from `current_analysis` ContextVar). The UI uses it to group
sub-events under the right Analysis node without inferring from time order.
"""

from __future__ import annotations

import asyncio
import json
from typing import Any

from api.analyses.types import Finding
from api.runtime.context import current_analysis, current_run_id

# run_id -> queue of events. Events are plain dicts.
_queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {}


# ── Queue management ──

def open_run(run_id: str) -> asyncio.Queue:
    q: asyncio.Queue = asyncio.Queue()
    _queues[run_id] = q
    return q


async def publish(run_id: str, event: dict[str, Any]) -> None:
    q = _queues.get(run_id)
    if q is not None:
        await q.put(event)


async def publish_current(event: dict[str, Any]) -> None:
    """Publish to the run identified by the `current_run_id` contextvar.
    No-op when called outside a run."""
    rid = current_run_id.get()
    if rid is not None:
        await publish(rid, event)


async def close(run_id: str) -> None:
    q = _queues.get(run_id)
    if q is not None:
        await q.put(None)


def drop(run_id: str) -> None:
    _queues.pop(run_id, None)


# ── SSE formatting + streaming ──

def sse_format(event: dict[str, Any]) -> str:
    kind = event.get("type", "message")
    payload = json.dumps({k: v for k, v in event.items() if k != "type"}, default=str)
    return f"event: {kind}\ndata: {payload}\n\n"


async def stream(run_id: str):
    """Async generator producing SSE-formatted strings until the run ends."""
    q = _queues.get(run_id)
    if q is None:
        yield sse_format({"type": "error", "message": f"unknown run_id {run_id}"})
        return
    try:
        while True:
            event = await q.get()
            if event is None:
                yield sse_format({"type": "end"})
                return
            yield sse_format(event)
    finally:
        drop(run_id)


# ── Event factories ──
# Each returns the dict payload to pass into `publish()` /
# `publish_current()`. Names mirror the `type` field for grep-ability.
# Factories that fire from inside an Analysis pick up the analysis name from
# the `current_analysis` ContextVar so callers don't have to pass it.

def run_start(question: str) -> dict[str, Any]:
    return {"type": "run_start", "question": question}


def run_end(*, answer: str | None = None, error: str | None = None) -> dict[str, Any]:
    return {"type": "run_end", "answer": answer or "", "error": error}


def plan_ready(rationale: str, steps: list[dict[str, Any]]) -> dict[str, Any]:
    return {"type": "plan_ready", "rationale": rationale, "steps": steps}


def node_update(node: str, update: Any) -> dict[str, Any]:
    return {
        "type": "node_update",
        "node": node,
        "keys": sorted(update.keys()) if isinstance(update, dict) else [],
    }


def analysis_start(name: str, args: dict[str, Any], why: str, *,
                   step_id: str | None = None) -> dict[str, Any]:
    return {"type": "analysis_start", "analysis": name, "step_id": step_id,
            "args": args, "why": why}


def analysis_end(name: str, finding: Finding, *,
                 step_id: str | None = None) -> dict[str, Any]:
    return {
        "type": "analysis_end",
        "analysis": name,
        "step_id": step_id,
        "summary": finding.summary,
        "error": finding.error,
        "row_count": len(finding.rows),
    }


def analysis_fallback(*, primary: str, fallback: str, reason: str,
                      step_id: str | None = None) -> dict[str, Any]:
    return {
        "type": "analysis_fallback",
        "primary": primary,
        "fallback": fallback,
        "reason": reason,
        "step_id": step_id,
    }


def synth_start() -> dict[str, Any]:
    return {"type": "synth_start"}


def tool_call_start(tool: str, *, input: Any = None) -> dict[str, Any]:
    return {
        "type": "tool_call_start",
        "tool": tool,
        "analysis": current_analysis.get(),
        "input": input,
    }


def tool_call_end(tool: str, *, output: Any = None,
                  error: str | None = None) -> dict[str, Any]:
    return {
        "type": "tool_call_end",
        "tool": tool,
        "analysis": current_analysis.get(),
        "output": output,
        "error": error,
    }


def llm_call(label: str, *, system_len: int, user_len: int) -> dict[str, Any]:
    """Lightweight breadcrumb for an LLM call. Doesn't carry the prompt
    contents (those go through Langfuse) — just the fact and rough size."""
    return {
        "type": "llm_call",
        "label": label,
        "analysis": current_analysis.get(),
        "system_chars": system_len,
        "user_chars": user_len,
    }