"""Fault-injecting HTTP/1.1 reverse proxy for loopback subjects only.

Idea lineage: the 2026-09-26 W3a QA proxy (drop_reply / fail rules); this is a

rewrite with a typed fault plan, per-request records, and client-kill hooks.

It refuses non-loopback upstreams: fault injection never targets oak.space.

Fault kinds (each rule fires once, on the n-th occurrence of an exact

"METHOD /path" key, or on every request for the global kinds):

- fail_status        reply ``status`` without forwarding (typed server error)

- drop_after_send    forward, then close without replying (unknown outcome)

- truncate_reply     forward, send headers + half the body, close (partial body)

- kill_before_forward  run the kill hook (client dies), never forward

- kill_after_upstream  forward, run the kill hook, close (server applied it)

- slow_reply  (global) relay each reply at ``bytes_per_s``

- jitter      (global) random 0..``max_ms`` delay before forwarding (seeded):

              concurrent requests complete out of order (reordering)

"""

from __future__ import annotations

import hashlib

import http.client

import ipaddress

import random

import threading

import time

from dataclasses import dataclass, field

from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from typing import Callable, Optional

PER_REQUEST_KINDS = ("fail_status", "drop_after_send", "truncate_reply", "kill_before_forward", "kill_after_upstream")

GLOBAL_KINDS = ("slow_reply", "jitter")

MUTATING_METHODS = frozenset({"POST", "PUT", "PATCH", "DELETE"})

HOP_HEADERS = frozenset({"host", "transfer-encoding", "content-length", "connection", "keep-alive"})

@dataclass

class Fault:

    kind: str

    key: Optional[str] = None  # "METHOD /path" (query stripped); None for global kinds

    occurrence: int = 1

    status: int = 503

    bytes_per_s: int = 32 * 1024

    max_ms: int = 150

    seed: int = 0

    fired: bool = False

    def describe(self) -> str:

        if self.kind in GLOBAL_KINDS:

            detail = f"bytes_per_s={self.bytes_per_s}" if self.kind == "slow_reply" else f"max_ms={self.max_ms},seed={self.seed}"

            return f"{self.kind}({detail})"

        extra = f",status={self.status}" if self.kind == "fail_status" else ""

        return f"{self.kind}@{self.key}#{self.occurrence}{extra}"

@dataclass

class Record:

    index: int

    method: str

    path: str

    key: str

    occurrence: int

    request_bytes: int

    request_sha256: str

    status: Optional[int] = None

    response_bytes: Optional[int] = None

    tag: str = ""

    t_arrive: float = 0.0

    t_done: float = 0.0

def request_key(method: str, path: str) -> str:

    return f"{method} {path.split('?', 1)[0]}"

def _assert_loopback(host: str) -> None:

    try:

        if ipaddress.ip_address(host).is_loopback:

            return

    except ValueError:

        if host == "localhost":

            return

    raise ValueError(f"fault proxy refuses non-loopback upstream {host!r}")

