from __future__ import annotations

import base64
import hashlib
import json
import os
import signal
import socket
import struct
import subprocess
import sys
import tempfile
import threading
import time
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from types import SimpleNamespace
from urllib.error import HTTPError
from urllib.request import urlopen

from scripts.agent_benchmark.browser_cdp import (
    MAX_MESSAGE_BYTES,
    WEBSOCKET_GUID,
    BrowserError,
    BrowserRenderer,
    _CDP,
    _StaticServer,
    _terminate_owned_process_group,
)
from scripts.agent_benchmark.web_validation import _runtime_gates


def _frame(payload: bytes, *, opcode: int = 1, fin: bool = True, masked: bool = False) -> bytes:
    first = (0x80 if fin else 0) | opcode
    size = len(payload)
    if size < 126:
        header = bytes((first, (0x80 if masked else 0) | size))
    elif size <= 65535:
        header = bytes((first, (0x80 if masked else 0) | 126)) + struct.pack("!H", size)
    else:
        header = bytes((first, (0x80 if masked else 0) | 127)) + struct.pack("!Q", size)
    if not masked:
        return header + payload
    mask = b"mask"
    return header + mask + bytes(value ^ mask[index % 4] for index, value in enumerate(payload))


def _exact(connection: socket.socket, size: int) -> bytes:
    chunks = bytearray()
    while len(chunks) < size:
        chunk = connection.recv(size - len(chunks))
        if not chunk:
            raise EOFError
        chunks.extend(chunk)
    return bytes(chunks)


def _client_frame(connection: socket.socket) -> tuple[int, bytes]:
    first, second = _exact(connection, 2)
    size = second & 0x7F
    if size == 126:
        size = struct.unpack("!H", _exact(connection, 2))[0]
    elif size == 127:
        size = struct.unpack("!Q", _exact(connection, 8))[0]
    if not second & 0x80:
        raise AssertionError("client frame was not masked")
    mask = _exact(connection, 4)
    payload = _exact(connection, size)
    return first & 0x0F, bytes(
        value ^ mask[index % 4] for index, value in enumerate(payload)
    )


class _WebSocketFixture:
    def __init__(self, script, *, valid_accept: bool = True):
        self.script = script
        self.valid_accept = valid_accept
        self.error: BaseException | None = None
        self.listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.listener.bind(("127.0.0.1", 0))
        self.listener.listen(1)
        self.thread = threading.Thread(target=self._serve, daemon=True)
        self.thread.start()

    @property
    def url(self) -> str:
        return f"ws://127.0.0.1:{self.listener.getsockname()[1]}/devtools/page/1"

    def _serve(self) -> None:
        try:
            connection, _ = self.listener.accept()
            with connection:
                request = bytearray()
                while b"\r\n\r\n" not in request:
                    request.extend(connection.recv(4096))
                key = ""
                for line in bytes(request).split(b"\r\n"):
                    if line.lower().startswith(b"sec-websocket-key:"):
                        key = line.split(b":", 1)[1].strip().decode("ascii")
                accept = base64.b64encode(
                    hashlib.sha1((key + WEBSOCKET_GUID).encode("ascii")).digest()
                ).decode("ascii")
                if not self.valid_accept:
                    accept = "invalid"
                connection.sendall(
                    (
                        "HTTP/1.1 101 Switching Protocols\r\n"
                        "Upgrade: websocket\r\nConnection: keep-alive, Upgrade\r\n"
                        f"Sec-WebSocket-Accept: {accept}\r\n\r\n"
                    ).encode("ascii")
                )
                if self.valid_accept:
                    self.script(connection)
        except BaseException as exc:  # surfaced by close()
            self.error = exc
        finally:
            self.listener.close()

    def close(self) -> None:
        self.thread.join(2)
        if self.thread.is_alive():
            self.listener.close()
            self.thread.join(2)
        if self.error is not None:
            raise self.error


