import json
import os
import tempfile
import threading
import time
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer


class TmpCase(unittest.TestCase):
    """A test case with a fresh temporary directory in self.dir."""

    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.dir = tmp.name

    def write(self, name, text):
        path = os.path.join(self.dir, name)
        with open(path, "w") as f:
            f.write(text)
        return path


def wait_for(predicate, timeout=5.0):
    end = time.monotonic() + timeout
    while not predicate():
        if time.monotonic() > end:
            raise AssertionError("timed out waiting")
        time.sleep(0.01)


class StubVllm:
    """A fake OpenAI-compatible server on a free port. Records every POST body. Streamed
    chats answer "Hel", "lo", "!"; others answer "ok". Set .delay (seconds, or a function of
    the request body) to make answers slow, .pause (the same) to stream the pieces that far
    apart, and .fail to an HTTP code to reject chats."""

    def __init__(self):
        self.bodies = []
        self.delay = 0.0
        self.pause = 0.0
        self.fail = None
        self.metrics = None  # Prometheus text served at GET /metrics
        stub = self

        class Handler(BaseHTTPRequestHandler):
            def do_GET(self):
                if self.path == "/metrics" and stub.metrics is not None:
                    return self.send(200, "text/plain", stub.metrics)
                self.send(200, "application/json", {"data": [{"id": "stub-model"}]})

            def do_POST(self):
                body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
                stub.bodies.append(body)
                time.sleep(stub.delay(body) if callable(stub.delay) else stub.delay)
                if stub.fail:
                    return self.send(stub.fail, "application/json", {"error": "rejected"})
                if not body.get("stream"):
                    return self.send(200, "application/json",
                                     {"choices": [{"message": {"content": "ok"}}]})
                events = [{"choices": [{"delta": {"content": p}}]} for p in ("Hel", "lo", "!")]
                pause = stub.pause(body) if callable(stub.pause) else stub.pause
                if pause:  # stream piece by piece; HTTP/1.0 ends the body by closing
                    return self.trickle([f"data: {json.dumps(e)}\n\n" for e in events], pause)
                text = "".join(f"data: {json.dumps(e)}\n\n" for e in events) + "data: [DONE]\n\n"
                self.send(200, "text/event-stream", text)

            def trickle(self, chunks, pause):
                try:
                    self.send_response(200)
                    self.send_header("Content-Type", "text/event-stream")
                    self.end_headers()
                    for i, chunk in enumerate(chunks + ["data: [DONE]\n\n"]):
                        time.sleep(pause if i else 0)
                        self.wfile.write(chunk.encode())
                        self.wfile.flush()
                except (BrokenPipeError, ConnectionResetError):
                    pass  # the client aborted

            def send(self, code, kind, payload):
                data = (payload if isinstance(payload, str) else json.dumps(payload)).encode()
                try:
                    self.send_response(code)
                    self.send_header("Content-Type", kind)
                    self.send_header("Content-Length", str(len(data)))
                    self.end_headers()
                    self.wfile.write(data)
                except (BrokenPipeError, ConnectionResetError):
                    pass  # the client aborted

            def log_message(self, *args):
                pass

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

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