"""Transcript paging and what the transcript keeps, through the real daemon."""

from dataclasses import dataclass
import http.client
import http.server
from collections.abc import Callable
import json
import resource
import socket
import threading
import time
import unittest
import urllib.parse

from harness import Albedo, Provider, Reply, exclusive, python, text

# More client connections than a macOS shell's default open-file limit.
CONNECTIONS = 300


@dataclass
class HeldCatalog:
    url: str
    release: Callable[[], None]


def held_catalog():
    """A local models catalog whose every request waits until it is released."""
    released = threading.Event()

    class Handler(http.server.BaseHTTPRequestHandler):
        def log_message(self, format, *args):
            pass

        def do_GET(self):
            released.wait()

    server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
    server.daemon_threads = True
    threading.Thread(target=server.serve_forever, daemon=True).start()

    def release():
        released.set()
        server.shutdown()
        server.server_close()

    return HeldCatalog(f"http://127.0.0.1:{server.server_port}/api.json", release)


def read_response(connection):
    """One HTTP response off a keep-alive socket: its status code."""
    head = b""
    while b"\r\n\r\n" not in head:
        chunk = connection.recv(65536)
        if not chunk:
            raise AssertionError("connection closed before a response")
        head += chunk
    lines = head.split(b"\r\n\r\n", 1)
    body = lines[1] if len(lines) > 1 else b""
    length = 0
    for line in lines[0].split(b"\r\n")[1:]:
        name, _, value = line.lower().partition(b":")
        if name.strip() == b"content-length":
            length = int(value.strip())
    while len(body) < length:
        chunk = connection.recv(65536)
        if not chunk:
            raise AssertionError("connection closed mid-body")
        body += chunk
    return int(lines[0].split(b" ")[1])


