from __future__ import annotations

import asyncio
import hmac
import hashlib
import os
import random
import ssl
import struct
import time
import logging

from typing import Dict, FrozenSet, NamedTuple, Optional, Tuple
from .stats import stats

try:
    from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
except ImportError:
    X25519PrivateKey = None


log = logging.getLogger('tg-mtproto-proxy')

TLS_RECORD_HANDSHAKE = 0x16
TLS_RECORD_CCS = 0x14
TLS_RECORD_APPDATA = 0x17

TLS_VERSION_10 = b'\x03\x01'
TLS_VERSION_12 = b'\x03\x03'
TLS_VERSION_13 = b'\x03\x04'

CLIENT_RANDOM_OFFSET = 11
CLIENT_RANDOM_LEN = 32
SESSION_ID_OFFSET = 44
SESSION_ID_LEN = 32

TIMESTAMP_TOLERANCE = 120

TLS_APPDATA_MAX = 16384

EXT_SUPPORTED_VERSIONS = 0x002b
EXT_KEY_SHARE = 0x0033

GROUP_X25519 = 0x001d
GROUP_X25519MLKEM768 = 0x11ec

X25519_KEY_LEN = 32
X25519MLKEM768_CLIENT_SHARE_LEN = 1216
MLKEM768_CIPHERTEXT_LEN = 1088

# TLS 1.3 suites in the order a typical server picks them from the client list
CIPHER_PREFERENCE = (0x1301, 0x1303, 0x1302)

# Bounds for the single ApplicationData record that follows CCS; the low bound
# matches the smallest record telemt emits, the high one is the TLS limit
APPDATA_MIN = 64
# Without a site profile stay near the 64-byte record telemt passed DPI with,
# far from the 1900-2100 bytes of the classic MTProxy answer
FALLBACK_APPDATA_RANGE = (64, 160)

SERVER_HELLO_DELAY = (0.008, 0.024)

PROFILE_PROBE_TIMEOUT = 5


_CCS_FRAME = b'\x14\x03\x03\x00\x01\x01'

_SH_RANDOM_OFF = 11

# Size of the first encrypted record the masking site sends after ServerHello;
# set once at startup by probe_appdata_len()
_profiled_appdata_len: Optional[int] = None


class ClientHelloInfo(NamedTuple):
    cipher_suites: FrozenSet[int]
    key_shares: Dict[int, int]  # named group -> key_exchange length


_EMPTY_HELLO_INFO = ClientHelloInfo(frozenset(), {})


def verify_client_hello(data: bytes, secret: bytes) -> Optional[Tuple[bytes, bytes, int]]:
    n = len(data)
    # 5 (record hdr) + 6 (hs type+len+version) + 32 (random) = 43
    if n < 43:
        return None
    if data[0] != TLS_RECORD_HANDSHAKE:
        return None
    if data[5] != 0x01:
        return None

    client_random = bytes(data[CLIENT_RANDOM_OFFSET:CLIENT_RANDOM_OFFSET + CLIENT_RANDOM_LEN])

    zeroed = bytearray(data)
    zeroed[CLIENT_RANDOM_OFFSET:CLIENT_RANDOM_OFFSET + CLIENT_RANDOM_LEN] = b'\x00' * CLIENT_RANDOM_LEN

    expected = hmac.new(secret, bytes(zeroed), hashlib.sha256).digest()

    if not hmac.compare_digest(expected[:28], client_random[:28]):
        return None

    ts_xor = bytes(client_random[28 + i] ^ expected[28 + i] for i in range(4))
    timestamp = struct.unpack('<I', ts_xor)[0]

    now = int(time.time())
    if abs(now - timestamp) > TIMESTAMP_TOLERANCE:
        return None

    session_id = b'\x00' * SESSION_ID_LEN
    if n >= SESSION_ID_OFFSET + SESSION_ID_LEN and data[43] == 0x20:
        session_id = bytes(data[SESSION_ID_OFFSET:SESSION_ID_OFFSET + SESSION_ID_LEN])

    return client_random, session_id, timestamp


def parse_client_hello(data: bytes) -> ClientHelloInfo:
    """Extract offered cipher suites and key_share groups from a ClientHello.

    Returns empty info on malformed input so the caller falls back to the
    plain X25519 / TLS_AES_128_GCM_SHA256 answer.
    """
    try:
        pos = 5 + 4 + 2 + CLIENT_RANDOM_LEN
        pos += 1 + data[pos]

        suites_len = struct.unpack_from('>H', data, pos)[0]
        pos += 2
        suites_end = pos + suites_len
        if suites_end > len(data):
            return _EMPTY_HELLO_INFO
        suites = frozenset(
            struct.unpack_from('>H', data, off)[0]
            for off in range(pos, suites_end - 1, 2))
        pos = suites_end

        pos += 1 + data[pos]

        ext_total = struct.unpack_from('>H', data, pos)[0]
        pos += 2
        ext_end = min(pos + ext_total, len(data))
        key_shares: Dict[int, int] = {}
        while pos + 4 <= ext_end:
            ext_type, ext_len = struct.unpack_from('>HH', data, pos)
            pos += 4
            if ext_type == EXT_KEY_SHARE and pos + 2 <= ext_end:
                shares_end = min(pos + 2 + struct.unpack_from('>H', data, pos)[0],
                                 ext_end)
                off = pos + 2
                while off + 4 <= shares_end:
                    group, key_len = struct.unpack_from('>HH', data, off)
                    off += 4 + key_len
                    if off <= shares_end:
                        key_shares.setdefault(group, key_len)
            pos += ext_len
    except (IndexError, struct.error):
        return _EMPTY_HELLO_INFO

    return ClientHelloInfo(suites, key_shares)


