"""Tests for keyproxy.py. Nothing here touches the internet: the "AI provider" is a fake server on 127.0.0.1."""
import io
import json
import logging
import re
import socket
import textwrap
import threading
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from wsgiref.simple_server import WSGIRequestHandler, make_server

import pytest

import keyproxy
from keyproxy import Config, PathFreeHandler, TokenBucket, config_from_env, make_app

KEY = "EXAMPLE-demo-key-0000-not-a-real-key"
GOOD = json.dumps({"model": "demo-small", "messages": [{"role": "user", "content": "hi"}]}).encode()


class QuietServer(ThreadingHTTPServer):
    def handle_error(self, request, client_address):  # the slow-provider test hangs up on purpose
        pass


class FakeProvider:
    """A local stand-in for the AI provider. `mode` picks how it misbehaves; `seen` records what arrived."""

    def __init__(self, mode="ok", redirect_to=None):
        self.mode, self.redirect_to, self.seen = mode, redirect_to, []
        outer = self

        class Handler(BaseHTTPRequestHandler):
            def do_POST(self):
                self.answer(self.rfile.read(int(self.headers["Content-Length"])))

            def do_GET(self):  # urllib turns a redirected POST into a GET
                self.answer(b"null")

            def answer(self, body):
                outer.seen.append({"auth": self.headers.get("Authorization"), "body": json.loads(body)})
                if outer.mode == "slow":
                    threading.Event().wait(0.5)
                echo = self.headers.get("Authorization", "")
                status, payload, headers, message = 200, {"answer": "pretend answer"}, {}, None
                if outer.mode == "echo-key":  # a careless provider that repeats the header it was given
                    payload = {"answer": "pretend answer", "debug": echo}
                if outer.mode == "error-with-key":
                    status, payload = 401, {"error": f"bad credential {echo}"}
                if outer.mode == "key-in-reason":  # the key rides in the status line, which urllib keeps as HTTPError.reason
                    status, message = 401, f"bad credential {echo}"
                if outer.mode == "redirect":
                    status, headers = 302, {"Location": outer.redirect_to}
                data = json.dumps(payload).encode()
                self.send_response(status, message)
                for name, value in headers.items():
                    self.send_header(name, value)
                self.send_header("Content-Length", str(len(data)))
                self.end_headers()
                self.wfile.write(data)

            def log_message(self, *args):
                pass

        self.server = QuietServer(("127.0.0.1", 0), Handler)
        self.url = f"http://127.0.0.1:{self.server.server_port}/v1/chat"
        threading.Thread(target=self.server.serve_forever, args=(0.01,), daemon=True).start()

    def close(self):
        self.server.shutdown()
        self.server.server_close()


@pytest.fixture
def provider():
    made = []

    def make(**kwargs):
        made.append(FakeProvider(**kwargs))
        return made[-1]

    yield make
    for p in made:
        p.close()


class Clock:
    def __init__(self):
        self.now = 1000.0

    def __call__(self):
        return self.now


def cfg_for(provider, **kwargs):
    return Config(KEY, provider.url, frozenset({"demo-small"}), **kwargs)


class NoRead(io.BytesIO):
    def read(self, *args):
        raise AssertionError("the body was read")


def call(app, body=GOOD, method="POST", path="/api/chat", client="203.0.113.7", length="auto", **extra):
    """Run a WSGI app once. Returns (status code, headers dict, body bytes)."""
    environ = {"REQUEST_METHOD": method, "PATH_INFO": path, "REMOTE_ADDR": client,
               "wsgi.input": body if isinstance(body, io.BytesIO) else io.BytesIO(body), **extra}
    if length == "auto":
        length = str(len(body.getvalue() if isinstance(body, io.BytesIO) else body))
    if length is not None:
        environ["CONTENT_LENGTH"] = length
    out = {}

    def start_response(status, headers):
        out["status"], out["headers"] = int(status.split()[0]), dict(headers)

    parts = app(environ, start_response)
    return out["status"], out["headers"], b"".join(parts)


# ---- the token bucket

def test_bucket_allows_a_burst_then_says_how_long_to_wait():
    bucket = TokenBucket(capacity=3, rate=1.0, now=0.0)
    assert [bucket.take(0.0) for _ in range(3)] == [0.0, 0.0, 0.0]
    assert bucket.take(0.0) == 1.0  # empty: one ticket arrives in 1 second
    assert bucket.take(0.5) == pytest.approx(0.5)  # a refused request spends nothing; half a ticket has arrived
    assert bucket.take(1.0) == 0.0


def test_bucket_never_holds_more_than_its_capacity():
    bucket = TokenBucket(capacity=3, rate=1.0, now=0.0)
    for _ in range(3):
        bucket.take(0.0)
    assert [bucket.take(1000.0) for _ in range(4)] == [0.0, 0.0, 0.0, 1.0]  # waiting 1,000 s buys only 3