class DaemonTest(unittest.TestCase):
    def test_history_pages_recover_older_turns_once(self):
        provider = Provider(lambda _request: text("answer"))
        self.addCleanup(provider.close)
        with Albedo(provider) as app:
            session = app.session()
            for message in ("oldest prompt", "middle prompt", "newest prompt"):
                app.prompt(session, message).close()
                app.idle(session)

            pages = []
            before = None
            while True:
                suffix = "" if before is None else f"&before={before}"
                with app.api(f"/sessions/{session}/history?rows=2{suffix}") as response:
                    page = json.load(response)
                pages.append(page)
                if not page["more"]:
                    break
                self.assertIsInstance(page["before"], int)
                self.assertNotEqual(page["before"], before)
                before = page["before"]

            prompts = [
                [event["text"] for event in page["events"] if event["type"] == "user"]
                for page in pages
            ]
            self.assertEqual(
                [turn for turn in prompts if turn],
                [["newest prompt"], ["middle prompt"], ["oldest prompt"]],
            )
            self.assertFalse(pages[-1]["more"])

    @exclusive
    def test_a_clean_restart_leaves_no_registry_crash_loop(self):
        # Shutdown closes the store before the VM halts, and the supervisor
        # may restart the registry inside that window; its init used to panic
        # on the closed store ten times before the supervisor gave up. What
        # holds the VM open long enough to lose the race is a models catalog
        # fetch still in flight, so the restart runs against a local catalog
        # that never answers; the fixture config itself stays, so later
        # fixtures keep their provider.
        provider = Provider(lambda _request: text("answer"))
        self.addCleanup(provider.close)
        catalog = held_catalog()
        self.addCleanup(catalog.release)
        with Albedo(provider) as app:
            settings = json.loads((app.home / "extensions.json").read_text())
            settings["models"] = {"url": catalog.url}
            (app.home / "extensions.json").write_text(json.dumps(settings))
            (app.home / "models.json").unlink(missing_ok=True)
            log = app.home / "daemon.log"
            before = log.stat().st_size if log.exists() else 0
            app.restart()
            app.restart()
            tail = log.read_text(errors="replace")[before:]
            for marker in ("Noproc", "reached_max_restart_intensity", "callee exited"):
                self.assertNotIn(marker, tail)

    @exclusive
    def test_requests_during_a_restart_are_answered_not_crashed(self):
        # Shutdown closes each session and then the store before the VM halts,
        # and requests keep arriving meanwhile: ones queued behind the drain
        # used to die with the registry ("callee exited"), and later ones
        # found no registry at all. Kernels make the drain take a while, a
        # held catalog fetch keeps the VM up past the store, and a second
        # shutdown must not start a second drain.
        def script(request):
            called = any(
                message.get("role") == "tool" for message in request["messages"]
            )
            return text("done") if called else python("x = 1")

        provider = Provider(script)
        self.addCleanup(provider.close)
        catalog = held_catalog()
        self.addCleanup(catalog.release)
        with Albedo(provider) as app:
            settings = json.loads((app.home / "extensions.json").read_text())
            settings["models"] = {"url": catalog.url}
            (app.home / "extensions.json").write_text(json.dumps(settings))
            (app.home / "models.json").unlink(missing_ok=True)
            sessions = [app.session() for _ in range(3)]
            for session in sessions:
                app.prompt(session, "keep a variable").close()
            for session in sessions:
                app.idle(session)
            log = app.home / "daemon.log"
            before = log.stat().st_size if log.exists() else 0

            stop = threading.Event()
            paths = [("GET", "/sessions", None), ("GET", "/models/fixture", None)]
            for session in sessions:
                paths += [
                    ("GET", f"/sessions/{session}/status", None),
                    ("GET", f"/sessions/{session}/tree", None),
                    ("GET", f"/sessions/{session}/children", None),
                    ("PATCH", f"/sessions/{session}", {"name": ""}),
                ]

            # One kept-alive connection per worker, reopened only when the
            # daemon drops it: a fresh socket per request exhausts the host's
            # ephemeral ports within seconds, which fails every other
            # connection on the machine, live sessions included.
            def hammer(base, token):
                address = urllib.parse.urlsplit(base)
                headers = {
                    "Authorization": "Bearer " + token,
                    "Content-Type": "application/json",
                }
                connection = None
                while not stop.is_set():
                    for method, path, body in paths:
                        try:
                            connection = connection or http.client.HTTPConnection(
                                address.hostname, address.port, timeout=20
                            )
                            payload = None if body is None else json.dumps(body)
                            connection.request(method, path, payload, headers)
                            connection.getresponse().read()
                        except (http.client.HTTPException, OSError):
                            if connection:
                                connection.close()
                            connection = None
                            stop.wait(0.01)
                if connection:
                    connection.close()

            workers = [
                threading.Thread(
                    target=hammer, args=(app.base, app.connection["token"])
                )
                for _ in range(6)
            ]
            for worker in workers:
                worker.start()
            try:
                # restart() asks again while the first drain runs.
                app.api("/shutdown", {}).close()
                app.restart()
            finally:
                stop.set()
                for worker in workers:
                    worker.join()
            tail = log.read_text(errors="replace")[before:]
            for marker in (
                "Noproc",
                "callee exited",
                "Callee subject had no owner",
                "reached_max_restart_intensity",
            ):
                self.assertNotIn(marker, tail)

    def test_a_late_request_body_keeps_the_keep_alive_connection(self):
        # mist parses every later TCP read as the next keep-alive request, so
        # a route that answers without reading its request body leaves those
        # bytes to that parser: a body arriving after its headers reads as a
        # new request, fails, and the connection closes without a response,
        # taking the client's next request with it. The router reads the body
        # up front, so a split write cannot cost the connection.
        provider = Provider(lambda _request: text("answer"))
        self.addCleanup(provider.close)
        with Albedo(provider) as app:
            address = urllib.parse.urlsplit(app.base)
            token = app.connection["token"]
            session = app.session()
            interrupt = f"/sessions/{session}/interrupt"

            def sent(method, path, body):
                # The headers and the body as two writes, so the route answers
                # in between, before the body is read.
                head = (
                    f"{method} {path} HTTP/1.1\r\n"
                    f"Host: {address.hostname}\r\n"
                    f"Authorization: Bearer {token}\r\n"
                    f"Content-Type: application/json\r\n"
                    f"Content-Length: {len(body)}\r\n\r\n"
                ).encode()
                connection = socket.create_connection(
                    (address.hostname, address.port), timeout=20
                )
                try:
                    connection.sendall(head)
                    time.sleep(0.05)
                    connection.sendall(body)
                    first = read_response(connection)
                    connection.sendall(
                        (
                            f"GET /health HTTP/1.1\r\n"
                            f"Host: {address.hostname}\r\n"
                            f"Authorization: Bearer {token}\r\n\r\n"
                        ).encode()
                    )
                    return first, read_response(connection)
                finally:
                    connection.close()

            # Every shape that answers without reading its body: a plain
            # route, an ignored POST body, an unknown route, and a route that
            # answers a session actor call.
            for method, path, status in (
                ("GET", "/health", 200),
                ("POST", "/auth/credentials/migration", 200),
                ("POST", "/nope", 404),
                ("POST", interrupt, 200),
            ):
                with self.subTest(f"{method} {path}"):
                    self.assertEqual(sent(method, path, b"{}"), (status, 200))

    @exclusive
    def test_a_daemon_started_from_a_macos_shell_answers_hundreds_of_connections(self):
        # macOS starts a shell at a soft limit of 256 open files. The CLI, like
        # any Go program, lifts its own limit but starts children on the
        # original one, and a few hundred client connections would then run
        # the daemon out of files and take its listener down mid-accept.
        hard = resource.getrlimit(resource.RLIMIT_NOFILE)[1]

        def prepare(app):
            app.daemon.open_files = (256, hard)

        with Albedo(prepare=prepare) as app:
            headers = {"Authorization": "Bearer " + app.connection["token"]}
            connections = []
            for _ in range(CONNECTIONS):
                connection = http.client.HTTPConnection(
                    "127.0.0.1", app.connection["port"], timeout=20
                )
                connection.connect()
                self.addCleanup(connection.close)
                connections.append(connection)
            for connection in connections:
                connection.request("GET", "/health", headers=headers)
                response = connection.getresponse()
                response.read()
                self.assertEqual(response.status, 200)

    @exclusive
    def test_a_listener_out_of_files_comes_back_on_the_daemons_port(self):
        # Accepts past the open-file limit crash glisten's acceptors until mist
        # restarts the whole listener, on the port it was built with. A port
        # the OS picks at each start would bring it back somewhere daemon.json
        # does not say, leaving a live daemon no client can reach.
        def prepare(app):
            app.daemon.open_files = (128, 128)

        with Albedo(prepare=prepare) as app:
            pid, port = app.connection["pid"], app.connection["port"]
            flood = []
            for _ in range(CONNECTIONS):
                try:
                    flood.append(
                        socket.create_connection(("127.0.0.1", port), timeout=20)
                    )
                except OSError:
                    pass
            # Nothing the daemon sends ends these sockets: each one ends when
            # the restarted listener's tree drops the connection it accepted,
            # or resets the one still waiting in its old socket's backlog.
            for connection in flood:
                try:
                    self.assertEqual(connection.recv(1), b"")
                except ConnectionResetError:
                    pass
                connection.close()

            record = json.loads((app.home / "daemon.json").read_text())
            self.assertEqual((record["pid"], record["port"]), (pid, port))
            deadline = time.monotonic() + 20
            while True:
                try:
                    with app.api("/health") as response:
                        self.assertEqual(response.status, 200)
                    break
                except OSError:
                    self.assertLess(
                        time.monotonic(), deadline, "the daemon's port stayed closed"
                    )
                    time.sleep(0.1)

    def test_a_thoughts_duration_is_kept_with_the_transcript(self):
        # summarized thinking streams once it is written: here the response
        # opens, thinks 0.6s, and its summary comes 0.3s before the answer,
        # so the thought is timed from the opening, not the summary
        def chunk(delta, finish=None):
            return {
                "id": "fixture",
                "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
            }

        reply = Reply(
            "text",
            delay=0.3,
            events=[
                chunk({"role": "assistant"}),
                chunk({"reasoning_content": "weighing it"}),
                chunk({"content": "answer"}),
                chunk({}, "stop"),
            ],
        )
        provider = Provider(lambda _request: reply)
        self.addCleanup(provider.close)
        with Albedo(provider) as app:
            session = app.session()
            app.prompt(session, "think first").close()
            app.idle(session)

            for source, events in (
                ("stream", app.events(session)),
                ("history", app.history(session)["events"]),
            ):
                thoughts = [event for event in events if event["type"] == "thinking"]
                self.assertEqual(
                    [event["text"] for event in thoughts], ["weighing it"], source
                )
                self.assertGreaterEqual(thoughts[0].get("elapsedMs", 0), 500, source)