def _select_cipher_suite(info: ClientHelloInfo) -> int:
    for suite in CIPHER_PREFERENCE:
        if suite in info.cipher_suites:
            return suite
    return CIPHER_PREFERENCE[0]


def _x25519_public_key() -> bytes:
    if X25519PrivateKey is not None:
        return X25519PrivateKey.generate().public_key().public_bytes_raw()
    # Real public keys are field elements below 2^255; keep the top bit clear
    key = bytearray(os.urandom(X25519_KEY_LEN))
    key[-1] &= 0x7f
    return bytes(key)


def _server_key_share(info: ClientHelloInfo) -> Tuple[int, bytes]:
    # Mirror the hybrid post-quantum group when the client offers it, as
    # current browsers and servers do; the client never derives keys from it
    if info.key_shares.get(GROUP_X25519MLKEM768) == X25519MLKEM768_CLIENT_SHARE_LEN:
        return GROUP_X25519MLKEM768, (os.urandom(MLKEM768_CIPHERTEXT_LEN)
                                      + _x25519_public_key())
    return GROUP_X25519, _x25519_public_key()


def _appdata_len() -> int:
    if _profiled_appdata_len is not None:
        return _profiled_appdata_len
    return random.randint(*FALLBACK_APPDATA_RANGE)


def build_server_hello(secret: bytes, client_random: bytes, session_id: bytes,
                       info: ClientHelloInfo = _EMPTY_HELLO_INFO) -> bytes:
    """ServerHello + CCS + one ApplicationData record.

    Telegram clients parse exactly these three records and check the HMAC
    over all of them, so the shape must stay fixed while sizes may vary.
    """
    group, key_exchange = _server_key_share(info)
    extensions = (
        struct.pack('>HHH', EXT_SUPPORTED_VERSIONS, 2, 0x0304)
        + struct.pack('>HHHH', EXT_KEY_SHARE, 4 + len(key_exchange),
                      group, len(key_exchange))
        + key_exchange
    )

    body = (
        TLS_VERSION_12
        + b'\x00' * 32
        + bytes([len(session_id)]) + session_id
        + struct.pack('>H', _select_cipher_suite(info))
        + b'\x00'
        + struct.pack('>H', len(extensions)) + extensions
    )
    message = b'\x02' + struct.pack('>I', len(body))[1:] + body
    server_hello = (bytes([TLS_RECORD_HANDSHAKE]) + TLS_VERSION_12
                    + struct.pack('>H', len(message)) + message)

    appdata_len = _appdata_len()
    app_record = (bytes([TLS_RECORD_APPDATA]) + TLS_VERSION_12
                  + struct.pack('>H', appdata_len) + os.urandom(appdata_len))

    response = server_hello + _CCS_FRAME + app_record

    hmac_input = client_random + response
    server_random = hmac.new(secret, hmac_input, hashlib.sha256).digest()

    final = bytearray(response)
    final[_SH_RANDOM_OFF:_SH_RANDOM_OFF + 32] = server_random

    return bytes(final)


def set_profiled_appdata_len(length: Optional[int]) -> None:
    global _profiled_appdata_len
    if length is None:
        _profiled_appdata_len = None
    else:
        _profiled_appdata_len = max(APPDATA_MIN, min(length, TLS_APPDATA_MAX))


def first_appdata_record_len(data: bytes) -> Optional[int]:
    """Length of the first ApplicationData record in a server's first flight.

    Returns None while more data is needed, 0 if the flight cannot contain
    one (alert, TLS 1.2 handshake, garbage).
    """
    pos = 0
    while pos + 5 <= len(data):
        rtype = data[pos]
        rec_len = struct.unpack_from('>H', data, pos + 3)[0]
        if rtype == TLS_RECORD_APPDATA:
            return rec_len
        if rtype not in (TLS_RECORD_HANDSHAKE, TLS_RECORD_CCS):
            return 0
        pos += 5 + rec_len
    return None


