#!/usr/bin/env python3
"""Force a real Pi HTTP provider response to prove response-hook timing."""

from __future__ import annotations

import json
import os
import tempfile
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

from pi_gate0_run import PiRpc, jsonl

HERE = Path(__file__).resolve().parent


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, _format: str, *_args: object) -> None:
        return

    def do_POST(self) -> None:  # noqa: N802
        length = int(self.headers.get("content-length", "0"))
        self.rfile.read(length)
        chunks = [
            {
                "id": "gate0-response",
                "object": "chat.completion.chunk",
                "created": 1,
                "model": "gate0-model",
                "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
            },
            {
                "id": "gate0-response",
                "object": "chat.completion.chunk",
                "created": 1,
                "model": "gate0-model",
                "choices": [
                    {"index": 0, "delta": {"content": "TIMING_OK"}, "finish_reason": None}
                ],
            },
            {
                "id": "gate0-response",
                "object": "chat.completion.chunk",
                "created": 1,
                "model": "gate0-model",
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
                "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
            },
        ]
        body = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
        encoded = body.encode()
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Content-Length", str(len(encoded)))
        self.send_header("X-Gate0-Response", "headers-before-stream")
        self.end_headers()
        self.wfile.write(encoded)
        self.wfile.flush()


def main() -> None:
    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    port = server.server_address[1]

    with tempfile.TemporaryDirectory(prefix="gate0-p2-timing-") as temp:
        root = Path(temp)
        workspace = root / "workspace"
        workspace.mkdir()
        log = root / "hooks.jsonl"
        env = os.environ.copy()
        env.update(
            {
                "GATE0_PI_LOG": str(log),
                "GATE0_LOCAL_PROVIDER_URL": f"http://127.0.0.1:{port}/v1",
                "MOSAIC_PI_FORCE_SKILLS": "",
                "PI_SKIP_VERSION_CHECK": "1",
            }
        )
        command = [
            "mosaic",
            "yolo",
            "pi",
            "--mode",
            "rpc",
            "--no-session",
            "--no-extensions",
            "--no-context-files",
            "--no-prompt-templates",
            "--provider",
            "gate0-local",
            "--model",
            "gate0-model",
            "--extension",
            str(HERE / "pi_gate0_extension.ts"),
        ]
        pi = PiRpc(command, workspace, env)
        try:
            pi.prompt_and_settle("timing", "Reply with TIMING_OK")
            records = jsonl(log)
            selected = [
                record
                for record in records
                if record["event"] in {"before_provider_request", "after_provider_response", "message_end"}
                and (record["event"] != "message_end" or record.get("role") == "assistant")
            ]
            print("$ python3 docs/compaction-refresh/probes/p2_provider_timing_run.py")
            print(f"local_http_endpoint=http://127.0.0.1:{port}/v1/chat/completions")
            for record in selected:
                print(json.dumps(record, sort_keys=True))
            after = next(record for record in selected if record["event"] == "after_provider_response")
            message = next(record for record in selected if record["event"] == "message_end")
            if not (
                after["seq"] < message["seq"]
                and after["assistantContentAvailableAtThisHook"] is False
                and message["assistantContentObserved"] is True
            ):
                raise AssertionError("provider response/content observation ordering failed")
            print("machine_assertions=PASS")
            print(f"after_provider_response_seq={after['seq']}")
            print(f"message_end_seq={message['seq']}")
            print(f"headers_hook_precedes_completed_message={after['seq'] < message['seq']}")
        finally:
            pi.close()

    server.shutdown()
    server.server_close()


if __name__ == "__main__":
    main()