"""Plaintext-to-TLS HTTP/1.1 forwarding proxy that counts requests and bytes.

Point a client at ``http://127.0.0.1:PORT`` and every request is forwarded to

``https://UPSTREAM`` over TLS, with the ``Host`` header rewritten. The proxy

records one entry per HTTP request: method, path (query string dropped so no

token-bearing query parameters are ever logged), status, request/response

bytes (HTTP/1.1 framing: request line + headers + body), and timing.

Request/response headers are never recorded. An optional ``inject_headers``

map lets a caller add headers (e.g. auth for a proxied Git remote) without

putting credentials on a command line; injected values are never logged.

Counts are HTTP/1.1 plaintext bytes. A direct HTTPS client may negotiate HTTP/2

(HPACK-compressed headers, multiplexing), so proxied *timing* is diagnostic;

request counts and body bytes are the instrument's product.

Stdlib only (ADR-0001).

"""

from __future__ import annotations

import socket

import ssl

import threading

import time

from dataclasses import dataclass, field, asdict

from typing import Optional

@dataclass

class RequestRecord:

    method: str

    path: str

    status: Optional[int]

    request_bytes: int

    response_bytes: int

    start_ms: float

    ttfb_ms: Optional[float]

    end_ms: Optional[float]

    connection_id: int

    error: Optional[str] = None

    def to_dict(self) -> dict:

        return asdict(self)

@dataclass

class TunnelRecord:

    """A CONNECT tunnel (opaque TLS): bytes and host only, no request count."""

    host: str

    port: int

    client_to_upstream_bytes: int

    upstream_to_client_bytes: int

    start_ms: float

    end_ms: Optional[float]

    def to_dict(self) -> dict:

        return asdict(self)

@dataclass

class HttpCountStats:

    requests: int

    connections: int

    request_bytes: int

    response_bytes: int

    records: list = field(default_factory=list)

    tunnels: list = field(default_factory=list)

    @property

    def tunnel_bytes(self) -> int:

        return sum(t.client_to_upstream_bytes + t.upstream_to_client_bytes for t in self.tunnels)

    def tunnel_hosts(self) -> dict:

        out: dict = {}

        for t in self.tunnels:

            h = out.setdefault(t.host, {"connections": 0, "up_bytes": 0, "down_bytes": 0})

            h["connections"] += 1

            h["up_bytes"] += t.client_to_upstream_bytes

            h["down_bytes"] += t.upstream_to_client_bytes

        return out

    @property

    def total_bytes(self) -> int:

        return self.request_bytes + self.response_bytes

    def by_route(self) -> dict:

        """Histogram of ``METHOD /path-shape`` -> count (hex ids collapsed)."""

        out: dict = {}

        for r in self.records:

            key = f"{r.method} {route_shape(r.path)}"

            out[key] = out.get(key, 0) + 1

        return dict(sorted(out.items(), key=lambda kv: (-kv[1], kv[0])))

def route_shape(path: str) -> str:

    """Collapse long hex/numeric path segments so routes aggregate."""

    parts = []

    for seg in path.split("/"):

        if len(seg) >= 12 and all(c in "0123456789abcdefABCDEF" for c in seg):

            parts.append(":hash")

        elif seg.isdigit():

            parts.append(":n")

        else:

            parts.append(seg)

    return "/".join(parts)

class _Reader:

    """Buffered reader over a socket that counts consumed bytes."""

    def __init__(self, sock: socket.socket) -> None:

        self.sock = sock

        self.buf = b""

    def _fill(self) -> bool:

        chunk = self.sock.recv(256 * 1024)

        if not chunk:

            return False

        self.buf += chunk

        return True

    def readline(self) -> bytes:

        while b"\r\n" not in self.buf:

            if not self._fill():

                out, self.buf = self.buf, b""

                return out

        idx = self.buf.index(b"\r\n") + 2

        out, self.buf = self.buf[:idx], self.buf[idx:]

        return out

    def read_some(self, limit: int) -> bytes:

        if not self.buf:

            if not self._fill():

                return b""

        out, self.buf = self.buf[:limit], self.buf[limit:]

        return out

    def read_head(self) -> Optional[list]:

        lines = []

        while True:

            line = self.readline()

            if not line:

                return None if not lines else lines

            if line == b"\r\n":

                return lines

            lines.append(line)

def _header(lines: list, name: bytes) -> Optional[bytes]:

    lname = name.lower()

    for line in lines[1:]:

        k, _, v = line.partition(b":")

        if k.strip().lower() == lname:

            return v.strip()

    return None

