"""
BIBIVIDI — HTTPS CONNECT tunnel
Listens on a separate TCP port, handles HTTP CONNECT requests:
  - AI API hosts  → TLS MITM → inner request forwarded to HTTP proxy (port 8080)
  - Other hosts   → transparent TCP pipe

Architecture:
  Client ──CONNECT api.anthropic.com:443──► Tunnel (port 8443)
             │  200 Connection Established
             │  TLS handshake (per-domain cert signed by BIBIVIDI CA)
             └──► inner HTTP request ──► HTTP proxy (port 8080) ──► real API
"""
from __future__ import annotations

import asyncio
import logging
import os
import ssl
import tempfile

import httpx

from bvd.config import settings
from bvd.core.certs import cert_pem_for_domain, get_or_generate_certs

logger = logging.getLogger(__name__)

# ssl.SSLContext cache per domain
_ssl_ctx_cache: dict[str, ssl.SSLContext] = {}


def _ssl_ctx_for_host(hostname: str) -> ssl.SSLContext:
    """Build (or return cached) server-side SSLContext for a given hostname."""
    if hostname not in _ssl_ctx_cache:
        cert_pem, key_pem = cert_pem_for_domain(hostname)
        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)

        # SSLContext.load_cert_chain() needs file paths — use a temp combined PEM file
        fd, tmp_path = tempfile.mkstemp(suffix=".pem")
        try:
            os.write(fd, cert_pem + key_pem)
            os.close(fd)
            ctx.load_cert_chain(tmp_path)
        finally:
            os.unlink(tmp_path)

        _ssl_ctx_cache[hostname] = ctx
    return _ssl_ctx_cache[hostname]


# ── Connection handler ────────────────────────────────────────────────────────

async def handle_client(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
) -> None:
    try:
        first_line = await asyncio.wait_for(reader.readline(), timeout=10.0)
        if not first_line:
            return

        if first_line.upper().startswith(b"CONNECT "):
            await _handle_connect(reader, writer, first_line)
        else:
            # Plain HTTP on the tunnel port — forward to the HTTP proxy
            rest = await reader.read(65_536)
            await _forward_bytes_to_proxy(writer, first_line + rest)

    except (TimeoutError, ConnectionResetError, BrokenPipeError):
        pass
    except Exception as exc:
        logger.debug("Tunnel client error: %s", exc)
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass


# ── CONNECT handling ──────────────────────────────────────────────────────────

async def _handle_connect(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    first_line: bytes,
) -> None:
    parts = first_line.decode(errors="replace").strip().split()
    if len(parts) < 2:
        writer.write(b"HTTP/1.1 400 Bad Request\r\n\r\n")
        return

    host_port = parts[1]
    hostname, _, port_str = host_port.partition(":")
    port = int(port_str) if port_str.isdigit() else 443

    # Drain remaining CONNECT headers
    while True:
        line = await asyncio.wait_for(reader.readline(), timeout=5.0)
        if line in (b"\r\n", b"\n", b""):
            break

    # 200 before TLS handshake (regardless of MITM or transparent)
    writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
    await writer.drain()

    is_ai_target = any(t in hostname for t in settings.target_apis)

    if is_ai_target:
        await _mitm_tls(reader, writer, hostname)
    else:
        await _transparent_pipe(reader, writer, hostname, port)


async def _mitm_tls(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    hostname: str,
) -> None:
    """Upgrade the connection to TLS (server side) and process inner requests."""
    ssl_ctx = _ssl_ctx_for_host(hostname)
    try:
        # Python 3.11+: start_tls upgrades both read and write sides
        await writer.start_tls(ssl_ctx)
    except ssl.SSLError as exc:
        logger.debug("TLS handshake failed for %s: %s", hostname, exc)
        return

    await _process_inner_http(reader, writer, hostname)


