import ssl
import asyncio
import httpx
from typing import List
from pathlib import Path


from aiorwlock import RWLock
from google.protobuf.message import Message

from GozargahNodeBridge.controller import NodeAPIError, Health
from GozargahNodeBridge.common import service_pb2 as service
from GozargahNodeBridge.abstract_node import GozargahNode
from GozargahNodeBridge.utils import string_to_temp_file


class Node(GozargahNode):
    def __init__(self, address: str, port: int, client_cert: str, client_key: str, server_ca: str, extra: dict = {}):
        super().__init__(client_cert, client_key, server_ca, extra)

        url = f"https://{address.strip('/')}:{port}/"

        self._client = httpx.AsyncClient(
            http2=True,
            verify=self.ctx,
            headers={"Content-Type": "application/x-protobuf"},
            base_url=url,
            timeout=httpx.Timeout(None),
            limits=httpx.Limits(max_keepalive_connections=1),
        )

        self._node_lock = RWLock()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.stop()
        await self._client.aclose()
        self._cleanup_temp_files()

    def __del__(self):
        self._cleanup_temp_files()

    def _serialize_protobuf(self, proto_message: Message) -> bytes:
        """Serialize a protobuf message to bytes."""
        return proto_message.SerializeToString()

    def _deserialize_protobuf(self, proto_class: type[Message], data: bytes) -> Message:
        """Deserialize bytes into a protobuf message."""
        proto_instance = proto_class()
        proto_instance.ParseFromString(data)
        return proto_instance

    def _handle_error(error: Exception):
        if isinstance(error, httpx.RemoteProtocolError):
            raise NodeAPIError(code=-1, detail=f"Server closed connection: {error}")
        elif isinstance(error, httpx.HTTPStatusError):
            raise NodeAPIError(code=error.response.status_code, detail=f"HTTP error: {error.response.text}")
        elif isinstance(error, httpx.ConnectError) or isinstance(error, httpx.ReadTimeout):
            raise NodeAPIError(code=-1, detail=f"Connection error: {error}")
        else:
            raise NodeAPIError(0, str(error))

    async def _make_request(
        self,
        method: str,
        endpoint: str,
        timeout: int,
        proto_message: Message = None,
        proto_response_class: type[Message] = None,
    ) -> Message:
        """Handle common REST API call logic with protobuf support (async)."""
        request_data = None

        if proto_message:
            request_data = self._serialize_protobuf(proto_message)

        try:
            response = await self._client.request(
                method=method,
                url=endpoint,
                content=request_data,
                timeout=timeout,
            )
            response.raise_for_status()

            if proto_response_class:
                return self._deserialize_protobuf(proto_response_class, response.content)
            return response.content

        except Exception as e:
            self._handle_error(e)

    async def start(self, config: str, backend_type: service.BackendType, users: List[service.User]):
        try:
            await self.connected()
            await self.stop()
        except:
            pass

        async with self._node_lock.writer_lock:
            response = await self._make_request(
                method="POST",
                endpoint="start",
                timeout=15,
                proto_message=service.Backend(type=backend_type, config=config, users=users),
                proto_response_class=service.BaseInfoResponse,
            )

            self._client.headers["Authorization"] = f"Bearer {response.session_id}"
            await self.connect(
                response.node_version,
                response.core_version,
                [
                    asyncio.create_task(self._check_node_health()),
                    asyncio.create_task(self.sync_user()),
                    asyncio.create_task(self._fetch_logs()),
                ],
            )

        return response

    async def stop(self) -> None:
        await self.connected()
        async with self._node_lock.writer_lock:
            await self.disconnect()
            await self._make_request(method="PUT", endpoint="stop", timeout=10)

    async def info(self) -> service.BaseInfoResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(method="GET", endpoint="info", timeout=10)

    async def get_system_stats(self) -> service.SystemStatsResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(method="GET", endpoint="stats/system", timeout=10)

    async def get_backend_stats(self) -> service.BackendStatsResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(method="GET", endpoint="stats/backend", timeout=10)

    async def get_outbounds_stats(self, reset: bool = True) -> service.StatResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(
                method="GET",
                endpoint="stats/outbounds",
                timeout=10,
                proto_message=service.StatRequest(reset=reset),
                proto_response_class=service.StatResponse,
            )

    async def get_outbound_stats(self, tag: str, reset: bool = True) -> service.StatResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(
                method="GET",
                endpoint="stats/outbound",
                timeout=10,
                proto_message=service.StatRequest(reset=reset, name=tag),
                proto_response_class=service.StatResponse,
            )

    async def get_inbounds_stats(self, reset: bool = True) -> service.StatResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(
                method="GET",
                endpoint="stats/inbounds",
                timeout=10,
                proto_message=service.StatRequest(reset=reset),
                proto_response_class=service.StatResponse,
            )

    async def get_inbound_stats(self, tag: str, reset: bool = True) -> service.StatResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(
                method="GET",
                endpoint="stats/inbound",
                timeout=10,
                proto_message=service.StatRequest(reset=reset, name=tag),
                proto_response_class=service.StatResponse,
            )

    async def get_users_stats(self, reset: bool = True) -> service.StatResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(
                method="GET",
                endpoint="stats/users",
                timeout=10,
                proto_message=service.StatRequest(reset=reset),
                proto_response_class=service.StatResponse,
            )

    async def get_user_stats(self, email: str, reset: bool = True) -> service.StatResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(
                method="GET",
                endpoint="stats/user",
                timeout=10,
                proto_message=service.StatRequest(reset=reset, name=email),
                proto_response_class=service.StatResponse,
            )

    async def user_online_stats(self, email: str) -> service.OnlineStatResponse | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(
                method="GET",
                endpoint=f"stats/user/online",
                timeout=10,
                proto_message=service.StatRequest(name=email),
                proto_response_class=service.OnlineStatResponse,
            )

    async def sync_users(self, users: List[service.User]) -> service.Empty | None:
        await self.connected()
        async with self._node_lock.reader_lock:
            return await self._make_request(
                method="POST",
                endpoint="users/sync",
                timeout=30,
                proto_message=service.Users(users=users),
                proto_response_class=service.Empty,
            )

    async def _check_node_health(self):
        while True:
            last_health = await self.get_health()

            try:
                await self.get_backend_stats()
                if last_health != Health.HEALTHY:
                    await self.set_health(Health.HEALTHY)
            except Exception as e:
                if last_health != Health.BROKEN:
                    await self.set_health(Health.BROKEN)

            await asyncio.sleep(2)

    async def _fetch_logs(self):
        while True:
            health = await self.get_health()
            if health == Health.BROKEN:
                await asyncio.sleep(5)
                continue
            elif health == Health.NOT_CONNECTED:
                return

            try:
                async with self._client.stream("GET", "/logs", timeout=None) as response:
                    buffer = b""

                    async for chunk in response.aiter_bytes():
                        buffer += chunk

                        while b"\n" in buffer:
                            line, buffer = buffer.split(b"\n", 1)
                            line = line.decode().strip()

                            if line:
                                if self._logs_queue:
                                    await self._logs_queue.put(line)

            except Exception:
                asyncio.sleep(3)
                continue

    async def sync_user(self) -> None:
        while True:
            health = await self.get_health()

            if health == Health.BROKEN:
                await asyncio.sleep(5)
                continue
            elif health == Health.NOT_CONNECTED:
                return

            async with self._lock.reader_lock:
                if self._user_queue is None or self._notify_queue is None:
                    return

                user_task = asyncio.create_task(self._user_queue.get())
                notify_task = asyncio.create_task(self._notify_queue.get())

            try:
                # Wait for first completed operation
                done, pending = await asyncio.wait([user_task, notify_task], return_when=asyncio.FIRST_COMPLETED)

            except asyncio.CancelledError:
                # Cleanup if task is cancelled
                user_task.cancel()
                notify_task.cancel()
                raise

            # Cancel pending tasks to avoid leaks
            for task in pending:
                task.cancel()

            if notify_task in done:
                continue

            # Handle UserChan message
            if user_task in done:
                user = user_task.result()

                try:
                    await self._make_request(
                        method="PUT",
                        endpoint="user/sync",
                        timeout=10,
                        proto_message=user,
                        proto_response_class=service.Empty,
                    )
                except Exception:
                    asyncio.sleep(5)
                    continue
