import gzip
import inspect
import json
import logging
import struct
import typing

from httpcore import (
    ConnectionPool,
    AsyncConnectionPool,
    RemoteProtocolError,
    Response,
)
from enum import Flag, Enum
from typing import Callable, Optional, Dict, Any, Generator, Tuple
from google.protobuf import json_format


class EnvelopeFlags(Flag):
    compressed = 0b00000001
    end_stream = 0b00000010


class Code(Enum):
    canceled = "canceled"
    unknown = "unknown"
    invalid_argument = "invalid_argument"
    deadline_exceeded = "deadline_exceeded"
    not_found = "not_found"
    already_exists = "already_exists"
    permission_denied = "permission_denied"
    resource_exhausted = "resource_exhausted"
    failed_precondition = "failed_precondition"
    aborted = "aborted"
    out_of_range = "out_of_range"
    unimplemented = "unimplemented"
    internal = "internal"
    unavailable = "unavailable"
    data_loss = "data_loss"
    unauthenticated = "unauthenticated"


def make_error_from_http_code(http_code: int):
    error_code_map = {
        400: Code.invalid_argument,
        401: Code.unauthenticated,
        403: Code.permission_denied,
        404: Code.not_found,
        409: Code.already_exists,
        413: Code.resource_exhausted,
        429: Code.resource_exhausted,
        499: Code.canceled,
        500: Code.internal,
        501: Code.unimplemented,
        502: Code.unavailable,
        503: Code.unavailable,
        504: Code.deadline_exceeded,
        505: Code.unimplemented,
    }

    return error_code_map.get(http_code, Code.unknown)


class ConnectException(Exception):
    def __init__(self, status: Code, message: str):
        self.status = status
        self.message = message


envelope_header_length = 5
envelope_header_pack = ">BI"


def encode_envelope(*, flags: EnvelopeFlags, data):
    return encode_envelope_header(flags=flags.value, data=data) + data


def encode_envelope_header(*, flags, data):
    return struct.pack(envelope_header_pack, flags, len(data))


def decode_envelope_header(header):
    flags, data_len = struct.unpack(envelope_header_pack, header)
    return EnvelopeFlags(flags), data_len


def error_for_response(http_resp: Response):
    try:
        error = json.loads(http_resp.content)
        return make_error(error)
    except (json.decoder.JSONDecodeError, KeyError):
        error = {"code": http_resp.status, "message": http_resp.content.decode("utf-8")}
        return make_error(error)


def make_error(error):
    status = None
    try:
        code_value = error.get("code")
        # return error code from http status code
        if isinstance(code_value, int):
            status = make_error_from_http_code(code_value)
        else:
            status = Code(code_value)
    except (KeyError, ValueError):
        status = Code.unknown

    return ConnectException(status, error.get("message", ""))


def _sync_retry(func, exc, retries):
    def retry(*args, **kwargs):
        for _ in range(retries):
            try:
                return func(*args, **kwargs)
            except exc:
                continue

        return func(*args, **kwargs)

    return retry


def _async_retry(func, exc, retries):
    async def retry(*args, **kwargs):
        for _ in range(retries):
            try:
                return await func(*args, **kwargs)
            except exc:
                continue

        return await func(*args, **kwargs)

    return retry


def _retry(exc: typing.Type[Exception], retries: int):
    def decorator(func):
        if inspect.iscoroutinefunction(func):
            return _async_retry(func, exc, retries)

        return _sync_retry(func, exc, retries)

    return decorator


class GzipCompressor:
    name = "gzip"
    decompress = gzip.decompress
    compress = gzip.compress


class JSONCodec:
    content_type = "json"

    @staticmethod
    def encode(msg):
        return json_format.MessageToJson(msg).encode("utf8")

    @staticmethod
    def decode(data, *, msg_type):
        msg = msg_type()
        json_format.Parse(data.decode("utf8"), msg, ignore_unknown_fields=True)
        return msg


class ProtobufCodec:
    content_type = "proto"

    @staticmethod
    def encode(msg):
        return msg.SerializeToString()

    @staticmethod
    def decode(data, *, msg_type):
        msg = msg_type()
        msg.ParseFromString(data)
        return msg