class BrowserProtocolTest(unittest.TestCase):
    def test_loopback_handler_rejects_escape_and_all_symlinks(self):
        with tempfile.TemporaryDirectory() as raw:
            root = Path(raw)
            (root / "nested").mkdir()
            (root / "index.html").write_text("ok", encoding="utf-8")
            (root / "nested" / "ok.txt").write_text("nested", encoding="utf-8")
            outside = root.parent / f"browser-cdp-outside-{root.name}.txt"
            outside.write_text("no", encoding="utf-8")
            (root / "internal-link").symlink_to(root / "index.html")
            (root / "outside-link").symlink_to(outside)
            (root / "nested-link").symlink_to(root / "nested", target_is_directory=True)
            server = _StaticServer(root)
            thread = threading.Thread(target=server.serve_forever, daemon=True)
            thread.start()
            try:
                base = f"http://127.0.0.1:{server.server_address[1]}"
                with urlopen(base + "/index.html") as response:
                    self.assertEqual(response.read(), b"ok")
                for path in (
                    "/internal-link",
                    "/outside-link",
                    "/nested-link/ok.txt",
                    "/%2e%2e/index.html",
                    "/nested/%2e%2e/index.html",
                ):
                    with self.subTest(path=path), self.assertRaises(HTTPError):
                        urlopen(base + path)
            finally:
                server.shutdown()
                server.server_close()
                thread.join(2)
                outside.unlink(missing_ok=True)

    def test_fragment_ping_event_and_response_correlation(self):
        def script(connection: socket.socket) -> None:
            opcode, payload = _client_frame(connection)
            command = json.loads(payload)
            self.assertEqual(opcode, 1)
            connection.sendall(_frame(b'{"method":"Page.ready","params":{"ok":true}}'))
            connection.sendall(_frame(b"ping", opcode=9))
            response = json.dumps({"id": command["id"], "result": {"ok": True}}).encode()
            split = len(response) // 2
            connection.sendall(_frame(response[:split], fin=False))
            connection.sendall(_frame(response[split:], opcode=0))
            pong, pong_payload = _client_frame(connection)
            self.assertEqual((pong, pong_payload), (10, b"ping"))

        server = _WebSocketFixture(script)
        client = _CDP(server.url, time.monotonic() + 2)
        try:
            self.assertEqual(client.call("Page.enable"), {"ok": True})
            self.assertEqual(client.events[0]["method"], "Page.ready")
        finally:
            client.close()
            server.close()

    def test_fire_and_call_responses_can_be_interleaved(self):
        def script(connection: socket.socket) -> None:
            _, first = _client_frame(connection)
            _, second = _client_frame(connection)
            fire, call = json.loads(first), json.loads(second)
            connection.sendall(_frame(json.dumps({"id": fire["id"], "result": {}}).encode()))
            connection.sendall(_frame(b'{"method":"Fetch.paused","params":{}}'))
            connection.sendall(_frame(json.dumps({"id": call["id"], "result": {"done": 1}}).encode()))

        server = _WebSocketFixture(script)
        client = _CDP(server.url, time.monotonic() + 2)
        try:
            client.fire("Fetch.continueRequest", {"requestId": "request"})
            self.assertEqual(client.call("Runtime.evaluate"), {"done": 1})
            self.assertEqual(client.events[0]["method"], "Fetch.paused")
        finally:
            client.close()
            server.close()

    def test_extended_frame_lengths_are_canonical_and_bounded(self):
        for size in (200, 70_000):
            with self.subTest(size=size):
                value = "x" * size

                def script(connection: socket.socket, value=value) -> None:
                    _, payload = _client_frame(connection)
                    command = json.loads(payload)
                    response = json.dumps(
                        {"id": command["id"], "result": {"value": value}},
                        separators=(",", ":"),
                    ).encode()
                    connection.sendall(_frame(response))

                server = _WebSocketFixture(script)
                client = _CDP(server.url, time.monotonic() + 2)
                try:
                    self.assertEqual(
                        client.call("Runtime.evaluate")["value"], value
                    )
                finally:
                    client.close()
                    server.close()

    def test_handshake_and_malformed_frame_matrix_fail_closed(self):
        bad_accept = _WebSocketFixture(lambda _connection: None, valid_accept=False)
        with self.assertRaises(BrowserError):
            _CDP(bad_accept.url, time.monotonic() + 1)
        bad_accept.close()

        cases = {
            "masked": _frame(b"{}", masked=True),
            "binary": _frame(b"{}", opcode=2),
            "continuation": _frame(b"{}", opcode=0),
            "malformed-json": _frame(b"{"),
            "close": _frame(b"", opcode=8),
            "oversized": bytes((0x81, 127)) + struct.pack("!Q", MAX_MESSAGE_BYTES + 1),
            "noncanonical": bytes((0x81, 126)) + struct.pack("!H", 1) + b"x",
        }
        for name, response in cases.items():
            with self.subTest(name=name):
                def script(connection: socket.socket, response=response) -> None:
                    _client_frame(connection)
                    connection.sendall(response)

                server = _WebSocketFixture(script)
                client = _CDP(server.url, time.monotonic() + 1)
                try:
                    with self.assertRaises(BrowserError):
                        client.call("Runtime.evaluate")
                finally:
                    client.close()
                    server.close()

    def test_uncorrelated_response_and_deadline_fail_closed(self):
        def wrong_id(connection: socket.socket) -> None:
            _, payload = _client_frame(connection)
            ident = json.loads(payload)["id"]
            connection.sendall(_frame(json.dumps({"id": ident + 1, "result": {}}).encode()))

        server = _WebSocketFixture(wrong_id)
        client = _CDP(server.url, time.monotonic() + 1)
        try:
            with self.assertRaises(BrowserError):
                client.call("Runtime.evaluate")
        finally:
            client.close()
            server.close()

        def no_response(connection: socket.socket) -> None:
            _client_frame(connection)
            time.sleep(0.2)

        server = _WebSocketFixture(no_response)
        client = _CDP(server.url, time.monotonic() + 0.05)
        try:
            with self.assertRaises(BrowserError):
                client.call("Runtime.evaluate")
        finally:
            client.close()
            server.close()

    def test_owned_process_group_is_terminated_and_reaped(self):
        source = (
            "import signal,subprocess,sys,time;"
            "signal.signal(signal.SIGTERM,signal.SIG_IGN);"
            "p=subprocess.Popen([sys.executable,'-c',"
            "'import signal,time;signal.signal(signal.SIGTERM,signal.SIG_IGN);time.sleep(60)']);"
            "print(p.pid,flush=True);time.sleep(60)"
        )
        process = subprocess.Popen(
            [sys.executable, "-c", source],
            stdout=subprocess.PIPE,
            text=True,
            start_new_session=True,
        )
        assert process.stdout is not None
        child_pid = int(process.stdout.readline().strip())
        try:
            _terminate_owned_process_group(process, grace=1)
            self.assertIsNotNone(process.poll())
            with self.assertRaises(ProcessLookupError):
                os.killpg(process.pid, 0)
            with self.assertRaises(ProcessLookupError):
                os.kill(child_pid, 0)
        finally:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            process.wait(timeout=2)
            process.stdout.close()

    def test_renderer_error_reaps_its_owned_process_group(self):
        with tempfile.TemporaryDirectory(dir=Path.cwd()) as raw:
            root = Path(raw)
            (root / "index.html").write_text("<main>unused</main>", encoding="utf-8")
            pid_path = root / "browser-pids"
            binary = root / "fake-browser"
            binary.write_text(
                "#!/usr/bin/env python3\n"
                "import os, signal, subprocess, sys, time\n"
                "signal.signal(signal.SIGTERM, signal.SIG_IGN)\n"
                "child = subprocess.Popen([sys.executable, '-c', "
                "'import signal,time;signal.signal(signal.SIGTERM,signal.SIG_IGN);time.sleep(60)'])\n"
                f"with open({str(pid_path)!r}, 'w', encoding='ascii') as handle:\n"
                "    handle.write(f'{os.getpid()} {child.pid}\\n')\n"
                "    handle.flush()\n"
                "    os.fsync(handle.fileno())\n"
                "time.sleep(60)\n",
                encoding="utf-8",
            )
            binary.chmod(0o700)
            pids: list[int] = []
            try:
                with self.assertRaisesRegex(BrowserError, "browser_cdp_unavailable"):
                    BrowserRenderer(str(binary)).render(
                        workspace_root=root,
                        output_root=root,
                        viewports=(
                            SimpleNamespace(
                                id="mobile.small+wide", width=375, height=700
                            ),
                        ),
                        timeout_seconds=1,
                    )
                pids = [int(value) for value in pid_path.read_text().split()]
                self.assertEqual(len(pids), 2)
                for pid in pids:
                    with self.subTest(pid=pid), self.assertRaises(ProcessLookupError):
                        os.kill(pid, 0)
            finally:
                for pid in pids:
                    try:
                        os.kill(pid, signal.SIGKILL)
                    except ProcessLookupError:
                        pass


