import argparse
import datetime
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request

TASKS = {
    "support": '只返回 JSON，不添加 Markdown 或解释。把以下工单分类为 billing、authentication 或 invoice。T01：已经扣费但余额没有增加。T02：API 请求返回 401，提示认证失败。T03：想了解开具发票的流程。返回格式：{"tickets":[{"id":"T01","category":"..."},{"id":"T02","category":"..."},{"id":"T03","category":"..."}]}。每个工单只出现一次。',
    "rag": '只依据下列资料回答，不添加外部知识。DOC-A：付款后 7 天内且未发货可以申请取消；已发货需走退货流程。DOC-B：订单付款 2 天，当前状态为待发货。问题：该订单是否满足申请取消的条件？只返回 JSON，格式：{"eligible":true或false,"days":规则规定的天数,"source_ids":["来源ID"]}，引用支持结论的全部资料，不添加 Markdown 或解释。',
    "invoice": '测试订单 INV-DEMO-001：商品小计 42.50 元，运费 7.50 元，优惠 10.00 元，币种人民币。计算实付金额（商品小计 + 运费 - 优惠）。只返回 JSON，格式：{"invoice_id":"...","currency":"CNY","total":数字}，不要 Markdown 或解释。',
}
ENDPOINTS = {"chat": "/chat/completions", "messages": "/messages", "responses": "/responses"}


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, request, response, code, message, headers, new_url):
        return None


def create_request(task, protocol, model):
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._/-]{0,99}", model):
        raise ValueError("请提供有效且准确的模型 ID。")
    prompt = TASKS[task]
    if protocol == "responses":
        body = {"model": model, "input": prompt, "max_output_tokens": 1024, "stream": False}
    else:
        body = {"model": model, "messages": [{"role": "user", "content": prompt}], "max_tokens": 1024, "stream": False}
    return "https://api.quickrouter.ai/v1" + ENDPOINTS[protocol], body


def redact(text, key):
    return re.sub(r"sk-[A-Za-z0-9_-]+", "[REDACTED]", text.replace(key, "[REDACTED]") if key else text)


def execute(task, protocol, model, key):
    if not key.strip():
        raise ValueError("请先在本地安全设置 QUICKROUTER_API_KEY。")
    url, body = create_request(task, protocol, model)
    headers = {"Content-Type": "application/json", "Authorization": "Bearer " + key}
    if protocol == "messages":
        headers["anthropic-version"] = "2023-06-01"
    request = urllib.request.Request(url, data=json.dumps(body).encode("utf-8"), headers=headers, method="POST")
    started = time.perf_counter()
    try:
        with urllib.request.build_opener(NoRedirect()).open(request, timeout=120) as response:
            raw = response.read(1_000_001)
            status = response.status
    except urllib.error.HTTPError as error:
        raise ValueError("HTTP " + str(error.code) + "：请到对应报错教程核对认证、协议、型号或额度。") from None
    except (urllib.error.URLError, TimeoutError):
        raise ValueError("请求失败或超时。请先核对控制台用量，再决定是否重试。") from None
    if len(raw) > 1_000_000:
        raise ValueError("响应超出短任务大小上限，请在控制台检查请求。")
    try:
        parsed = json.loads(redact(raw.decode("utf-8"), key))
    except (UnicodeDecodeError, json.JSONDecodeError):
        raise ValueError("响应不是有效 JSON。") from None
    if not isinstance(parsed, dict) or parsed.get("error"):
        raise ValueError("API 未返回有效成功响应，请核对控制台记录。")
    return {
        "task": task, "prompt_version": "2026-10-06-v1", "protocol": protocol,
        "requested_model": model, "observed_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "http_status": status, "elapsed_ms": round((time.perf_counter() - started) * 1000),
        "billed_usd": None, "response": parsed,
    }


def main():
    parser = argparse.ArgumentParser(description="QuickRouter 固定任务测试。默认仅预览；--run 发送一次可能收费的请求。")
    parser.add_argument("--task", choices=TASKS, default="support")
    parser.add_argument("--protocol", choices=ENDPOINTS, default="chat")
    parser.add_argument("--model", required=True)
    parser.add_argument("--run", action="store_true", help="执行一次请求，不自动重试")
    args = parser.parse_args()
    try:
        url, body = create_request(args.task, args.protocol, args.model)
        result = execute(args.task, args.protocol, args.model, os.environ.get("QUICKROUTER_API_KEY", "")) if args.run else {"mode": "preview", "url": url, "body": body, "network_request_sent": False}
        print(json.dumps(result, ensure_ascii=False, indent=2))
    except ValueError as error:
        print(str(error), file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())