class HttpCountingProxy:

    def __init__(

        self,

        upstream_host: str,

        *,

        upstream_port: int = 443,

        listen_host: str = "127.0.0.1",

        listen_port: int = 0,

        inject_headers: Optional[dict] = None,

        use_tls: bool = True,

    ) -> None:

        self.upstream_host = upstream_host

        self.upstream_port = upstream_port

        self.listen_host = listen_host

        self.listen_port = listen_port

        self.inject_headers = dict(inject_headers or {})

        self.use_tls = use_tls

        self._ctx = ssl.create_default_context()

        self._ctx.set_alpn_protocols(["http/1.1"])

        self._lock = threading.Lock()

        self._stop = threading.Event()

        self._listener: Optional[socket.socket] = None

        self._threads: list = []

        self._socks: set = set()

        self._records: list = []

        self._tunnels: list = []

        self._connections = 0

        self._t0 = time.monotonic()

    # -- lifecycle -----------------------------------------------------

    @property

    def address(self) -> tuple:

        assert self._listener is not None

        h, p = self._listener.getsockname()[:2]

        return str(h), int(p)

    @property

    def url(self) -> str:

        h, p = self.address

        return f"http://{h}:{p}"

    def start(self) -> "HttpCountingProxy":

        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        s.bind((self.listen_host, self.listen_port))

        s.listen(64)

        s.settimeout(0.2)

        self._listener = s

        t = threading.Thread(target=self._accept, daemon=True)

        t.start()

        self._threads.append(t)

        return self

    def stop(self) -> None:

        self._stop.set()

        if self._listener is not None:

            try:

                self._listener.close()

            except OSError:

                pass

        with self._lock:

            socks = list(self._socks)

        for s in socks:

            _close(s)

    def __enter__(self):

        return self.start()

    def __exit__(self, *exc):

        self.stop()

    def reset(self) -> None:

        with self._lock:

            self._records = []

            self._tunnels = []

            self._connections = 0

            self._t0 = time.monotonic()

    def snapshot(self) -> HttpCountStats:

        with self._lock:

            recs = list(self._records)

            tuns = list(self._tunnels)

            conns = self._connections

        return HttpCountStats(

            requests=len(recs),

            connections=conns,

            request_bytes=sum(r.request_bytes for r in recs),

            response_bytes=sum(r.response_bytes for r in recs),

            records=recs,

            tunnels=tuns,

        )

    # -- internals -----------------------------------------------------

    def _now_ms(self) -> float:

        return (time.monotonic() - self._t0) * 1000.0

    def _accept(self) -> None:

        assert self._listener is not None

        while not self._stop.is_set():

            try:

                c, _ = self._listener.accept()

            except socket.timeout:

                continue

            except OSError:

                break

            t = threading.Thread(target=self._serve, args=(c,), daemon=True)

            with self._lock:

                self._socks.add(c)

                self._threads.append(t)

            t.start()

    def _connect_upstream(self) -> socket.socket:

        raw = socket.create_connection((self.upstream_host, self.upstream_port), timeout=60)

        raw.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

        if not self.use_tls:

            return raw

        return self._ctx.wrap_socket(raw, server_hostname=self.upstream_host)

    def _serve(self, client: socket.socket) -> None:

        client.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

        up: Optional[socket.socket] = None

        with self._lock:

            self._connections += 1

            conn_id = self._connections

        cr = _Reader(client)

        ur: Optional[_Reader] = None

        try:

            while not self._stop.is_set():

                head = cr.read_head()

                if not head:

                    break

                if head[0].upper().startswith(b"CONNECT "):

                    self._tunnel(head, cr, client)

                    return

                if up is None:

                    up = self._connect_upstream()

                    with self._lock:

                        self._socks.add(up)

                    ur = _Reader(up)

                assert ur is not None

                rec = self._one_exchange(head, cr, client, up, ur, conn_id)

                with self._lock:

                    self._records.append(rec)

                if rec.error or rec.status is None:

                    break

        except OSError:

            pass

        finally:

            _close(client)

            if up is not None:

                _close(up)

            with self._lock:

                self._socks.discard(client)

                if up is not None:

                    self._socks.discard(up)

    def _tunnel(self, head, cr: "_Reader", client: socket.socket) -> None:

        target = head[0].split(b" ")[1].decode("latin-1")

        host, _, port_s = target.rpartition(":")

        port = int(port_s or 443)

        rec = TunnelRecord(host, port, 0, 0, self._now_ms(), None)

        try:

            up = socket.create_connection((host, port), timeout=60)

        except OSError:

            client.sendall(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")

            return

        with self._lock:

            self._socks.add(up)

            self._tunnels.append(rec)

        client.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")

        if cr.buf:

            up.sendall(cr.buf)

            rec.client_to_upstream_bytes += len(cr.buf)

            cr.buf = b""

        def pump(src, dst, attr):

            try:

                while True:

                    b = src.recv(256 * 1024)

                    if not b:

                        break

                    with self._lock:

                        setattr(rec, attr, getattr(rec, attr) + len(b))

                    dst.sendall(b)

            except OSError:

                pass

            finally:

                try:

                    dst.shutdown(socket.SHUT_WR)

                except OSError:

                    pass

        t = threading.Thread(target=pump, args=(up, client, "upstream_to_client_bytes"), daemon=True)

        t.start()

        pump(client, up, "client_to_upstream_bytes")

        t.join()

        rec.end_ms = self._now_ms()

        _close(up)

        with self._lock:

            self._socks.discard(up)

    def _one_exchange(self, head, cr, client, up, ur, conn_id) -> RequestRecord:

        start = self._now_ms()

        reqline = head[0].decode("latin-1").rstrip("\r\n")

        parts = reqline.split(" ")

        method = parts[0] if parts else "?"

        target = parts[1] if len(parts) > 1 else "/"

        path = target.split("?", 1)[0]

        out = [head[0]]

        injected = {k.lower() for k in self.inject_headers}

        for line in head[1:]:

            k = line.split(b":", 1)[0].strip().lower()

            if k == b"host":

                port = "" if self.upstream_port in (80, 443) else f":{self.upstream_port}"

                out.append(f"Host: {self.upstream_host}{port}\r\n".encode())

            elif k.decode("latin-1") in injected:

                continue

            else:

                out.append(line)

        for k, v in self.inject_headers.items():

            out.append(f"{k}: {v}\r\n".encode("latin-1"))

        out.append(b"\r\n")

        blob = b"".join(out)

        up.sendall(blob)

        req_bytes = len(blob)

        req_bytes += _relay_body(head, cr, up, is_response=False, method=method)

        # Response (skip 1xx informational heads).

        resp_bytes = 0

        ttfb = None

        status = None

        while True:

            rhead = ur.read_head()

            if not rhead:

                return RequestRecord(method, path, None, req_bytes, resp_bytes, start, ttfb, self._now_ms(), conn_id, "upstream closed")

            if ttfb is None:

                ttfb = self._now_ms() - start

            try:

                status = int(rhead[0].split(b" ")[1])

            except (IndexError, ValueError):

                status = None

            rblob = b"".join(rhead) + b"\r\n"

            client.sendall(rblob)

            resp_bytes += len(rblob)

            if status is not None and 100 <= status < 200 and status != 101:

                continue

            break

        resp_bytes += _relay_body(rhead, ur, client, is_response=True, method=method, status=status)

        return RequestRecord(method, path, status, req_bytes, resp_bytes, start, ttfb, self._now_ms(), conn_id)

def _relay_body(head, reader: _Reader, dest: socket.socket, *, is_response: bool, method: str, status: Optional[int] = None) -> int:

    if is_response and (method == "HEAD" or status in (204, 304) or (status is not None and 100 <= status < 200)):

        return 0

    te = _header(head, b"transfer-encoding")

    cl = _header(head, b"content-length")

    n = 0

    if te is not None and b"chunked" in te.lower():

        while True:

            size_line = reader.readline()

            if not size_line:

                return n

            dest.sendall(size_line)

            n += len(size_line)

            size = int(size_line.split(b";")[0].strip() or b"0", 16)

            if size == 0:

                # trailers until blank line

                while True:

                    t = reader.readline()

                    dest.sendall(t)

                    n += len(t)

                    if t in (b"\r\n", b""):

                        return n

            remaining = size + 2  # data + CRLF

            while remaining > 0:

                chunk = reader.read_some(remaining)

                if not chunk:

                    return n

                dest.sendall(chunk)

                n += len(chunk)

                remaining -= len(chunk)

    if cl is not None:

        remaining = int(cl)

        while remaining > 0:

            chunk = reader.read_some(remaining)

            if not chunk:

                break

            dest.sendall(chunk)

            n += len(chunk)

            remaining -= len(chunk)

        return n

    if is_response:

        # read until close

        while True:

            chunk = reader.read_some(256 * 1024)

            if not chunk:

                break

            dest.sendall(chunk)

            n += len(chunk)

        try:

            dest.shutdown(socket.SHUT_WR)

        except OSError:

            pass

    return n

def _close(s: socket.socket) -> None:

    try:

        s.shutdown(socket.SHUT_RDWR)

    except OSError:

        pass

    try:

        s.close()

    except OSError:

        pass