async def probe_appdata_len(domain: str, port: int = 443,
                            timeout: float = PROFILE_PROBE_TIMEOUT) -> Optional[int]:
    """Measure the first encrypted record the masking site sends after ServerHello.

    Only the raw first flight is read; the TLS handshake is never completed.
    """
    if not getattr(ssl, 'HAS_TLSv1_3', False):
        log.warning("Fake TLS profile: local ssl lacks TLS 1.3, "
                    "using random record sizes")
        return None

    writer = None
    try:
        ctx = ssl.create_default_context()
        ctx.set_ecdh_curve('X25519')
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
        ctx.minimum_version = ssl.TLSVersion.TLSv1_3
        ctx.set_alpn_protocols(['h2', 'http/1.1'])
        outgoing = ssl.MemoryBIO()
        tls = ctx.wrap_bio(ssl.MemoryBIO(), outgoing, server_hostname=domain)
        try:
            tls.do_handshake()
        except ssl.SSLWantReadError:
            pass
        client_hello = outgoing.read()

        # Share one deadline across connecting, sending, and reading the flight.
        loop = asyncio.get_running_loop()
        deadline = loop.time() + timeout
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(domain, port), timeout=timeout)
        writer.write(client_hello)
        await asyncio.wait_for(writer.drain(), timeout=deadline - loop.time())

        flight = bytearray()
        while True:
            remaining = deadline - loop.time()
            if remaining <= 0:
                raise asyncio.TimeoutError
            chunk = await asyncio.wait_for(
                reader.read(16384), timeout=remaining)
            if not chunk:
                return None
            flight.extend(chunk)
            length = first_appdata_record_len(bytes(flight))
            if length is not None:
                return length or None
    except Exception as exc:
        log.warning("Fake TLS profile: cannot probe %s:%d: %r",
                    domain, port, exc)
        return None
    finally:
        if writer is not None:
            writer.close()


def wrap_tls_record(data: bytes) -> bytes:
    parts = []
    offset = 0
    while offset < len(data):
        chunk = data[offset:offset + TLS_APPDATA_MAX]
        parts.append(
            b'\x17\x03\x03'
            + struct.pack('>H', len(chunk))
            + chunk
        )
        offset += len(chunk)
    return b''.join(parts)


class FakeTlsStream:
    __slots__ = ('_reader', '_writer', '_read_buf', '_read_left')

    def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        self._reader = reader
        self._writer = writer
        self._read_buf = bytearray()
        self._read_left = 0

    async def readexactly(self, n: int) -> bytes:
        while len(self._read_buf) < n:
            payload = await self._read_tls_payload()
            if not payload:
                raise asyncio.IncompleteReadError(bytes(self._read_buf), n)
            self._read_buf.extend(payload)
        result = bytes(self._read_buf[:n])
        del self._read_buf[:n]
        return result

    async def read(self, n: int) -> bytes:
        if self._read_buf:
            chunk = bytes(self._read_buf[:n])
            del self._read_buf[:n]
            return chunk
        payload = await self._read_tls_payload()
        if not payload:
            return b''
        if len(payload) > n:
            self._read_buf.extend(payload[n:])
            return payload[:n]
        return payload

    async def _read_tls_payload(self) -> bytes:
        if self._read_left > 0:
            data = await self._reader.read(self._read_left)
            if not data:
                return b''
            self._read_left -= len(data)
            return data

        while True:
            hdr = await self._reader.readexactly(5)
            rtype = hdr[0]
            rec_len = struct.unpack('>H', hdr[3:5])[0]

            if rtype == TLS_RECORD_CCS:
                if rec_len > 0:
                    await self._reader.readexactly(rec_len)
                continue

            if rtype != TLS_RECORD_APPDATA:
                return b''

            data = await self._reader.read(min(rec_len, 65536))
            if not data:
                return b''
            remaining = rec_len - len(data)
            if remaining > 0:
                self._read_left = remaining
            return data

    def write(self, data: bytes) -> None:
        self._writer.write(wrap_tls_record(data))

    async def drain(self) -> None:
        await self._writer.drain()

    def close(self) -> None:
        self._writer.close()

    async def wait_closed(self) -> None:
        await self._writer.wait_closed()

    def get_extra_info(self, name, default=None):
        return self._writer.get_extra_info(name, default)

    @property
    def transport(self):
        return self._writer.transport

    def is_closing(self):
        return self._writer.is_closing()


async def proxy_to_masking_domain(reader, writer, initial_data: bytes,
                                    domain: str, label: str) -> None:
    try:
        up_reader, up_writer = await asyncio.wait_for(
            asyncio.open_connection(domain, 443), timeout=10)
    except Exception as exc:
        log.warning("[%s] masking: cannot connect to %s:443: %s",
                  label, domain, repr(exc))
        return

    log.debug("[%s] masking -> %s:443", label, domain)
    stats.connections_masked += 1

    try:
        if initial_data:
            up_writer.write(initial_data)
            await up_writer.drain()

        async def _relay(src, dst):
            try:
                while True:
                    chunk = await src.read(16384)
                    if not chunk:
                        break
                    dst.write(chunk)
                    await dst.drain()
            except (ConnectionResetError, BrokenPipeError, OSError,
                    asyncio.CancelledError):
                pass
            finally:
                try:
                    dst.close()
                    await dst.wait_closed()
                except Exception:
                    pass

        await asyncio.gather(
            _relay(reader, up_writer),
            _relay(up_reader, writer),
        )
    except Exception:
        pass
    finally:
        try:
            up_writer.close()
        except Exception:
            pass