class _Counter(BaseHTTPRequestHandler):
    count = 0

    def do_GET(self):
        type(self).count += 1
        self.send_response(200)
        self.end_headers()

    def log_message(self, *_args):
        pass


class BrowserIntegrationTest(unittest.TestCase):
    _VIEWPORTS = (
        SimpleNamespace(id="desktop", width=900, height=700),
        SimpleNamespace(id="mobile.small+wide", width=375, height=700),
    )

    @staticmethod
    def _valid_page(
        root: Path,
        extra_image: str = "",
        *,
        focus_css: str = "",
        autofocus: bool = False,
    ) -> None:
        (root / "assets").mkdir()
        for name in ("a.svg", "b.svg"):
            (root / "assets" / name).write_text(
                "<svg xmlns='http://www.w3.org/2000/svg' width='80' height='60'/>",
                encoding="utf-8",
            )
        focus_attribute = " autofocus" if autofocus else ""
        (root / "index.html").write_text(
            "<link rel='stylesheet' href='styles.css'>"
            "<main><h1>Ready</h1><img src='assets/a.svg' alt='A'>"
            "<img src='assets/b.svg' alt='B'>"
            f"{extra_image}<a href='#x'{focus_attribute}>go</a></main>"
            "<script src='script.js'></script>",
            encoding="utf-8",
        )
        (root / "styles.css").write_text(
            f"body{{color:#111;background:#fff}}{focus_css}", encoding="utf-8"
        )
        (root / "script.js").write_text("document.body.dataset.ready='1'", encoding="utf-8")

    def test_valid_page_emits_complete_two_viewport_observations(self):
        with tempfile.TemporaryDirectory() as raw:
            root = Path(raw)
            self._valid_page(root)
            render = BrowserRenderer().render(
                workspace_root=root,
                output_root=root,
                viewports=self._VIEWPORTS,
                timeout_seconds=20,
            )
            self.assertEqual(
                [view.id for view in render.viewports],
                ["desktop", "mobile.small+wide"],
            )
            self.assertTrue(all((root / view.screenshot).stat().st_size > 0 for view in render.viewports))
            self.assertTrue(all(len(view.image_facts) == 2 for view in render.viewports))
            self.assertFalse([item for item in render.requests if item["kind"] == "external"])

    def test_denied_page_dispatches_no_external_request(self):
        _Counter.count = 0
        with tempfile.TemporaryDirectory() as raw:
            root = Path(raw)
            counter = ThreadingHTTPServer(("127.0.0.1", 0), _Counter)
            thread = threading.Thread(target=counter.serve_forever, daemon=True)
            thread.start()
            try:
                self._valid_page(
                    root,
                    f"<img src='http://127.0.0.1:{counter.server_address[1]}/leak' alt='blocked'>",
                )
                render = BrowserRenderer().render(
                    workspace_root=root,
                    output_root=root,
                    viewports=self._VIEWPORTS,
                    timeout_seconds=20,
                )
                self.assertEqual(_Counter.count, 0)
                self.assertTrue(any(item["kind"] == "external" for item in render.requests))
            finally:
                counter.shutdown()
                counter.server_close()
                thread.join(2)

    def test_focus_visibility_uses_computed_indicator(self):
        manifest = SimpleNamespace(
            fixture=SimpleNamespace(
                assets=tuple(
                    SimpleNamespace(workspace_path=f"assets/{name}.svg")
                    for name in ("a", "b")
                )
            ),
            viewports=self._VIEWPORTS,
        )
        cases = (
            ("suppressed", "a:focus{outline:none;box-shadow:none}", False, False),
            (
                "transparent-outline",
                "a:focus{outline:4px solid rgba(0,85,255,0);box-shadow:none}",
                False,
                False,
            ),
            (
                "transparent-shadow",
                "a:focus{outline:none;box-shadow:0 0 0 4px rgba(0,85,255,0)}",
                False,
                False,
            ),
            (
                "transparent-border",
                "a:focus{outline:none;box-shadow:none;border:4px solid rgba(0,85,255,0)}",
                False,
                False,
            ),
            (
                "transparent-gradient",
                "a:focus{outline:none;box-shadow:none;background-image:linear-gradient(rgba(0,85,255,0),rgba(0,85,255,0))}",
                False,
                False,
            ),
            (
                "autofocus-visible-outline",
                "a:focus{outline:4px solid #05f;box-shadow:none}",
                True,
                True,
            ),
            (
                "visible-outline",
                "a:focus{outline:4px solid #05f;box-shadow:none}",
                True,
                False,
            ),
            (
                "visible-shadow",
                "a:focus{outline:none;box-shadow:0 0 0 4px #05f}",
                True,
                False,
            ),
            (
                "visible-border",
                "a:focus{outline:none;box-shadow:none;border:4px solid #05f}",
                True,
                False,
            ),
            (
                "visible-background",
                "a:focus{outline:none;box-shadow:none;background-color:#8cf}",
                True,
                False,
            ),
        )
        for label, focus_css, expected, autofocus in cases:
            with self.subTest(case=label), tempfile.TemporaryDirectory() as raw:
                root = Path(raw)
                self._valid_page(
                    root, focus_css=focus_css, autofocus=autofocus
                )
                render = BrowserRenderer().render(
                    workspace_root=root,
                    output_root=root,
                    viewports=self._VIEWPORTS,
                    timeout_seconds=20,
                )
                observed = [
                    control["focus_visible"]
                    for viewport in render.viewports
                    for control in viewport.accessibility["controls"]
                ]
                self.assertEqual(observed, [expected, expected])
                if not expected:
                    self.assertFalse(
                        _runtime_gates(manifest, render)["accessibility"]["passed"]
                    )