from __future__ import annotations

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable, Dict, Optional

from app.utils.call_timeline import CallTimeline
from app.ws_gateway.readiness import build_completed_stop_message, build_stop_message


@dataclass(slots=True)
class RoomLifecycleState:
    agent_participant: Optional[Any] = None
    done_payload: Dict[str, Any] = field(default_factory=dict)
    agent_connected: asyncio.Event = field(default_factory=asyncio.Event)
    agent_generation: int = 0
    debug_forward_tasks: set[asyncio.Task[Any]] = field(default_factory=set)

    def __post_init__(self) -> None:
        if self.agent_participant is not None:
            self.agent_connected.set()
            if self.agent_generation <= 0:
                self.agent_generation = 1


def _close_reason(reason: str, *, limit_bytes: int = 120) -> str:
    normalized = " ".join(str(reason or "").strip().split())
    if not normalized:
        return ""

    encoded = normalized.encode("utf-8")
    if len(encoded) <= limit_bytes:
        return normalized
    return encoded[:limit_bytes].decode("utf-8", "ignore")


async def close_websocket(ws: Any, *, code: int = 1000, reason: str = "") -> None:
    close_reason = _close_reason(reason)
    try:
        await ws.close(code=code, reason=close_reason)
    except TypeError:
        await ws.close()


def register_room_lifecycle_handlers(
    *,
    room: Any,
    logger: logging.Logger,
    timeline: CallTimeline,
    room_name: str,
    protocol: str,
    dispatch_started_at: float,
    agent_ready: asyncio.Event,
    call_done: asyncio.Event,
    state: RoomLifecycleState,
    is_agent: Callable[[Any], bool],
    on_agent_disconnect: Callable[[str, int], Awaitable[None]] | None = None,
    ws: Any | None = None,
    debug_events_enabled: bool = False,
    monotonic_fn: Callable[[], float] = time.monotonic,
) -> None:
    @room.on("data_received")
    def _on_data_received(data_packet: Any) -> None:
        try:
            topic = getattr(data_packet, "topic", "")
            if topic == "agent.debug":
                if not debug_events_enabled or ws is None:
                    return
                raw = data_packet.data
                if isinstance(raw, (bytes, bytearray)):
                    raw = raw.decode("utf-8", "ignore")
                message = json.loads(raw)
                message["stress_test"] = True

                async def _forward_debug_event() -> None:
                    try:
                        await ws.send_text(json.dumps(message, ensure_ascii=False))
                    except Exception:
                        logger.debug("debug event forwarding failed", exc_info=True)

                task = asyncio.create_task(
                    _forward_debug_event(), name="forward_agent_debug_event"
                )
                state.debug_forward_tasks.add(task)
                task.add_done_callback(state.debug_forward_tasks.discard)
                return
            if topic != "agent.stage":
                return

            raw = data_packet.data
            if isinstance(raw, (bytes, bytearray)):
                raw = raw.decode("utf-8", "ignore")

            message = json.loads(raw)
            if message.get("stage") == "DONE":
                state.done_payload = message
                timeline.emit(
                    "done_packet_received",
                    reason=message.get("reason", ""),
                    payload=message,
                )
                call_done.set()
        except Exception:
            logger.exception("DATA_PACKET_FAIL | topic=%s", getattr(data_packet, "topic", ""))

    @room.on("participant_connected")
    def _on_participant_connected(participant: Any) -> None:
        try:
            identity = str(getattr(participant, "identity", "") or "")
            if is_agent(participant):
                state.agent_participant = participant
                state.agent_generation += 1
                state.agent_connected.set()
                agent_ready.set()
                logger.info(
                    "LK_AGENT_JOIN | room=%s | agent_identity=%s | dt_s=%.2f",
                    room_name,
                    identity or "-",
                    monotonic_fn() - dispatch_started_at,
                )
                timeline.emit(
                    "agent_join",
                    agent_identity=identity or "",
                    dispatch_dt_ms=round((monotonic_fn() - dispatch_started_at) * 1000),
                )
            else:
                logger.debug("LK_PARTICIPANT_JOIN | room=%s | identity=%s", room_name, identity or "-")
        except Exception:
            logger.debug("participant_connected handler failed", exc_info=True)

    @room.on("participant_disconnected")
    def _on_participant_disconnected(participant: Any) -> None:
        try:
            identity = str(getattr(participant, "identity", "") or "")
            if is_agent(participant):
                current_identity = str(getattr(state.agent_participant, "identity", "") or "")
                cleared_current_agent = False
                if state.agent_participant is participant or current_identity == identity:
                    state.agent_participant = None
                    state.agent_generation += 1
                    state.agent_connected.clear()
                    cleared_current_agent = True
                logger.info("LK_AGENT_LEAVE | room=%s | agent_identity=%s", room_name, identity or "-")
                timeline.emit("agent_leave", agent_identity=identity or "")
                if cleared_current_agent and not call_done.is_set() and on_agent_disconnect is not None:
                    generation = state.agent_generation

                    async def _notify_disconnect() -> None:
                        try:
                            await on_agent_disconnect(identity, generation)
                        except Exception:
                            logger.exception(
                                "AGENT_DISCONNECT_HANDLER_FAIL | room=%s | protocol=%s | agent_identity=%s",
                                room_name,
                                protocol,
                                identity or "-",
                            )

                    asyncio.create_task(_notify_disconnect(), name="agent_disconnect_handler")
            else:
                logger.debug("LK_PARTICIPANT_LEAVE | room=%s | identity=%s", room_name, identity or "-")
        except Exception:
            logger.debug("participant_disconnected handler failed", exc_info=True)


async def watch_call_done(
    *,
    call_done: asyncio.Event,
    state: RoomLifecycleState,
    ws: Any,
    logger: logging.Logger,
    room_name: str,
    protocol: str,
    timeline: CallTimeline,
) -> None:
    await call_done.wait()

    payload = state.done_payload or {}
    reason = str(payload.get("reason") or "DONE")
    logger.info("CALL_DONE | room=%s | protocol=%s | reason=%s", room_name, protocol, reason)
    timeline.emit("call_done", reason=reason)

    if state.debug_forward_tasks:
        await asyncio.gather(*tuple(state.debug_forward_tasks), return_exceptions=True)

    stop_status = str(payload.get("status") or "").strip()
    if stop_status:
        stop_message = build_stop_message(
            status=stop_status,
            reason=reason,
            resource=str(payload.get("resource") or "").strip(),
            failed_resources=list(payload.get("failed_resources") or []),
            phase=str(payload.get("phase") or "in_session").strip(),
        )
    else:
        stop_message = build_completed_stop_message(reason)

    try:
        await ws.send_text(json.dumps(stop_message, ensure_ascii=False))
        timeline.emit("stop_sent", reason=reason)
    except Exception:
        pass

    try:
        await close_websocket(ws, code=1000, reason=f"call_done:{reason}")
    except Exception:
        pass
