from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING, Any

from relay.server.models import Group, Message, Peer

if TYPE_CHECKING:
    from relay.server.db import MessageRow, RelayDB


class RelayState:
    def __init__(self, db: RelayDB | None = None) -> None:
        self._db = db
        self.groups: dict[str, Group] = {}
        self.sessions: dict[str, tuple[str, list[str]]] = {}
        self.sse_queues: dict[str, asyncio.Queue[dict[str, Any]]] = {}
        self._peers: dict[str, dict[str, Peer]] = {}

    async def startup(self) -> None:
        if self._db is not None:
            await self._db.connect()
            await self._db.ensure_schema()

    def _ensure_group(self, group_id: str) -> Group:
        if group_id not in self.groups:
            self.groups[group_id] = Group(group_id=group_id)
        return self.groups[group_id]

    def _build_notification(self, message: Message) -> dict[str, Any]:
        return {
            "type": "message",
            "message_id": message.message_id,
            "group_id": message.group_id,
            "from_display": message.from_display,
            "preview": message.content[:100],
        }

    def _row_to_message(self, row: MessageRow) -> Message:
        return Message(
            message_id=row.message_id,
            from_peer=row.from_peer,
            from_display=row.from_display,
            content=row.content,
            timestamp=row.timestamp,
            group_id=row.group_id,
        )

    def _ensure_peer_cache(self, group_id: str) -> dict[str, Peer]:
        if group_id not in self._peers:
            self._peers[group_id] = {}
        return self._peers[group_id]

    async def _load_peers_into_cache(self, group_id: str) -> dict[str, Peer]:
        assert self._db is not None
        cache = self._ensure_peer_cache(group_id)
        rows = await self._db.get_peers(group_id)
        for row in rows:
            if row.peer_id not in cache:
                cache[row.peer_id] = Peer(
                    peer_id=row.peer_id,
                    hostname=row.hostname,
                    working_directory=row.working_directory,
                    git_branch=row.git_branch,
                    group_id=row.group_id,
                    online=row.peer_id in self.sse_queues,
                )
        return cache

    async def register(
        self,
        session_id: str,
        group_id: str,
        hostname: str,
        working_directory: str,
        git_branch: str | None,
        peer_id: str | None = None,
    ) -> Peer:
        if self._db is not None:
            kwargs: dict[str, Any] = {
                "hostname": hostname,
                "working_directory": working_directory,
                "git_branch": git_branch,
                "group_id": group_id,
            }
            if peer_id is not None:
                kwargs["peer_id"] = peer_id
            peer = Peer(**kwargs)
            if peer.peer_id in self.sse_queues:
                peer.online = True
            await self._db.upsert_peer(
                peer_id=peer.peer_id,
                group_id=group_id,
                hostname=hostname,
                working_directory=working_directory,
                git_branch=git_branch,
            )
            cache = self._ensure_peer_cache(group_id)
            cache[peer.peer_id] = peer
            if session_id in self.sessions:
                _, existing_groups = self.sessions[session_id]
                if group_id not in existing_groups:
                    existing_groups.append(group_id)
            else:
                self.sessions[session_id] = (peer.peer_id, [group_id])
            return peer

        group = self._ensure_group(group_id)
        kwargs = {
            "hostname": hostname,
            "working_directory": working_directory,
            "git_branch": git_branch,
            "group_id": group_id,
        }
        if peer_id is not None:
            kwargs["peer_id"] = peer_id
        peer = Peer(**kwargs)
        if peer.peer_id in self.sse_queues:
            peer.online = True
        group.peers[peer.peer_id] = peer
        if session_id in self.sessions:
            _, existing_groups = self.sessions[session_id]
            if group_id not in existing_groups:
                existing_groups.append(group_id)
        else:
            self.sessions[session_id] = (peer.peer_id, [group_id])
        return peer

    async def ensure_peer_session(
        self,
        session_id: str,
        peer_id: str,
        group_ids: list[str],
        hostname: str,
        working_directory: str,
    ) -> str:
        if session_id in self.sessions:
            existing_peer_id, existing_groups = self.sessions[session_id]
            for gid in group_ids:
                if gid not in existing_groups:
                    await self.register(
                        session_id, gid, hostname, working_directory, None, peer_id=existing_peer_id
                    )
            return existing_peer_id
        for gid in group_ids:
            await self.register(
                session_id, gid, hostname, working_directory, None, peer_id=peer_id
            )
        return peer_id

    async def resolve_session(self, session_id: str) -> tuple[str, list[str]]:
        if session_id not in self.sessions:
            raise ValueError(f"Unknown session: {session_id}")
        return self.sessions[session_id]

    async def send_message(self, group_id: str, from_peer_id: str, content: str) -> Message:
        if self._db is not None:
            cache = self._ensure_peer_cache(group_id)
            if from_peer_id not in cache:
                cache = await self._load_peers_into_cache(group_id)
            peer = cache[from_peer_id]
            message = Message(
                from_peer=from_peer_id,
                from_display=peer.display_name,
                content=content,
                group_id=group_id,
            )
            await self._db.insert_message(
                message_id=message.message_id,
                group_id=group_id,
                from_peer=from_peer_id,
                from_display=peer.display_name,
                content=content,
                timestamp=message.timestamp,
            )
            notification = self._build_notification(message)
            for pid in cache:
                if pid != from_peer_id and pid in self.sse_queues:
                    self.sse_queues[pid].put_nowait(notification)
            return message

        group = self.groups[group_id]
        peer = group.peers[from_peer_id]
        message = Message(
            from_peer=from_peer_id,
            from_display=peer.display_name,
            content=content,
            group_id=group_id,
        )
        group.messages.append(message)
        notification = self._build_notification(message)
        for pid in group.peers:
            if pid != from_peer_id and pid in self.sse_queues:
                self.sse_queues[pid].put_nowait(notification)
        return message

    async def read_new_messages(self, group_id: str, peer_id: str) -> list[Message]:
        if self._db is not None:
            cursor = await self._db.get_cursor(peer_id, group_id)
            rows = await self._db.get_messages_after(group_id, cursor)
            messages = [self._row_to_message(row) for row in rows]
            if rows:
                await self._db.set_cursor(peer_id, group_id, rows[-1].seq)
            return messages

        group = self.groups[group_id]
        peer = group.peers[peer_id]
        new_msgs = group.messages[peer.last_read_index:]
        peer.last_read_index = len(group.messages)
        return new_msgs

    async def list_peers(self, group_id: str) -> list[Peer]:
        if self._db is not None:
            cache = await self._load_peers_into_cache(group_id)
            return list(cache.values())

        return list(self.groups[group_id].peers.values())

    async def get_history(self, group_id: str) -> list[Message]:
        if self._db is not None:
            rows = await self._db.get_all_messages(group_id)
            return [self._row_to_message(row) for row in rows]

        group = self.groups[group_id]
        return list(group.messages)

    def _broadcast_presence(self, peer_id: str, event_type: str) -> None:
        for group_id, group in self.groups.items():
            if peer_id in group.peers:
                peer = group.peers[peer_id]
                notification = {
                    "type": event_type,
                    "message_id": f"{event_type}-{peer_id}",
                    "peer_id": peer_id,
                    "display_name": peer.display_name,
                    "group_id": group_id,
                }
                for pid, q in self.sse_queues.items():
                    if pid != peer_id:
                        q.put_nowait(notification)
        for group_id, cache in self._peers.items():
            if peer_id in cache:
                peer = cache[peer_id]
                notification = {
                    "type": event_type,
                    "message_id": f"{event_type}-{peer_id}",
                    "peer_id": peer_id,
                    "display_name": peer.display_name,
                    "group_id": group_id,
                }
                for pid, q in self.sse_queues.items():
                    if pid != peer_id:
                        q.put_nowait(notification)

    async def connect_sse(self, peer_id: str) -> asyncio.Queue[dict[str, Any]]:
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self.sse_queues[peer_id] = queue
        for group in self.groups.values():
            if peer_id in group.peers:
                group.peers[peer_id].online = True
        for group_cache in self._peers.values():
            if peer_id in group_cache:
                group_cache[peer_id].online = True
        self._broadcast_presence(peer_id, "peer_joined")
        return queue

    async def disconnect_sse(self, peer_id: str) -> None:
        self._broadcast_presence(peer_id, "peer_left")
        self.sse_queues.pop(peer_id, None)
        for group in self.groups.values():
            if peer_id in group.peers:
                group.peers[peer_id].online = False
        for group_cache in self._peers.values():
            if peer_id in group_cache:
                group_cache[peer_id].online = False

    async def leave_group(self, session_id: str, group_id: str) -> None:
        if session_id not in self.sessions:
            return
        peer_id, group_ids = self.sessions[session_id]
        if group_id not in group_ids:
            return
        group_ids.remove(group_id)
        if not group_ids:
            self.sessions.pop(session_id)
            self.sse_queues.pop(peer_id, None)
        if self._db is not None:
            await self._db.delete_peer(peer_id, group_id)
            if group_id in self._peers:
                self._peers[group_id].pop(peer_id, None)
        else:
            if group_id in self.groups:
                self.groups[group_id].peers.pop(peer_id, None)
