#!/usr/bin/env python3
"""OpenAI 兼容性接口测试脚本

用法:
    python3 scripts/openai_compat_test.py --base_url <url> [options]

示例:
    python3 scripts/openai_compat_test.py --base_url https://api.example.com/v1
    python3 scripts/openai_compat_test.py --base_url https://api.example.com/v1 --api_key sk-xxx --model gpt-4o
    python3 scripts/openai_compat_test.py --base_url https://api.example.com/v1 --stream --tools
"""

import json
import time
import ssl
import argparse
import urllib.request
import urllib.error

TIMEOUT = 30


def create_ssl_context():
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    return ctx


def http_request(url, method="GET", headers=None, body=None, ssl_ctx=None):
    req = urllib.request.Request(url, method=method)
    if headers:
        for k, v in headers.items():
            req.add_header(k, v)
    if body is not None:
        if isinstance(body, str):
            req.data = body.encode("utf-8")
        else:
            req.data = json.dumps(body).encode("utf-8")

    start = time.time()
    try:
        resp = urllib.request.urlopen(req, timeout=TIMEOUT, context=ssl_ctx)
        elapsed = time.time() - start
        return resp.getcode(), resp.read().decode("utf-8"), elapsed
    except urllib.error.HTTPError as e:
        elapsed = time.time() - start
        return e.code, e.read().decode("utf-8"), elapsed
    except Exception as e:
        elapsed = time.time() - start
        return None, str(e), elapsed


def http_stream_request(url, headers=None, body=None, ssl_ctx=None):
    req = urllib.request.Request(url, method="POST")
    if headers:
        for k, v in headers.items():
            req.add_header(k, v)
    if body is not None:
        req.data = json.dumps(body).encode("utf-8")

    start = time.time()
    try:
        resp = urllib.request.urlopen(req, timeout=TIMEOUT, context=ssl_ctx)
        status = resp.getcode()
        lines = []
        for raw_line in resp:
            line = raw_line.decode("utf-8").rstrip("\n\r")
            if line:
                lines.append(line)
        elapsed = time.time() - start
        return status, "\n".join(lines), elapsed
    except urllib.error.HTTPError as e:
        elapsed = time.time() - start
        return e.code, e.read().decode("utf-8"), elapsed
    except Exception as e:
        elapsed = time.time() - start
        return None, str(e), elapsed


def format_json(text):
    try:
        parsed = json.loads(text)
        return json.dumps(parsed, ensure_ascii=False, indent=2)
    except (json.JSONDecodeError, TypeError):
        return text


def build_headers(api_key):
    h = {"Content-Type": "application/json"}
    if api_key:
        h["Authorization"] = f"Bearer {api_key}"
    return h


def run_test(index, total, desc, url, method, headers, body, stream, ssl_ctx):
    print(f"\n[{index}/{total}] {desc}")
    print(f">>> {method} {url}")
    if body is not None:
        if isinstance(body, str):
            print(body)
        else:
            print(format_json(json.dumps(body, ensure_ascii=False)))

    if stream:
        status, data, elapsed = http_stream_request(url, headers, body, ssl_ctx)
    else:
        status, data, elapsed = http_request(url, method, headers, body, ssl_ctx)

    if status is not None:
        print(f"状态码: {status} | 耗时: {elapsed:.2f}s")
    else:
        print(f"请求失败 | 耗时: {elapsed:.2f}s")

    if stream and status and status < 300:
        # 流式响应按 SSE 行逐行输出
        for line in data.split("\n"):
            print(line)
    else:
        print(format_json(data))

    return status