# ---- what a good request does

def test_good_request_is_forwarded_with_the_key_and_the_server_sets_the_limits(provider):
    p = provider()
    body = json.dumps({"model": "demo-small", "messages": [{"role": "user", "content": "hi"}],
                       "max_tokens": 99999, "stream": True, "tools": [{"name": "x"}]}).encode()
    status, headers, out = call(make_app(cfg_for(p)), body)
    assert (status, json.loads(out)) == (200, {"answer": "pretend answer"})
    assert headers["Content-Type"] == "application/json"
    assert p.seen == [{"auth": f"Bearer {KEY}", "body": {
        "model": "demo-small", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 300}}]


def test_printing_the_config_does_not_print_the_key(provider):
    cfg = cfg_for(provider())
    assert KEY not in repr(cfg) and KEY not in str(cfg)


def test_config_comes_from_environment_variables_and_names_only_the_variables_when_one_is_missing():
    env = {"AI_API_KEY": KEY, "AI_UPSTREAM_URL": "http://127.0.0.1:1/x", "AI_MODELS": "demo-small, demo-large"}
    assert config_from_env(env).allowed_models == {"demo-small", "demo-large"}  # a space after the comma is not part of the name
    with pytest.raises(SystemExit) as exc:
        config_from_env({"AI_UPSTREAM_URL": "x", "AI_MODELS": "m"})
    assert "AI_API_KEY" in str(exc.value) and KEY not in str(exc.value)


@pytest.mark.parametrize("url", ["http://api.example.com/v1/chat", "api.example.com/v1/chat",
                                 "http://localhost.evil.example/v1/chat", "http://127.0.0.1.evil.example/v1/chat"])
def test_an_upstream_address_that_is_not_https_is_refused_unless_it_is_this_computer(url):
    """Guards a typo (http for https) that would send the key in clear text."""
    env = {"AI_API_KEY": KEY, "AI_UPSTREAM_URL": url, "AI_MODELS": "demo-small"}
    with pytest.raises(SystemExit) as exc:
        config_from_env(env)
    assert "https" in str(exc.value) and KEY not in str(exc.value)
    env["AI_UPSTREAM_URL"] = "https://api.example.com/v1/chat"
    assert config_from_env(env).upstream_url == env["AI_UPSTREAM_URL"]


# ---- rate limit

def test_sixth_request_in_a_burst_is_refused_and_never_reaches_the_provider(provider):
    p, clock = provider(), Clock()
    app = make_app(cfg_for(p, bucket_capacity=5, refill_per_second=1.0), clock)
    codes = [call(app)[0] for _ in range(6)]
    assert codes == [200] * 5 + [429]
    status, headers, _ = call(app)
    assert (status, headers["Retry-After"]) == (429, "1")
    assert len(p.seen) == 5
    clock.now += 1.0
    assert call(app)[0] == 200


def test_retry_after_rounds_a_fraction_of_a_second_up(provider):
    """Guards ceil() turning into int(): a wait of 0.4 s must say 1, because 0 tells the client to retry at once."""
    clock = Clock()
    app = make_app(cfg_for(provider(), bucket_capacity=1, refill_per_second=1.0), clock)
    assert call(app)[0] == 200
    clock.now += 0.6
    status, headers, _ = call(app)
    assert (status, headers["Retry-After"]) == (429, "1")


def test_each_client_address_has_its_own_bucket(provider):
    app = make_app(cfg_for(provider(), bucket_capacity=1), Clock())
    assert call(app, client="203.0.113.1")[0] == 200
    assert call(app, client="203.0.113.1")[0] == 429
    assert call(app, client="203.0.113.2")[0] == 200


def test_a_forwarded_for_header_cannot_buy_a_fresh_bucket(provider):
    app = make_app(cfg_for(provider(), bucket_capacity=1), Clock())
    assert call(app, HTTP_X_FORWARDED_FOR="198.51.100.1")[0] == 200
    assert call(app, HTTP_X_FORWARDED_FOR="198.51.100.2")[0] == 429


# ---- request size

def test_body_over_the_cap_is_refused_before_it_is_read(provider):
    p = provider()
    app = make_app(cfg_for(p, max_body_bytes=100))
    status, _, _ = call(app, NoRead(b"x" * 101), length="101")
    assert status == 413 and p.seen == []
    assert call(app, NoRead(b"x"), length="9999999999")[0] == 413
    assert call(app, NoRead(b"x"), length="9" * 5000)[0] == 413  # int() refuses more than 4,300 digits, which was a 500


@pytest.mark.parametrize("length", [None, "", "abc", "-5", "²", "1.5"])
def test_missing_or_odd_content_length_is_refused(provider, length):
    assert call(make_app(cfg_for(provider())), length=length)[0] == 411