async def _process_inner_http(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    hostname: str,
) -> None:
    """Read inner HTTP/1.1 requests and forward them to the local HTTP proxy."""
    proxy_url = f"http://{settings.proxy_host}:{settings.proxy_port}"

    async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as client:
        while True:
            try:
                # Read headers (until \r\n\r\n)
                header_buf = b""
                while b"\r\n\r\n" not in header_buf:
                    chunk = await asyncio.wait_for(reader.read(4096), timeout=30.0)
                    if not chunk:
                        return
                    header_buf += chunk

                header_section, _, rest = header_buf.partition(b"\r\n\r\n")
                header_lines = header_section.split(b"\r\n")
                if not header_lines:
                    return

                # Request line
                req_line = header_lines[0].decode(errors="replace")
                parts = req_line.split()
                if len(parts) < 2:
                    return
                method, path = parts[0], parts[1]

                # Headers
                headers: dict[str, str] = {"host": hostname}
                content_length = 0
                for hl in header_lines[1:]:
                    if b":" in hl:
                        k, _, v = hl.partition(b":")
                        key = k.decode().strip().lower()
                        val = v.decode().strip()
                        headers[key] = val
                        if key == "content-length":
                            content_length = int(val)

                # Body
                body = rest
                remaining = content_length - len(body)
                if remaining > 0:
                    body += await asyncio.wait_for(
                        reader.readexactly(remaining), timeout=30.0
                    )

                # Forward to local HTTP proxy
                target_url = f"{proxy_url}{path}" if path.startswith("/") else f"{proxy_url}/{path}"
                resp = await client.request(
                    method=method,
                    url=target_url,
                    headers=headers,
                    content=body,
                )

                # Write HTTP/1.1 response back to the client
                reason = resp.reason_phrase or "OK"
                resp_headers = (
                    f"HTTP/1.1 {resp.status_code} {reason}\r\n"
                    + "".join(
                        f"{k}: {v}\r\n"
                        for k, v in resp.headers.items()
                        if k.lower() not in {"transfer-encoding", "connection"}
                    )
                    + f"content-length: {len(resp.content)}\r\n"
                    + "\r\n"
                )
                writer.write(resp_headers.encode() + resp.content)
                await writer.drain()

                # HTTP/1.1 keep-alive: loop for next request if client sends one
                if headers.get("connection", "").lower() == "close":
                    break

            except (TimeoutError, asyncio.IncompleteReadError,
                    ConnectionResetError, BrokenPipeError):
                break
            except Exception as exc:
                logger.warning("Inner HTTPS request error (%s): %s", hostname, exc)
                try:
                    writer.write(b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\n\r\n")
                    await writer.drain()
                except Exception:
                    pass
                break


# ── Transparent pipe (non-AI hosts) ──────────────────────────────────────────

async def _transparent_pipe(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    hostname: str,
    port: int,
) -> None:
    """Pipe TCP bytes transparently to the real host without inspection."""
    try:
        rem_reader, rem_writer = await asyncio.wait_for(
            asyncio.open_connection(hostname, port),
            timeout=10.0,
        )
        await asyncio.gather(
            _pipe(reader, rem_writer),
            _pipe(rem_reader, writer),
            return_exceptions=True,
        )
    except Exception as exc:
        logger.debug("Transparent pipe error (%s:%d): %s", hostname, port, exc)
    finally:
        try:
            rem_writer.close()  # type: ignore[possibly-undefined]
        except Exception:
            pass


async def _pipe(src: asyncio.StreamReader, dst: asyncio.StreamWriter) -> None:
    while True:
        data = await src.read(65_536)
        if not data:
            break
        dst.write(data)
        await dst.drain()


async def _forward_bytes_to_proxy(writer: asyncio.StreamWriter, data: bytes) -> None:
    """Forward a plain HTTP request to the local HTTP proxy and return the response."""
    try:
        r, w = await asyncio.open_connection(settings.proxy_host, settings.proxy_port)
        w.write(data)
        await w.drain()
        response = await r.read(65_536)
        writer.write(response)
        await writer.drain()
        w.close()
    except Exception as exc:
        logger.debug("Forward to HTTP proxy error: %s", exc)


# ── Server lifecycle ──────────────────────────────────────────────────────────

async def start_tunnel_server(host: str, port: int) -> asyncio.AbstractServer:
    """Start the CONNECT tunnel server and return the server object."""
    get_or_generate_certs()
    server = await asyncio.start_server(handle_client, host, port)
    logger.info("BIBIVIDI CONNECT tunnel on %s:%d", host, port)
    return server
