from __future__ import annotations

import asyncio
import contextlib
from dataclasses import dataclass, field
from typing import Dict

from .protocol import AUTH, PING, PONG, STATUS_ERR, STATUS_OK, TCP_CLOSE, TCP_DATA, TCP_OPEN, TCP_STATUS, Frame, decode_json, encode_json, read_frame, write_frame


@dataclass
class TcpSession:
    session_id: int
    stream_id: int
    writer: asyncio.StreamWriter
    task: asyncio.Task


@dataclass
class TcpRelayChannel:
    reader: asyncio.StreamReader
    writer: asyncio.StreamWriter
    token: str
    tcp_sessions: Dict[tuple[int, int], TcpSession] = field(default_factory=dict)
    closed: bool = False
    send_lock: asyncio.Lock = field(default_factory=asyncio.Lock)

    async def run(self) -> None:
        try:
            auth = await read_frame(self.reader)
            if auth.kind != AUTH:
                return
            payload = decode_json(auth.payload) if auth.payload else {}
            if payload.get("token") != self.token:
                return
            await self.safe_send(Frame(AUTH, 0, 0, 0, STATUS_OK, encode_json({"status": "ok", "kind": payload.get("purpose", "normal"), "udp_only": False})))
            while True:
                frame = await read_frame(self.reader)
                await self.handle(frame)
        finally:
            await self.close()

    async def safe_send(self, frame: Frame) -> bool:
        if self.closed:
            return False
        try:
            async with self.send_lock:
                if self.closed:
                    return False
                await write_frame(self.writer, frame)
            return True
        except Exception:
            return False

    async def handle(self, frame: Frame) -> None:
        key = (frame.session_id, frame.stream_id)
        if frame.kind == PING:
            await self.safe_send(Frame(PONG, 0, 0, frame.seq, 0, b"pong"))
            return
        if frame.kind == TCP_OPEN:
            try:
                meta = decode_json(frame.payload) if frame.payload else {}
                family = int(meta.get("family", 0)) or 0
                reader, writer = await asyncio.open_connection(meta["host"], int(meta["port"]), family=family or 0)
                task = asyncio.create_task(self._tcp_pump(frame.session_id, frame.stream_id, reader))
                self.tcp_sessions[key] = TcpSession(frame.session_id, frame.stream_id, writer, task)
                await self.safe_send(Frame(TCP_STATUS, frame.session_id, frame.stream_id, 0, STATUS_OK, b"ok"))
            except Exception as exc:
                await self.safe_send(Frame(TCP_STATUS, frame.session_id, frame.stream_id, 0, STATUS_ERR, str(exc).encode()))
            return
        if frame.kind == TCP_DATA:
            session = self.tcp_sessions.get(key)
            if session:
                try:
                    session.writer.write(frame.payload)
                    await session.writer.drain()
                except Exception:
                    await self._close_tcp(key)
            return
        if frame.kind == TCP_CLOSE:
            await self._close_tcp(key)

    async def _tcp_pump(self, session_id: int, stream_id: int, reader: asyncio.StreamReader) -> None:
        try:
            while True:
                chunk = await reader.read(65536)
                if not chunk:
                    break
                sent = await self.safe_send(Frame(TCP_DATA, session_id, stream_id, 0, 0, chunk))
                if not sent:
                    break
        finally:
            if not self.closed:
                await self.safe_send(Frame(TCP_CLOSE, session_id, stream_id, 0, 0, b""))
            await self._close_tcp((session_id, stream_id), from_task=True)

    async def _close_tcp(self, key: tuple[int, int], from_task: bool = False) -> None:
        session = self.tcp_sessions.pop(key, None)
        if session is None:
            return
        if not from_task and session.task is not asyncio.current_task():
            session.task.cancel()
            with contextlib.suppress(Exception):
                await session.task
        session.writer.close()
        with contextlib.suppress(Exception):
            await session.writer.wait_closed()

    async def close(self) -> None:
        if self.closed:
            return
        self.closed = True
        for key in list(self.tcp_sessions):
            await self._close_tcp(key)
        self.writer.close()
        with contextlib.suppress(Exception):
            await self.writer.wait_closed()


class TcpRelayServer:
    def __init__(self, token: str) -> None:
        self.token = token

    async def start(self, host: str, port: int) -> None:
        async def accept(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
            await TcpRelayChannel(reader, writer, self.token).run()

        server = await asyncio.start_server(accept, host, port)
        async with server:
            await server.serve_forever()