"""In-memory pub/sub for pushing events to WebSocket clients."""

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass, field
from pathlib import Path

from mixersystem.studio.formatting import fmt_duration, fmt_cost
from mixersystem.studio.log_parser import parse_trace_line
from mixersystem.studio.trace_db import (
    has_trace_db,
    query_events_after,
    query_max_id,
    TraceEventRow,
)

_MAX_QUEUE_SIZE = 256


@dataclass
class StudioEvent:
    type: str
    payload: dict = field(default_factory=dict)
    run_id: str = ""
    session_folder: str = ""

    def to_json(self) -> str:
        return json.dumps({
            "type": self.type,
            "payload": self.payload,
            "run_id": self.run_id,
            "session_folder": self.session_folder,
        })


class EventBus:
    """Broadcast events to all subscribed WebSocket clients."""

    def __init__(self, project_root: Path | None = None) -> None:
        self._subscribers: set[asyncio.Queue[StudioEvent]] = set()
        self._watchers: dict[str, asyncio.Task] = {}
        self._watcher_offsets: dict[str, int] = {}       # text log byte offsets
        self._watcher_db_cursors: dict[str, int] = {}    # SQLite last-seen event id
        self._watcher_uses_db: dict[str, bool] = {}      # True if watcher uses SQLite
        self._project_root = project_root or Path.cwd()

    def subscribe(self) -> asyncio.Queue[StudioEvent]:
        queue: asyncio.Queue[StudioEvent] = asyncio.Queue(maxsize=_MAX_QUEUE_SIZE)
        self._subscribers.add(queue)
        return queue

    def unsubscribe(self, queue: asyncio.Queue[StudioEvent]) -> None:
        self._subscribers.discard(queue)

    def emit(self, event: StudioEvent) -> None:
        for queue in self._subscribers:
            if queue.full():
                try:
                    queue.get_nowait()  # drop oldest
                except asyncio.QueueEmpty:
                    pass
            try:
                queue.put_nowait(event)
            except asyncio.QueueFull:
                pass

    # ------------------------------------------------------------------
    # Public watcher API
    # ------------------------------------------------------------------

    def watch_trace_log(self, session_folder: str, run_id: str) -> None:
        """Spawn an async task that polls trace.db (preferred) or agent_trace.log."""
        key = f"{session_folder}:{run_id}"
        if key in self._watchers:
            return

        folder_path = self._project_root / session_folder
        if has_trace_db(folder_path):
            self._watcher_uses_db[key] = True
            task = asyncio.create_task(self._poll_trace_db(session_folder, run_id, key))
        else:
            self._watcher_uses_db[key] = False
            task = asyncio.create_task(self._poll_trace(session_folder, run_id, key))
        self._watchers[key] = task

    def stop_watching(self, session_folder: str, run_id: str) -> None:
        key = f"{session_folder}:{run_id}"
        uses_db = self._watcher_uses_db.get(key, False)
        if uses_db:
            self._flush_trace_db(session_folder, run_id, key)
        else:
            self._flush_trace(session_folder, run_id, key)
        task = self._watchers.pop(key, None)
        if task and not task.done():
            task.cancel()
        self._watcher_offsets.pop(key, None)
        self._watcher_db_cursors.pop(key, None)
        self._watcher_uses_db.pop(key, None)

    # ------------------------------------------------------------------
    # SQLite polling
    # ------------------------------------------------------------------

    def _emit_trace_event(self, row: TraceEventRow, run_id: str, session_folder: str) -> None:
        """Convert a TraceEventRow into a StudioEvent and emit it."""
        self.emit(StudioEvent(
            type="agent_event",
            payload={
                "timestamp_utc": row.timestamp_utc,
                "event_type": row.event_type,
                "agent_name": row.agent_name,
                "call_id": row.call_id,
                "model": row.model,
                "mode": row.mode,
                "status": row.status,
                "error_message": row.error_message,
                # Formatted strings for backward compat with LiveLog
                "duration": fmt_duration(row.duration_secs),
                "cost": fmt_cost(row.cost_usd),
                "duration_seconds": row.duration_secs,
                "cost_usd": row.cost_usd,
                # New native fields
                "duration_secs": row.duration_secs,
                "event_id": row.id,
                "provider": row.provider,
                "depth": row.depth,
                "is_resume": row.is_resume,
            },
            run_id=run_id,
            session_folder=session_folder,
        ))

    async def _poll_trace_db(self, session_folder: str, run_id: str, key: str) -> None:
        folder_path = self._project_root / session_folder

        # Wait for DB file to appear (agent may not have written yet)
        while not has_trace_db(folder_path):
            await asyncio.sleep(0.5)

        # Start from current max id
        cursor = query_max_id(folder_path)
        self._watcher_db_cursors[key] = cursor

        try:
            while True:
                await asyncio.sleep(0.5)
                if not has_trace_db(folder_path):
                    continue
                new_events = query_events_after(folder_path, cursor)
                if not new_events:
                    continue
                for row in new_events:
                    self._emit_trace_event(row, run_id, session_folder)
                    cursor = row.id
                self._watcher_db_cursors[key] = cursor
        except asyncio.CancelledError:
            pass
        finally:
            self._watchers.pop(key, None)
            self._watcher_db_cursors.pop(key, None)
            self._watcher_uses_db.pop(key, None)

    def _flush_trace_db(self, session_folder: str, run_id: str, key: str) -> None:
        """Synchronous final read from SQLite to catch events written right before stop."""
        folder_path = self._project_root / session_folder
        if not has_trace_db(folder_path):
            return
        cursor = self._watcher_db_cursors.get(key, 0)
        new_events = query_events_after(folder_path, cursor)
        for row in new_events:
            self._emit_trace_event(row, run_id, session_folder)

    # ------------------------------------------------------------------
    # Text log polling (legacy)
    # ------------------------------------------------------------------

    def _emit_text_events(self, lines: str, run_id: str, session_folder: str) -> None:
        """Parse text trace lines and emit events."""
        for line in lines.splitlines():
            event = parse_trace_line(line)
            if event is None:
                continue
            self.emit(StudioEvent(
                type="agent_event",
                payload={
                    "timestamp_utc": event.timestamp_utc,
                    "event_type": event.event_type,
                    "agent_name": event.agent_name,
                    "call_id": event.call_id,
                    "model": event.model,
                    "mode": event.mode,
                    "duration": event.duration,
                    "cost": event.cost,
                    "status": event.status,
                    "error_message": event.error_message,
                },
                run_id=run_id,
                session_folder=session_folder,
            ))

    def _flush_trace(self, session_folder: str, run_id: str, key: str) -> None:
        """Synchronously read any remaining trace lines and emit events."""
        trace_path = self._project_root / session_folder / "logs" / "agent_trace.log"
        offset = self._watcher_offsets.get(key, 0)
        if not trace_path.is_file():
            return
        size = trace_path.stat().st_size
        if size <= offset:
            return
        with open(trace_path, "r") as f:
            f.seek(offset)
            new_data = f.read()
            self._watcher_offsets[key] = f.tell()
        self._emit_text_events(new_data, run_id, session_folder)

    async def _poll_trace(self, session_folder: str, run_id: str, key: str) -> None:
        trace_path = self._project_root / session_folder / "logs" / "agent_trace.log"
        offset = 0

        # Start from current end if file already exists
        if trace_path.is_file():
            offset = trace_path.stat().st_size
        self._watcher_offsets[key] = offset

        try:
            while True:
                await asyncio.sleep(0.5)

                # Check if SQLite DB appeared (new session started writing)
                # If so, switch to DB polling
                folder_path = self._project_root / session_folder
                if has_trace_db(folder_path):
                    self._watcher_uses_db[key] = True
                    self._watcher_offsets.pop(key, None)
                    await self._poll_trace_db(session_folder, run_id, key)
                    return

                if not trace_path.is_file():
                    continue

                size = trace_path.stat().st_size
                if size <= offset:
                    continue

                with open(trace_path, "r") as f:
                    f.seek(offset)
                    new_data = f.read()
                    offset = f.tell()
                self._watcher_offsets[key] = offset
                self._emit_text_events(new_data, run_id, session_folder)
        except asyncio.CancelledError:
            pass
        finally:
            self._watchers.pop(key, None)
            self._watcher_offsets.pop(key, None)
