from __future__ import annotations

import socket
import ssl
import threading
from urllib.parse import urlparse

from h2.config import H2Configuration
from h2.connection import H2Connection
from h2.events import (
    ConnectionTerminated,
    DataReceived,
    ResponseReceived,
    StreamEnded,
    StreamReset,
    WindowUpdated,
)

from hypertube.utils.connectionData.Http2ConnectionData import Http2ConnectionData

# cache of clients keyed by connection identity
_clients = {}
_clients_lock = threading.Lock()
MAX_ATTEMPTS = 2


class Http2ClientDisposedError(RuntimeError):
    pass


class Http2ProtocolError(RuntimeError):
    pass


class Http2ResponseError(RuntimeError):
    pass


class Http2TransportError(RuntimeError):
    pass


class Http2Client:
    def __init__(self, connection_data: Http2ConnectionData):
        self.connection_data = connection_data
        self.uri = connection_data.hostname
        self._parsed_uri = urlparse(self.uri)
        if self._parsed_uri.scheme.lower() not in ("http", "https"):
            raise ValueError(
                "Http2Client requires 'http' (h2c) or 'https' (h2 over TLS) scheme. "
                f"Provided connection data: {connection_data}"
            )
        if not self._parsed_uri.hostname:
            raise ValueError(f"HTTP/2 connection data contains invalid URI: {self.uri}")

        self.timeout_in_seconds = connection_data.timeout_in_seconds
        self._send_receive_lock = threading.Lock()
        self._socket = None
        self._connection = None
        self._disposed = False

    @staticmethod
    def _connection_key(connection_data: Http2ConnectionData) -> str:
        return connection_data.to_canonical_key()

    @staticmethod
    def create_or_get_client(connection_data: Http2ConnectionData, expected_current=None):
        key = Http2Client._connection_key(connection_data)
        with _clients_lock:
            if expected_current is None:
                current = _clients.get(key)
                if current is not None:
                    return current
                created = Http2Client(connection_data)
                _clients[key] = created
                return created

            current = _clients.get(key)
            if current is expected_current:
                current.close()
                _clients.pop(key, None)
            elif current is not None:
                return current
            created = Http2Client(connection_data)
            _clients[key] = created
            return created

    @staticmethod
    def send_message(connection_data: Http2ConnectionData, message: bytes) -> bytes:
        if connection_data is None:
            raise ValueError("connection_data must not be None")
        if message is None:
            raise ValueError("message must not be None")

        client = Http2Client.create_or_get_client(connection_data, None)

        last_ex = None
        for attempt in range(MAX_ATTEMPTS):
            try:
                with client._send_receive_lock:
                    return client.send_byte_array(message)
            except (OSError, TimeoutError, Http2ClientDisposedError, Http2TransportError) as ex:
                last_ex = ex
                if attempt == MAX_ATTEMPTS - 1:
                    break
                client = Http2Client.create_or_get_client(connection_data, client)

        raise RuntimeError(
            f"HTTP/2 send to {connection_data.hostname} failed after {MAX_ATTEMPTS} attempts. "
            f"Last error: {last_ex}"
        ) from last_ex

    def send_byte_array(self, message: bytes) -> bytes:
        self._ensure_open()
        if message is None:
            raise ValueError("message must not be None")

        self._ensure_connected()
        sock = self._socket
        connection = self._connection
        if sock is None or connection is None:
            raise Http2TransportError(f"HTTP/2 connection to {self.uri} is not open.")

        response_body = bytearray()
        response_status = None
        response_reason = ""
        stream_ended = False

        stream_id = connection.get_next_available_stream_id()
        connection.send_headers(
            stream_id,
            self._build_headers(message),
            end_stream=len(message) == 0,
        )
        self._send_pending_data(sock, connection)

        if message:
            offset = 0
            while offset < len(message):
                flow_control_window = connection.local_flow_control_window(stream_id)
                max_chunk_size = min(flow_control_window, connection.max_outbound_frame_size)
                if max_chunk_size <= 0:
                    stream_ended, response_status, response_reason = self._receive_events(
                        sock,
                        connection,
                        stream_id,
                        response_body,
                        response_status,
                        response_reason,
                    )
                    if stream_ended:
                        raise Http2TransportError(f"HTTP/2 stream ended before request body was sent. URL: {self.uri}")
                    continue

                chunk = message[offset : offset + max_chunk_size]
                offset += len(chunk)
                connection.send_data(stream_id, chunk, end_stream=offset == len(message))
                self._send_pending_data(sock, connection)

        while not stream_ended:
            stream_ended, response_status, response_reason = self._receive_events(
                sock,
                connection,
                stream_id,
                response_body,
                response_status,
                response_reason,
            )

        if response_status is None:
            raise Http2ProtocolError(f"HTTP/2 request to {self.uri} completed without a response status.")
        if response_status < 200 or response_status >= 300:
            raise Http2ResponseError(
                f"HTTP/2 request to {self.uri} failed with status {response_status} {response_reason}. "
            )
        return bytes(response_body)

    def _ensure_connected(self) -> None:
        if self._socket is not None and self._connection is not None:
            return

        sock = self._open_socket()
        configuration = H2Configuration(client_side=True, header_encoding="utf-8")
        connection = H2Connection(config=configuration)
        connection.initiate_connection()
        self._send_pending_data(sock, connection)
        self._socket = sock
        self._connection = connection

    def _open_socket(self):
        host = self._parsed_uri.hostname
        port = self._parsed_uri.port or (443 if self._parsed_uri.scheme.lower() == "https" else 80)
        socket_timeout = None if self.timeout_in_seconds < 0 else self.timeout_in_seconds
        raw_socket = socket.create_connection((host, port), timeout=socket_timeout)
        raw_socket.settimeout(socket_timeout)

        if self._parsed_uri.scheme.lower() == "https":
            context = ssl.create_default_context()
            context.set_alpn_protocols(["h2"])
            try:
                tls_socket = context.wrap_socket(raw_socket, server_hostname=host)
            except Exception:
                raw_socket.close()
                raise
            if tls_socket.selected_alpn_protocol() != "h2":
                tls_socket.close()
                raise Http2ProtocolError(f"The endpoint did not negotiate HTTP/2. URL: {self.uri}")
            return tls_socket

        return raw_socket

    def _build_headers(self, message: bytes):
        custom_headers = self.connection_data.get_headers()
        content_type = "application/octet-stream"
        normalized_headers = []
        for key, value in custom_headers.items():
            if key is None or value is None:
                continue
            normalized_key = str(key).lower()
            if normalized_key.startswith(":"):
                continue
            if normalized_key == "content-type":
                content_type = str(value)
                continue
            normalized_headers.append((normalized_key, str(value)))

        path = self._parsed_uri.path or "/"
        if self._parsed_uri.query:
            path = f"{path}?{self._parsed_uri.query}"

        headers = [
            (":method", "POST"),
            (":scheme", self._parsed_uri.scheme.lower()),
            (":authority", self._authority_header()),
            (":path", path),
            ("content-type", content_type),
            ("content-length", str(len(message))),
        ]
        headers.extend(normalized_headers)
        return headers

    def _authority_header(self) -> str:
        host = self._parsed_uri.hostname or ""
        port = self._parsed_uri.port
        if port is None:
            return host
        return f"{host}:{port}"

    def _receive_events(
        self,
        sock,
        connection: H2Connection,
        stream_id: int,
        response_body: bytearray,
        response_status,
        response_reason: str,
    ):
        data = sock.recv(65535)
        if not data:
            raise Http2TransportError(f"HTTP/2 connection to {self.uri} closed before the response completed.")

        stream_ended = False
        for event in connection.receive_data(data):
            if isinstance(event, ResponseReceived) and event.stream_id == stream_id:
                for key, value in event.headers:
                    if key == ":status":
                        response_status = int(value)
                    elif key == "x-http2-reason":
                        response_reason = value
            elif isinstance(event, DataReceived) and event.stream_id == stream_id:
                response_body.extend(event.data)
                connection.acknowledge_received_data(event.flow_controlled_length, event.stream_id)
            elif isinstance(event, StreamEnded) and event.stream_id == stream_id:
                stream_ended = True
            elif isinstance(event, StreamReset) and event.stream_id == stream_id:
                raise Http2TransportError(f"HTTP/2 stream was reset by the remote endpoint. Error code: {event.error_code}")
            elif isinstance(event, ConnectionTerminated):
                raise Http2TransportError(
                    f"HTTP/2 connection was terminated by the remote endpoint. Error code: {event.error_code}"
                )
            elif isinstance(event, WindowUpdated):
                pass

        self._send_pending_data(sock, connection)
        return stream_ended, response_status, response_reason

    @staticmethod
    def _send_pending_data(sock, connection: H2Connection) -> None:
        pending = connection.data_to_send()
        if pending:
            sock.sendall(pending)

    def _ensure_open(self) -> None:
        if self._disposed:
            raise Http2ClientDisposedError(
                f"HTTP/2 client is disposed for {self.uri}. "
                "The connection may have been closed and should be recreated."
            )

    def close(self) -> None:
        sock = self._socket
        connection = self._connection
        if sock is not None and connection is not None:
            try:
                connection.close_connection()
                self._send_pending_data(sock, connection)
            except Exception:
                pass
        if sock is not None:
            try:
                sock.close()
            except Exception:
                pass
        self._socket = None
        self._connection = None
        self._disposed = True

    @staticmethod
    def close_client(connection_data) -> None:
        key = Http2Client._connection_key(connection_data)
        with _clients_lock:
            client = _clients.pop(key, None)
        if client is not None:
            client.close()

    @staticmethod
    def get_state(connection_data):
        key = Http2Client._connection_key(connection_data)
        with _clients_lock:
            if key not in _clients:
                return None
            return True
