#!/usr/bin/env python3
"""Tiny streaming inference API simulator with TTFT and total-latency measurement."""
from __future__ import annotations
import argparse
import json
import threading
import time
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

HOST = "127.0.0.1"
PORT = 18081
TOKENS = ["Streaming", " separates", " first-output", " latency", " from", " total", " completion."]


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.0"
    server_version = "GNU-Learning-Inference/1"

    def log_message(self, fmt, *args):
        return

    def do_POST(self):
        if self.path != "/v1/stream":
            self._json_error(404, "not_found")
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
            body = json.loads(self.rfile.read(length) or b"{}")
        except (ValueError, json.JSONDecodeError):
            self._json_error(400, "invalid_json")
            return
        prompt = str(body.get("prompt", "")).strip()
        if not prompt:
            self._json_error(400, "prompt_required")
            return
        try:
            queue_ms = max(0, min(int(body.get("queue_ms", 700)), 5000))
        except (TypeError, ValueError):
            self._json_error(400, "queue_ms_must_be_integer")
            return
        self.send_response(200)
        self.send_header("Content-Type", "application/x-ndjson")
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        time.sleep(queue_ms / 1000)
        for index, token in enumerate(TOKENS):
            event = {"type": "token", "index": index, "text": token}
            self.wfile.write((json.dumps(event) + "\n").encode())
            self.wfile.flush()
            time.sleep(0.06)
        self.wfile.write((json.dumps({"type": "done", "tokens": len(TOKENS)}) + "\n").encode())
        self.wfile.flush()

    def _json_error(self, status: int, error: str):
        data = json.dumps({"error": error}).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)


def start_server(port: int):
    server = ThreadingHTTPServer((HOST, port), Handler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return server, thread


def measure(port: int, queue_ms: int):
    payload = json.dumps({"prompt": "explain batching", "queue_ms": queue_ms}).encode()
    request = urllib.request.Request(f"http://{HOST}:{port}/v1/stream", data=payload, headers={"Content-Type": "application/json"}, method="POST")
    started = time.perf_counter()
    first = None
    done = False
    events = 0
    with urllib.request.urlopen(request, timeout=10) as response:
        for raw in response:
            event = json.loads(raw)
            if event.get("type") == "token" and first is None:
                first = time.perf_counter()
            if event.get("type") == "done":
                done = True
            events += 1
    finished = time.perf_counter()
    assert first is not None and done
    return (first - started) * 1000, (finished - started) * 1000, events


def self_test() -> int:
    server, thread = start_server(0)
    port = server.server_address[1]
    try:
        ttft_fast, total_fast, events_fast = measure(port, 120)
        ttft_slow, total_slow, events_slow = measure(port, 700)
        assert events_fast == len(TOKENS) + 1 == events_slow
        assert ttft_slow > ttft_fast + 350
        assert total_slow > total_fast + 350
        print(f"PASS stream_events={events_fast} done_event=true")
        print(f"PASS fast_ttft_ms={ttft_fast:.0f} fast_total_ms={total_fast:.0f}")
        print(f"PASS slow_ttft_ms={ttft_slow:.0f} slow_total_ms={total_slow:.0f}")
        print("PASS queue_delay_increases_ttft_without_changing_token_interval")
    finally:
        server.shutdown(); server.server_close(); thread.join(timeout=2)
    return 0


def serve() -> int:
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    print(f"Teaching streaming API listening on http://{HOST}:{PORT}/v1/stream")
    print("Use curl -N so each NDJSON event is shown as it arrives. Ctrl-C to stop.")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopping cleanly.")
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--self-test", action="store_true")
    group.add_argument("--serve", action="store_true")
    args = parser.parse_args()
    raise SystemExit(self_test() if args.self_test else serve())