def test_body_exactly_at_the_cap_passes(provider):
    app = make_app(cfg_for(provider(), max_body_bytes=len(GOOD)))
    assert call(app)[0] == 200


@pytest.mark.parametrize("body, length, status", [(GOOD, str(len(GOOD)).zfill(12), 200), (b"", "0" * 12, 400)],
                         ids=["padded-real-length", "padded-zero"])
def test_a_content_length_with_leading_zeros_is_the_number_it_spells(provider, body, length, status):
    """Guards the nine-digit bound counting padding (a valid length got 413) and int("") on an all-zero length."""
    assert call(make_app(cfg_for(provider())), body, length=length)[0] == status


# ---- model allowlist and shape

def test_a_model_that_is_not_on_the_list_is_refused_and_not_forwarded(provider):
    p = provider()
    body = json.dumps({"model": "demo-huge-and-costly", "messages": [{"role": "user", "content": "hi"}]}).encode()
    status, _, out = call(make_app(cfg_for(p)), body)
    assert (status, json.loads(out)) == (400, {"error": "model not allowed"}) and p.seen == []


@pytest.mark.parametrize("body", [
    b"", b"not json", b"[]", b"7", b'{"model": "demo-small"}', b'{"messages": [{"role": "user"}]}',
    b'{"model": ["demo-small"], "messages": [{"role": "user"}]}',
    b'{"model": "demo-small", "messages": []}', b'{"model": "demo-small", "messages": "hi"}', b"\xff\xfe"])
def test_malformed_json_is_refused_without_a_crash(provider, body):
    p = provider()
    assert call(make_app(cfg_for(p)), body)[0] == 400 and p.seen == []


def test_a_parser_that_runs_out_of_stack_gives_400_not_a_crash(monkeypatch):
    """Guards Python 3.11 and earlier, where json.loads raises RecursionError on about a thousand open brackets."""
    def too_deep(*args):
        raise RecursionError

    monkeypatch.setattr(json, "loads", too_deep)
    cfg = Config(KEY, "http://127.0.0.1:9/v1/chat", frozenset({"demo-small"}))
    assert call(make_app(cfg), b"[" * 4000)[0] == 400


def test_only_post_on_the_one_route_is_served(provider):
    app = make_app(cfg_for(provider()))
    assert call(app, path="/")[0] == 404 and call(app, path="/api/chat/")[0] == 404
    status, headers, _ = call(app, method="GET")
    assert (status, headers["Allow"]) == (405, "POST")


# ---- upstream trouble

def test_provider_error_becomes_a_plain_502(provider):
    p = provider(mode="error-with-key")
    status, _, out = call(make_app(cfg_for(p)))
    assert (status, json.loads(out)) == (502, {"error": "upstream failed"})


def test_provider_that_never_answers_gives_502_after_the_timeout(provider):
    p = provider(mode="slow")
    assert call(make_app(cfg_for(p, timeout_seconds=0.05)))[0] == 502


def test_nothing_listening_gives_502():
    cfg = Config(KEY, "http://127.0.0.1:9/v1/chat", frozenset({"demo-small"}))
    assert call(make_app(cfg))[0] == 502


def test_a_redirect_is_not_followed_so_the_key_is_not_sent_elsewhere(provider):
    elsewhere = provider()
    redirecting = provider(mode="redirect", redirect_to=elsewhere.url)
    status, _, _ = call(make_app(cfg_for(redirecting)))
    assert status == 502
    assert redirecting.seen and elsewhere.seen == []


# ---- the key never comes back out

def test_a_provider_that_echoes_the_key_is_scrubbed_and_the_test_is_not_vacuous(provider):
    p = provider(mode="echo-key")
    status, _, out = call(make_app(cfg_for(p)))
    assert p.seen[0]["auth"] == f"Bearer {KEY}"  # the provider really did receive it, and repeated it
    assert status == 200 and KEY.encode() not in out and b"[removed]" in out


@pytest.mark.parametrize("mode", ["ok", "echo-key", "error-with-key", "key-in-reason", "slow", "redirect"])
def test_the_key_is_in_no_response_body_header_or_log_line(provider, caplog, mode):
    caplog.set_level(logging.DEBUG)
    elsewhere = provider()
    p = provider(mode=mode, redirect_to=elsewhere.url)
    clock = Clock()
    app = make_app(cfg_for(p, bucket_capacity=1, max_body_bytes=200, timeout_seconds=0.05), clock)
    seen = []
    seen.append(call(app))  # a good request, whatever the provider then does
    seen.append(call(app, client="203.0.113.9"))
    seen.append(call(app))  # runs out of tickets: 429
    seen.append(call(app, b"x" * 300, client="203.0.113.10"))  # 413
    seen.append(call(app, b"not json", client="203.0.113.11"))  # 400
    seen.append(call(app, method="GET", client="203.0.113.12"))  # 405
    seen.append(call(app, path=f"/{KEY}", client="203.0.113.13"))  # a path that contains the key: 404
    seen.append(call(app, length=None, client="203.0.113.14"))  # 411
    assert {s[0] for s in seen} >= {400, 404, 405, 411, 413, 429}
    for status, headers, body in seen:
        assert KEY.encode() not in body, status
        assert all(KEY not in f"{name}: {value}" for name, value in headers.items()), status
    assert caplog.records
    for record in caplog.records:
        text = " ".join([record.getMessage(), str(record.args), str(record.exc_text), record.pathname])
        assert KEY not in text


