#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业13 · 模拟 OpenAI 兼容上游（可限流出 429 / 可注入 5xx / 超时）

用法：
  # primary：低并发上限，方便压测打到 429
  python hw13_01_模拟上游.py --port 18101 --name primary --max-concurrent 4

  # backup：更高上限（或不同行为）
  python hw13_01_模拟上游.py --port 18102 --name backup --max-concurrent 20

  # 演示 5xx：请求带 Header X-Mock-Fail: 500
  # 演示超时：Header X-Mock-Sleep: 35（秒）
"""
from __future__ import annotations

import argparse
import json
import threading
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse


class State:
    def __init__(self, name: str, max_concurrent: int, latency_ms: int):
        self.name = name
        self.max_concurrent = max_concurrent
        self.latency_ms = latency_ms
        self.lock = threading.Lock()
        self.in_flight = 0
        self.total_ok = 0
        self.total_429 = 0
        self.total_5xx = 0


def make_handler(state: State):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self, fmt: str, *args) -> None:
            print(f"[{state.name}] {self.address_string()} - {fmt % args}")

        def _json(self, code: int, obj: dict) -> None:
            raw = json.dumps(obj, ensure_ascii=False).encode("utf-8")
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(raw)))
            self.end_headers()
            self.wfile.write(raw)

        def do_GET(self) -> None:
            path = urlparse(self.path).path
            if path in ("/health", "/v1/health"):
                self._json(
                    200,
                    {
                        "status": "ok",
                        "name": state.name,
                        "in_flight": state.in_flight,
                        "max_concurrent": state.max_concurrent,
                        "total_ok": state.total_ok,
                        "total_429": state.total_429,
                    },
                )
                return
            if path == "/v1/models":
                self._json(
                    200,
                    {
                        "object": "list",
                        "data": [
                            {
                                "id": f"mock-{state.name}",
                                "object": "model",
                                "owned_by": "homework13",
                            }
                        ],
                    },
                )
                return
            self._json(404, {"error": {"message": f"not found: {path}", "type": "not_found"}})

        def do_POST(self) -> None:
            path = urlparse(self.path).path
            if path not in ("/v1/chat/completions", "/chat/completions"):
                self._json(404, {"error": {"message": f"not found: {path}", "type": "not_found"}})
                return

            length = int(self.headers.get("Content-Length") or 0)
            _ = self.rfile.read(length) if length else b"{}"

            # 注入故障（学习 Runbook 用）
            fail = (self.headers.get("X-Mock-Fail") or "").strip()
            sleep_s = float(self.headers.get("X-Mock-Sleep") or 0)

            acquired = False
            with state.lock:
                if state.in_flight >= state.max_concurrent:
                    state.total_429 += 1
                    self._json(
                        429,
                        {
                            "error": {
                                "message": (
                                    f"mock upstream '{state.name}' overloaded: "
                                    f"in_flight={state.in_flight} >= max={state.max_concurrent}"
                                ),
                                "type": "rate_limit_error",
                                "code": "429",
                            }
                        },
                    )
                    return
                state.in_flight += 1
                acquired = True

            try:
                if sleep_s > 0:
                    time.sleep(sleep_s)
                elif state.latency_ms > 0:
                    time.sleep(state.latency_ms / 1000.0)

                if fail in ("500", "502", "503"):
                    with state.lock:
                        state.total_5xx += 1
                    self._json(
                        int(fail),
                        {
                            "error": {
                                "message": f"mock injected {fail} on '{state.name}'",
                                "type": "server_error",
                                "code": fail,
                            }
                        },
                    )
                    return

                cid = f"chatcmpl-hw13-{uuid.uuid4().hex[:12]}"
                with state.lock:
                    state.total_ok += 1
                self._json(
                    200,
                    {
                        "id": cid,
                        "object": "chat.completion",
                        "created": int(time.time()),
                        "model": f"mock-{state.name}",
                        "choices": [
                            {
                                "index": 0,
                                "finish_reason": "stop",
                                "message": {
                                    "role": "assistant",
                                    "content": f"ok from {state.name}",
                                },
                            }
                        ],
                        "usage": {"prompt_tokens": 8, "completion_tokens": 4, "total_tokens": 12},
                    },
                )
            finally:
                if acquired:
                    with state.lock:
                        state.in_flight -= 1

    return Handler


def main() -> int:
    ap = argparse.ArgumentParser(description="作业13 模拟上游")
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, required=True)
    ap.add_argument("--name", default="primary", help="primary / backup 等标签")
    ap.add_argument("--max-concurrent", type=int, default=4, help="超过则返回 429")
    ap.add_argument("--latency-ms", type=int, default=80, help="成功请求人为延迟")
    args = ap.parse_args()

    state = State(args.name, args.max_concurrent, args.latency_ms)
    server = ThreadingHTTPServer((args.host, args.port), make_handler(state))
    print(
        f"[OK] mock upstream name={args.name} http://127.0.0.1:{args.port}/v1 "
        f"max_concurrent={args.max_concurrent} latency_ms={args.latency_ms}"
    )
    print("[INFO] POST /v1/chat/completions  | GET /health | GET /v1/models")
    print("[INFO] Header X-Mock-Fail:500 注入5xx；X-Mock-Sleep:35 注入超时")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n[STOP]")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())