class Client:
    def __init__(
        self,
        *,
        pool: Optional[ConnectionPool] = None,
        async_pool: Optional[AsyncConnectionPool] = None,
        url: str,
        response_type,
        compressor=None,
        json: Optional[bool] = False,
        headers: Optional[Dict[str, str]] = None,
        logger: Optional[logging.Logger] = None,
    ):
        if headers is None:
            headers = {}

        self.pool = pool
        self.async_pool = async_pool
        self.url = url
        self._codec = JSONCodec if json else ProtobufCodec
        self._response_type = response_type
        self._compressor = compressor
        self._headers = headers
        self._connection_retries = 3
        self._logger = logger

    def _log_request(self) -> None:
        if self._logger is not None:
            self._logger.info(f"Request: POST {self.url}")

    def _log_response(self, status: int) -> None:
        if self._logger is None:
            return
        if status >= 400:
            self._logger.error(f"Response: {status} {self.url}")
        else:
            self._logger.info(f"Response: {status} {self.url}")

    def _log_stream_message(self) -> None:
        if self._logger is not None:
            self._logger.debug(f"Response stream: {self.url}")

    def _prepare_unary_request(
        self,
        req,
        request_timeout=None,
        headers: Optional[dict] = None,
        **opts,
    ) -> dict:
        data = self._codec.encode(req)

        if self._compressor is not None:
            data = self._compressor.compress(data)

        if headers is None:
            headers = {}

        extensions = (
            None
            if request_timeout is None
            else {
                "timeout": {
                    "connect": request_timeout,
                    "pool": request_timeout,
                    "read": request_timeout,
                    "write": request_timeout,
                }
            }
        )

        return {
            "method": "POST",
            "url": self.url,
            "content": data,
            "extensions": extensions,
            "headers": {
                **self._headers,
                **headers,
                **opts.get("headers", {}),
                "connect-protocol-version": "1",
                "content-encoding": (
                    "identity" if self._compressor is None else self._compressor.name
                ),
                "content-type": f"application/{self._codec.content_type}",
            },
        }

    def _process_unary_response(
        self,
        http_resp: Response,
    ):
        self._log_response(http_resp.status)

        if http_resp.status != 200:
            raise error_for_response(http_resp)

        content = http_resp.content

        if self._compressor is not None:
            content = self._compressor.decompress(content)

        return self._codec.decode(
            content,
            msg_type=self._response_type,
        )

    @_retry(RemoteProtocolError, 3)
    async def acall_unary(
        self,
        req,
        request_timeout=None,
        headers: Optional[dict] = None,
        **opts,
    ):
        if self.async_pool is None:
            raise ValueError("async_pool is required")

        req_data = self._prepare_unary_request(
            req,
            request_timeout,
            headers,
            **opts,
        )

        self._log_request()
        res = await self.async_pool.request(**req_data)
        return self._process_unary_response(res)

    @_retry(RemoteProtocolError, 3)
    def call_unary(
        self,
        req,
        request_timeout=None,
        headers: Optional[dict] = None,
        **opts,
    ):
        if self.pool is None:
            raise ValueError("pool is required")

        req_data = self._prepare_unary_request(
            req,
            request_timeout,
            headers,
            **opts,
        )

        self._log_request()
        res = self.pool.request(**req_data)
        return self._process_unary_response(res)

    def _create_stream_timeout(self, timeout: Optional[float]):
        if timeout:
            return {"connect-timeout-ms": str(int(timeout * 1000))}
        return {}

    def _prepare_server_stream_request(
        self,
        req,
        request_timeout=None,
        timeout=None,
        headers: Optional[dict] = None,
        **opts,
    ) -> dict:
        headers = headers or {}
        data = self._codec.encode(req)
        flags = EnvelopeFlags(0)

        # `request_timeout` bounds connection setup and request sending, but NOT the
        # stream read: a stream can stay open for the whole command `timeout` (minutes
        # or, when disabled, indefinitely), so we deliberately leave `read` unset.
        # The command `timeout` is enforced server-side via the `connect-timeout-ms`
        # header (see `_create_stream_timeout`), which returns a clean `deadline_exceeded`.
        # This mirrors the JS SDK, which has no per-chunk read timeout either — setting
        # `read` to the command `timeout` would race that server response and surface a
        # raw transport `ReadTimeout` instead.
        timeout_ext = {}
        if request_timeout is not None:
            timeout_ext["connect"] = request_timeout
            timeout_ext["pool"] = request_timeout
            timeout_ext["write"] = request_timeout
        extensions = {"timeout": timeout_ext} if timeout_ext else None

        if self._compressor is not None:
            data = self._compressor.compress(data)
            flags |= EnvelopeFlags.compressed

        stream_timeout = self._create_stream_timeout(timeout)

        return {
            "method": "POST",
            "url": self.url,
            "content": encode_envelope(
                flags=flags,
                data=data,
            ),
            "extensions": extensions,
            "headers": {
                **self._headers,
                **headers,
                **opts.get("headers", {}),
                **stream_timeout,
                "connect-protocol-version": "1",
                "connect-content-encoding": (
                    "identity" if self._compressor is None else self._compressor.name
                ),
                "content-type": f"application/connect+{self._codec.content_type}",
            },
        }

    # Note: no retry here — generator functions don't execute until iterated, so a
    # call-level retry never fires, and retrying mid-stream would replay delivered events.
    async def acall_server_stream(
        self,
        req,
        request_timeout=None,
        timeout=None,
        headers: Optional[dict] = None,
        **opts,
    ):
        if self.async_pool is None:
            raise ValueError("async_pool is required")

        req_data = self._prepare_server_stream_request(
            req,
            request_timeout,
            timeout,
            headers,
            **opts,
        )

        parser = ServerStreamParser(
            decode=self._codec.decode,
            response_type=self._response_type,
        )

        self._log_request()
        async with self.async_pool.stream(**req_data) as http_resp:
            if http_resp.status != 200:
                self._log_response(http_resp.status)
                await http_resp.aread()
                raise error_for_response(http_resp)

            async for chunk in http_resp.aiter_stream():
                for parsed in parser.parse(chunk):
                    self._log_stream_message()
                    yield parsed

    def call_server_stream(
        self,
        req,
        request_timeout=None,
        timeout=None,
        headers: Optional[dict] = None,
        **opts,
    ):
        if self.pool is None:
            raise ValueError("pool is required")

        req_data = self._prepare_server_stream_request(
            req,
            request_timeout,
            timeout,
            headers,
            **opts,
        )

        parser = ServerStreamParser(
            decode=self._codec.decode,
            response_type=self._response_type,
        )

        self._log_request()
        with self.pool.stream(**req_data) as http_resp:
            if http_resp.status != 200:
                self._log_response(http_resp.status)
                http_resp.read()
                raise error_for_response(http_resp)

            for chunk in http_resp.iter_stream():
                for parsed in parser.parse(chunk):
                    self._log_stream_message()
                    yield parsed

    def call_client_stream(self, req, **opts):
        raise NotImplementedError("client stream not supported")

    def acall_client_stream(self, req, **opts):
        raise NotImplementedError("client stream not supported")

    def call_bidi_stream(self, req, **opts):
        raise NotImplementedError("bidi stream not supported")

    def acall_bidi_stream(self, req, **opts):
        raise NotImplementedError("bidi stream not supported")