def main():
    parser = argparse.ArgumentParser(
        description="OpenAI 兼容性接口测试",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--base_url", required=True, help="API 基础地址 (如 https://api.example.com/v1)")
    parser.add_argument("--api_key", default="", help="API 密钥 (默认空)")
    parser.add_argument("--model", default="gpt-4o", help="模型名称 (默认 gpt-4o)")
    parser.add_argument("--vision", action="store_true", help="执行视觉相关测试")
    parser.add_argument("--stream", action="store_true", help="执行流式响应测试")
    parser.add_argument("--tools", action="store_true", help="执行工具调用测试")
    parser.add_argument("--logprobs", action="store_true", help="执行 logprobs 测试")
    parser.add_argument("--json_schema", action="store_true", help="执行 Structured Output 测试")
    parser.add_argument("--all", action="store_true", help="开启所有扩展测试")
    args = parser.parse_args()

    if args.all:
        args.vision = True
        args.stream = True
        args.tools = True
        args.logprobs = True
        args.json_schema = True

    base_url = args.base_url.rstrip("/")
    api_key = args.api_key
    model = args.model
    ssl_ctx = create_ssl_context()
    headers = build_headers(api_key)
    headers_bad_auth = build_headers("invalid-key-xxx")

    chat_url = f"{base_url}/chat/completions"

    # --- 收集用例: (描述, 方法, URL, 请求头, 请求体, 是否流式) ---
    cases = []

    # ---- Models API ----
    cases.append((
        "获取模型列表 (GET /models)",
        "GET", f"{base_url}/models", headers, None, False
    ))
    cases.append((
        "获取指定模型详情 (GET /models/{model})",
        "GET", f"{base_url}/models/{model}", headers, None, False
    ))
    cases.append((
        "获取不存在的模型 (GET /models/nonexistent-model-xxx)",
        "GET", f"{base_url}/models/nonexistent-model-xxx", headers, None, False
    ))

    # ---- Chat Completions: 正面用例 ----
    cases.append((
        "基本对话（仅 user）",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "system + user 对话",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "1+1="}
            ],
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "developer + user 对话",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [
                {"role": "developer", "content": "You are a helpful assistant."},
                {"role": "user", "content": "1+1="}
            ],
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "多轮对话（含 assistant 历史）",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [
                {"role": "user", "content": "Hi"},
                {"role": "assistant", "content": "Hello!"},
                {"role": "user", "content": "1+1="}
            ],
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "temperature + top_p",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 5,
            "temperature": 0.5,
            "top_p": 0.9
        }, False
    ))
    cases.append((
        "max_tokens 限制",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [{"role": "user", "content": "讲一个故事"}],
            "max_tokens": 10
        }, False
    ))
    cases.append((
        "stop sequences",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [{"role": "user", "content": "数数: 1,2,3,"}],
            "max_tokens": 20,
            "stop": ["5"]
        }, False
    ))
    cases.append((
        "n=2 多候选",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 5,
            "n": 2
        }, False
    ))
    cases.append((
        "seed 参数",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 5,
            "seed": 42
        }, False
    ))
    cases.append((
        "frequency_penalty + presence_penalty",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 5,
            "frequency_penalty": 0.5,
            "presence_penalty": 0.5
        }, False
    ))
    cases.append((
        "max_completion_tokens 参数",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [{"role": "user", "content": "讲一个故事"}],
            "max_completion_tokens": 10
        }, False
    ))
    cases.append((
        "JSON mode (response_format: json_object)",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [
                {"role": "system", "content": "以 JSON 格式回复: {\"answer\": \"ok\"}"},
                {"role": "user", "content": "test"}
            ],
            "max_tokens": 10,
            "response_format": {"type": "json_object"}
        }, False
    ))

    # ---- Chat Completions: 负面用例 ----
    cases.append((
        "缺少 model 参数",
        "POST", chat_url, headers, {
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "缺少 messages 参数",
        "POST", chat_url, headers, {
            "model": model,
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "messages 为空数组",
        "POST", chat_url, headers, {
            "model": model,
            "messages": [],
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "无效 API key",
        "POST", chat_url, headers_bad_auth, {
            "model": model,
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "不存在的模型 (chat)",
        "POST", chat_url, headers, {
            "model": "nonexistent-model-xxx",
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 5
        }, False
    ))
    cases.append((
        "畸形 JSON body",
        "POST", chat_url, headers, "invalid json{", False
    ))

    # ---- --vision ----
    if args.vision:
        image_url = (
            "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/"
            "Gfp-wisconsin-madison-the-nature-boardwalk.jpg/"
            "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
        )
        cases.append((
            "图片 URL 输入 + detail 参数 (--vision)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [
                    {"role": "system", "content": "简短描述图片"},
                    {"role": "user", "content": [
                        {"type": "text", "text": "用一个词描述这张图"},
                        {"type": "image_url", "image_url": {
                            "url": image_url, "detail": "low"
                        }}
                    ]}
                ],
                "max_tokens": 10
            }, False
        ))

    # ---- --stream ----
    if args.stream:
        cases.append((
            "基本流式 (--stream)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "Hi"}],
                "max_tokens": 5,
                "stream": True
            }, True
        ))
        cases.append((
            "流式 + include_usage (--stream)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "Hi"}],
                "max_tokens": 5,
                "stream": True,
                "stream_options": {"include_usage": True}
            }, True
        ))
        cases.append((
            "流式 + stop sequences (--stream)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "数数: 1,2,3,"}],
                "max_tokens": 20,
                "stream": True,
                "stop": ["5"]
            }, True
        ))

    # ---- --tools ----
    if args.tools:
        tool_weather = {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "获取指定城市的天气",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {"type": "string", "description": "城市名称"}
                    },
                    "required": ["location"]
                }
            }
        }
        cases.append((
            "工具调用 tool_choice: auto (--tools)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "北京天气怎么样？"}],
                "max_tokens": 50,
                "tools": [tool_weather],
                "tool_choice": "auto"
            }, False
        ))
        cases.append((
            "工具调用 tool_choice: required (--tools)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "北京天气怎么样？"}],
                "max_tokens": 50,
                "tools": [tool_weather],
                "tool_choice": "required"
            }, False
        ))
        cases.append((
            "指定函数调用 tool_choice: {name} (--tools)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "北京天气怎么样？"}],
                "max_tokens": 50,
                "tools": [tool_weather],
                "tool_choice": {
                    "type": "function",
                    "function": {"name": "get_weather"}
                }
            }, False
        ))
        cases.append((
            "多轮工具调用（构造 tool 结果）(--tools)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [
                    {"role": "user", "content": "北京天气怎么样？"},
                    {"role": "assistant", "content": None, "tool_calls": [{
                        "id": "call_001", "type": "function",
                        "function": {
                            "name": "get_weather",
                            "arguments": "{\"location\": \"Beijing\"}"
                        }
                    }]},
                    {"role": "tool", "tool_call_id": "call_001",
                     "content": "{\"temperature\": 22, \"condition\": \"晴\"}"}
                ],
                "max_tokens": 20,
                "tools": [tool_weather]
            }, False
        ))
        cases.append((
            "parallel_tool_calls: false (--tools)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "北京和上海的天气怎么样？"}],
                "max_tokens": 50,
                "tools": [tool_weather],
                "tool_choice": "auto",
                "parallel_tool_calls": False
            }, False
        ))

    # ---- --logprobs ----
    if args.logprobs:
        cases.append((
            "logprobs + top_logprobs (--logprobs)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "Hi"}],
                "max_tokens": 5,
                "logprobs": True,
                "top_logprobs": 2
            }, False
        ))

    # ---- --json-schema ----
    if args.json_schema:
        cases.append((
            "Structured Output json_schema (--json_schema)",
            "POST", chat_url, headers, {
                "model": model,
                "messages": [{"role": "user", "content": "1+1等于几？"}],
                "max_tokens": 20,
                "response_format": {
                    "type": "json_schema",
                    "json_schema": {
                        "name": "math_answer",
                        "strict": True,
                        "schema": {
                            "type": "object",
                            "properties": {
                                "answer": {"type": "number"},
                                "explanation": {"type": "string"}
                            },
                            "required": ["answer", "explanation"],
                            "additionalProperties": False
                        }
                    }
                }
            }, False
        ))

    # ---- 执行测试 ----
    total = len(cases)
    count_2xx = 0
    count_other = 0

    print("=" * 60)
    print("OpenAI 兼容性测试")
    print(f"目标: {base_url}")
    print(f"模型: {model}")
    print(f"时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    flags = []
    if args.vision:
        flags.append("vision")
    if args.stream:
        flags.append("stream")
    if args.tools:
        flags.append("tools")
    if args.logprobs:
        flags.append("logprobs")
    if args.json_schema:
        flags.append("json-schema")
    print(f"用例: {total} 个" + (f" | 扩展: {', '.join(flags)}" if flags else ""))
    print("=" * 60)

    for i, (desc, method, url, hdrs, body, stream) in enumerate(cases, 1):
        status = run_test(i, total, desc, url, method, hdrs, body, stream, ssl_ctx)
        if status is not None and 200 <= status < 300:
            count_2xx += 1
        else:
            count_other += 1

    print()
    print("=" * 60)
    print(f"测试完成 | 总计: {total} | HTTP 2xx: {count_2xx} | 非 2xx: {count_other}")
    print("=" * 60)


if __name__ == "__main__":
    main()