from __future__ import annotations

import asyncio
import base64
import contextlib
import json
import os
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from .nat import NatResult, try_upnp_map
from .state import SharedState

MessageHandler = Callable[[str, str, dict[str, Any]], Awaitable[None] | None]


@dataclass
class _Peer:
    peer_id: str
    reader: asyncio.StreamReader
    writer: asyncio.StreamWriter


@dataclass
class Playground:
    """Dead-simple peer-to-peer toolkit.

    Use one `Playground` object per running node.
    """

    app_name: str
    host: str = "0.0.0.0"
    port: int = 7337
    peer_id: str = field(default_factory=lambda: uuid.uuid4().hex[:10])
    state: SharedState = field(default_factory=SharedState)

    _server: asyncio.AbstractServer | None = field(init=False, default=None)
    _peers: dict[str, _Peer] = field(init=False, default_factory=dict)
    _handlers: dict[str, list[MessageHandler]] = field(init=False, default_factory=dict)

    async def start(self, open_port: bool = True) -> NatResult | None:
        self._server = await asyncio.start_server(self._on_accept, self.host, self.port)
        if open_port:
            return try_upnp_map(self.port)
        return None

    def quick_start(self, open_port: bool = True) -> "Playground":
        """Start this node immediately for sync-first scripts."""
        asyncio.run(self.start(open_port=open_port))
        return self

    async def stop(self) -> None:
        for peer in list(self._peers.values()):
            peer.writer.close()
            await peer.writer.wait_closed()
        self._peers.clear()

        if self._server:
            self._server.close()
            await self._server.wait_closed()
            self._server = None

    async def join(self, host: str, port: int) -> str:
        reader, writer = await asyncio.open_connection(host, port)
        await self._write(writer, {"kind": "hello", "peer_id": self.peer_id, "app": self.app_name})
        hello = await self._read(reader)
        remote_peer_id = hello["peer_id"]
        self._peers[remote_peer_id] = _Peer(remote_peer_id, reader, writer)
        asyncio.create_task(self._peer_loop(remote_peer_id))
        await self.push_state(remote_peer_id)
        return remote_peer_id

    async def send(self, kind: str, payload: dict[str, Any], to: str | None = None) -> None:
        message = {"kind": kind, "from": self.peer_id, "payload": payload}
        if to is None:
            for peer in list(self._peers.values()):
                await self._write(peer.writer, message)
            return

        peer = self._peers[to]
        await self._write(peer.writer, message)

    async def send_text(self, text: str, to: str | None = None) -> None:
        await self.send("text", {"text": text}, to=to)

    def on(self, kind: str, handler: MessageHandler) -> None:
        self._handlers.setdefault(kind, []).append(handler)

    async def share_file(self, file_path: str | os.PathLike[str], to: str | None = None) -> None:
        path = Path(file_path)
        data = path.read_bytes()
        payload = {
            "name": path.name,
            "bytes_b64": base64.b64encode(data).decode("ascii"),
        }
        await self.send("file", payload, to=to)

    async def request_state(self, to: str | None = None) -> None:
        await self.send("state_request", {}, to=to)

    async def push_state(self, to: str | None = None) -> None:
        await self.send("state_update", {"snapshot": self.state.snapshot()}, to=to)

    async def _on_accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        hello = await self._read(reader)
        remote_peer_id = hello["peer_id"]
        await self._write(writer, {"kind": "hello", "peer_id": self.peer_id, "app": self.app_name})
        self._peers[remote_peer_id] = _Peer(remote_peer_id, reader, writer)
        asyncio.create_task(self._peer_loop(remote_peer_id))

    async def _peer_loop(self, remote_peer_id: str) -> None:
        peer = self._peers[remote_peer_id]
        try:
            while True:
                msg = await self._read(peer.reader)
                kind = msg.get("kind", "")
                src = msg.get("from", remote_peer_id)
                payload = msg.get("payload", {})

                if kind == "state_request":
                    await self.push_state(src)
                    continue

                if kind == "state_update":
                    self.state.apply_remote_snapshot(payload.get("snapshot", {}))
                    continue

                for handler in self._handlers.get(kind, []):
                    out = handler(src, kind, payload)
                    if asyncio.iscoroutine(out):
                        await out
        except (asyncio.IncompleteReadError, ConnectionResetError, KeyError):
            pass
        finally:
            peer.writer.close()
            with contextlib.suppress(Exception):
                await peer.writer.wait_closed()
            self._peers.pop(remote_peer_id, None)

    async def _read(self, reader: asyncio.StreamReader) -> dict[str, Any]:
        line = await reader.readline()
        if not line:
            raise asyncio.IncompleteReadError(partial=b"", expected=1)
        return json.loads(line.decode("utf-8"))

    async def _write(self, writer: asyncio.StreamWriter, data: dict[str, Any]) -> None:
        writer.write(json.dumps(data, ensure_ascii=True).encode("utf-8") + b"\n")
        await writer.drain()