# ---- through a real server socket

def stderr_after_raw_request(handler_class, request_line, capsys, cfg):
    """Serve one raw request through a real wsgiref server and return what the server printed to stderr."""
    server = make_server("127.0.0.1", 0, make_app(cfg), handler_class=handler_class)
    threading.Thread(target=server.serve_forever, args=(0.01,), daemon=True).start()
    try:
        with socket.create_connection(("127.0.0.1", server.server_port), timeout=5) as conn:
            conn.sendall(f"{request_line}\r\nContent-Length: 0\r\n\r\n".encode())
            while conn.recv(65536):  # the server closes after it has answered and logged
                pass
    finally:
        server.shutdown()
        server.server_close()
    return capsys.readouterr().err


@pytest.mark.parametrize("request_line", [f"POST /{KEY} HTTP/1.1", f"GET /x {KEY}"], ids=["key-as-path", "key-in-a-bad-request-line"])
def test_the_servers_own_access_log_leaves_out_the_request_line(capsys, request_line):
    """Guards PathFreeHandler: the standard handler prints the whole request line, so a key used as the path was logged."""
    cfg = Config(KEY, "http://127.0.0.1:9/v1/chat", frozenset({"demo-small"}))
    assert KEY in stderr_after_raw_request(WSGIRequestHandler, request_line, capsys, cfg)  # the stock handler really prints it
    printed = stderr_after_raw_request(PathFreeHandler, request_line, capsys, cfg)
    assert KEY not in printed and re.search(r"client=127\.0\.0\.1 (status=\d+|protocol error)", printed)


def test_main_serves_with_the_path_free_handler(monkeypatch):
    """Guards the wiring: a main() that leaves handler_class out serves with the standard handler and logs the key."""
    class Stop(Exception):
        pass

    seen = {}

    def fake_make_server(host, port, app, handler_class=None):
        seen["handler_class"] = handler_class
        raise Stop  # no socket is opened

    monkeypatch.setattr(logging, "basicConfig", lambda **kwargs: None)
    monkeypatch.setattr("keyproxy.make_server", fake_make_server)
    monkeypatch.setenv("AI_API_KEY", KEY)
    monkeypatch.setenv("AI_UPSTREAM_URL", "https://api.example.com/v1/chat")
    monkeypatch.setenv("AI_MODELS", "demo-small")
    with pytest.raises(Stop):
        keyproxy.main()
    assert seen["handler_class"] is PathFreeHandler


def test_end_to_end_through_a_real_wsgi_server(provider):
    class Quiet(WSGIRequestHandler):
        def log_message(self, *args):
            pass

    p = provider()
    server = make_server("127.0.0.1", 0, make_app(cfg_for(p)), handler_class=Quiet)
    threading.Thread(target=server.serve_forever, args=(0.01,), daemon=True).start()
    try:
        request = urllib.request.Request(f"http://127.0.0.1:{server.server_port}/api/chat", GOOD,
                                         headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(request, timeout=5) as response:
            body = response.read()
            headers = str(response.headers)
    finally:
        server.shutdown()
        server.server_close()
    assert json.loads(body) == {"answer": "pretend answer"} and KEY not in headers
    assert p.seen[0]["auth"] == f"Bearer {KEY}"


def test_every_python_block_in_the_lesson_is_a_verbatim_excerpt_of_the_source():
    """The lesson quotes this file. If the code changes and the lesson does not, the reader sees code that was never tested."""
    lesson = (Path(__file__).parent.parent / "lesson.md").read_text()
    source = Path(__file__).with_name("keyproxy.py").read_text().splitlines()
    blocks = re.findall(r"^```python\n(.*?)^```", lesson, flags=re.S | re.M)
    assert blocks, "the lesson has no python blocks to check"
    for block in blocks:
        lines = block.rstrip("\n").splitlines()
        windows = (textwrap.dedent("\n".join(source[i:i + len(lines)])) for i in range(len(source) - len(lines) + 1))
        assert textwrap.dedent("\n".join(lines)) in windows, f"not in keyproxy.py:\n{block}"