class FaultProxy:

    def __init__(self, upstream_port: int, faults: Optional[list[Fault]] = None, upstream_host: str = "127.0.0.1",

                 kill_hook: Optional[Callable[[], None]] = None, upstream_timeout_s: float = 120.0,

                 listen_port: int = 0):

        _assert_loopback(upstream_host)

        self.listen_port = listen_port

        self.upstream_host = upstream_host

        self.upstream_port = upstream_port

        self.faults = list(faults or [])

        self.kill_hook = kill_hook

        self.upstream_timeout_s = upstream_timeout_s

        self.records: list[Record] = []

        self._occurrences: dict[str, int] = {}

        self._lock = threading.Lock()

        self._httpd: Optional[ThreadingHTTPServer] = None

        self._thread: Optional[threading.Thread] = None

        self._rng = random.Random(next((f.seed for f in self.faults if f.kind == "jitter"), 0))

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

    def start(self) -> "FaultProxy":

        proxy = self

        class Handler(BaseHTTPRequestHandler):

            protocol_version = "HTTP/1.1"

            def log_message(self, *args):  # noqa: D401 - silence

                pass

            def handle_any(self):

                proxy._handle(self)

            do_GET = do_POST = do_PUT = do_DELETE = do_PATCH = do_HEAD = handle_any

        class Server(ThreadingHTTPServer):

            daemon_threads = True

            def handle_error(self, request, client_address):  # noqa: D401 - injected faults reset sockets

                pass

        httpd = Server(("127.0.0.1", self.listen_port), Handler)

        self._httpd = httpd

        self._thread = threading.Thread(target=httpd.serve_forever, name="p4-faultproxy", daemon=True)

        self._thread.start()

        return self

    def stop(self) -> None:

        if self._httpd is not None:

            self._httpd.shutdown()

            self._httpd.server_close()

            self._httpd = None

    @property

    def port(self) -> int:

        assert self._httpd is not None

        return int(self._httpd.server_address[1])

    @property

    def url(self) -> str:

        return f"http://127.0.0.1:{self.port}"

    def __enter__(self) -> "FaultProxy":

        return self.start()

    def __exit__(self, *exc) -> None:

        self.stop()

    # -- evidence ------------------------------------------------------------

    def fired(self) -> list[Fault]:

        return [f for f in self.faults if f.fired]

    def keys_in_order(self) -> list[tuple[str, int]]:

        return [(r.key, r.occurrence) for r in sorted(self.records, key=lambda r: r.index)]

    def records_after(self, index: int) -> list[Record]:

        return [r for r in self.records if r.index > index]

    # -- core ----------------------------------------------------------------

    def _take(self, key: str, occurrence: int) -> Optional[Fault]:

        with self._lock:

            for fault in self.faults:

                if fault.fired or fault.kind in GLOBAL_KINDS:

                    continue

                if fault.key == key and fault.occurrence == occurrence:

                    fault.fired = True

                    return fault

        return None

    def _global(self, kind: str) -> Optional[Fault]:

        for fault in self.faults:

            if fault.kind == kind:

                fault.fired = True

                return fault

        return None

    @staticmethod

    def _read_body(handler: BaseHTTPRequestHandler) -> bytes:

        length = handler.headers.get("Content-Length")

        if length is not None:

            return handler.rfile.read(int(length))

        if handler.headers.get("Transfer-Encoding", "").lower() == "chunked":

            body = b""

            while True:

                size = int(handler.rfile.readline().strip().split(b";")[0], 16)

                if size == 0:

                    handler.rfile.readline()

                    return body

                body += handler.rfile.read(size)

                handler.rfile.readline()

        return b""

    def _handle(self, handler: BaseHTTPRequestHandler) -> None:

        body = self._read_body(handler)

        key = request_key(handler.command, handler.path)

        with self._lock:

            occurrence = self._occurrences.get(key, 0) + 1

            self._occurrences[key] = occurrence

            record = Record(len(self.records), handler.command, handler.path, key, occurrence, len(body),

                            hashlib.sha256(body).hexdigest()[:16], t_arrive=time.monotonic())

            self.records.append(record)

        fault = self._take(key, occurrence)

        jitter = self._global("jitter")

        if jitter is not None:

            with self._lock:

                delay = self._rng.uniform(0, jitter.max_ms) / 1000.0

            time.sleep(delay)

        try:

            if fault and fault.kind == "fail_status":

                payload = b'{"error":"p4 injected fault"}'

                handler.send_response(fault.status)

                handler.send_header("Content-Type", "application/json")

                handler.send_header("Content-Length", str(len(payload)))

                handler.end_headers()

                handler.wfile.write(payload)

                record.status, record.response_bytes, record.tag = fault.status, len(payload), "INJECTED_STATUS"

                return

            if fault and fault.kind == "kill_before_forward":

                record.tag = "KILLED_BEFORE_FORWARD"

                if self.kill_hook:

                    self.kill_hook()

                self._abort(handler)

                return

            status, headers, data = self._forward(handler, body)

            record.status, record.response_bytes = status, len(data)

            if fault and fault.kind == "drop_after_send":

                record.tag = "REPLY_DROPPED"

                self._abort(handler)

                return

            if fault and fault.kind == "kill_after_upstream":

                record.tag = "KILLED_AFTER_UPSTREAM"

                if self.kill_hook:

                    self.kill_hook()

                self._abort(handler)

                return

            handler.send_response(status)

            for name, value in headers:

                if name.lower() not in HOP_HEADERS:

                    handler.send_header(name, value)

            handler.send_header("Content-Length", str(len(data)))

            handler.end_headers()

            if fault and fault.kind == "truncate_reply":

                record.tag = "REPLY_TRUNCATED"

                handler.wfile.write(data[: len(data) // 2])

                handler.wfile.flush()

                self._abort(handler)

                return

            slow = self._global("slow_reply")

            if slow is not None and data:

                record.tag = "SLOW"

                step = max(1, slow.bytes_per_s // 20)

                for offset in range(0, len(data), step):

                    handler.wfile.write(data[offset: offset + step])

                    handler.wfile.flush()

                    time.sleep(0.05)

            else:

                handler.wfile.write(data)

        except (BrokenPipeError, ConnectionResetError, OSError) as exc:

            record.tag = (record.tag + " " if record.tag else "") + f"CLIENT_GONE:{type(exc).__name__}"

            handler.close_connection = True

        finally:

            record.t_done = time.monotonic()

    def _forward(self, handler: BaseHTTPRequestHandler, body: bytes):

        conn = http.client.HTTPConnection(self.upstream_host, self.upstream_port, timeout=self.upstream_timeout_s)

        try:

            headers = {k: v for k, v in handler.headers.items() if k.lower() not in HOP_HEADERS}

            headers["Content-Length"] = str(len(body))

            headers["Host"] = f"{self.upstream_host}:{self.upstream_port}"

            conn.request(handler.command, handler.path, body=body, headers=headers)

            response = conn.getresponse()

            data = response.read()

            return response.status, response.getheaders(), data

        finally:

            conn.close()

    @staticmethod

    def _abort(handler: BaseHTTPRequestHandler) -> None:

        handler.close_connection = True

        try:

            handler.wfile.flush()

        except OSError:

            pass

        try:

            import socket

            handler.connection.shutdown(socket.SHUT_RDWR)

        except OSError:

            pass