DataLen = int


class ServerStreamParser:
    def __init__(
        self,
        decode: Callable,
        response_type: Any,
    ):
        self.decode = decode
        self.response_type = response_type

        self.buffer: bytes = b""
        self._header: Optional[tuple[EnvelopeFlags, DataLen]] = None

    def shift_buffer(self, size: int):
        buffer = self.buffer[:size]
        self.buffer = self.buffer[size:]
        return buffer

    @property
    def header(self) -> Tuple[EnvelopeFlags, DataLen]:
        if self._header:
            return self._header

        header_data = self.shift_buffer(envelope_header_length)
        self._header = decode_envelope_header(header_data)

        return self._header

    @header.deleter
    def header(self):
        self._header = None

    def parse(self, chunk: bytes) -> Generator[Any, None, None]:
        self.buffer += chunk

        # Once the header is consumed, the remaining payload can be shorter
        # than the header length, so only require a full header when we still
        # need to read one.
        while self._header is not None or len(self.buffer) >= envelope_header_length:
            flags, data_len = self.header

            if data_len > len(self.buffer):
                break

            data = self.shift_buffer(data_len)

            if EnvelopeFlags.end_stream in flags:
                data = json.loads(data)

                if "error" in data:
                    raise make_error(data["error"])

                return

            yield self.decode(data, msg_type=self.response_type)
            del self.header