#!/usr/bin/env python3
"""
查询各平台API额度
"""

import os
import json
import requests
from datetime import datetime


def check_aliyun_quota():
    """查询阿里云百炼额度"""
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        return {"error": "DASHSCOPE_API_KEY not set"}

    try:
        url = "https://dashscope.aliyuncs.com/compatible-mode/v1/usage"
        headers = {"Authorization": f"Bearer {api_key}"}
        resp = requests.get(url, headers=headers, timeout=10)
        if resp.status_code == 200:
            data = resp.json()
            return {"provider": "aliyun", "status": "ok", "data": data}
        else:
            return {"provider": "aliyun", "error": resp.text}
    except Exception as e:
        return {"provider": "aliyun", "error": str(e)}


def check_openrouter_quota():
    """查询OpenRouter额度"""
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        return {"error": "OPENROUTER_API_KEY not set"}

    try:
        url = "https://openrouter.ai/api/v1/credits"
        headers = {"Authorization": f"Bearer {api_key}"}
        resp = requests.get(url, headers=headers, timeout=10)
        if resp.status_code == 200:
            data = resp.json()
            return {"provider": "openrouter", "status": "ok", "data": data}
        else:
            return {"provider": "openrouter", "error": resp.text}
    except Exception as e:
        return {"provider": "openrouter", "error": str(e)}


def check_deepseek_quota():
    """查询DeepSeek额度"""
    api_key = os.getenv("DEEPSEEK_API_KEY")
    if not api_key:
        return {"error": "DEEPSEEK_API_KEY not set"}

    try:
        url = "https://api.deepseek.com/v1/remaining_quota"
        headers = {"Authorization": f"Bearer {api_key}"}
        resp = requests.get(url, headers=headers, timeout=10)
        if resp.status_code == 200:
            data = resp.json()
            return {"provider": "deepseek", "status": "ok", "data": data}
        else:
            return {"provider": "deepseek", "error": resp.text}
    except Exception as e:
        return {"provider": "deepseek", "error": str(e)}


def check_groq_quota():
    """查询Groq额度（通过OpenAI兼容接口）"""
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        return {"error": "GROQ_API_KEY not set"}

    try:
        # Groq没有直接的额度查询API，返回提示
        return {
            "provider": "groq",
            "status": "no_api",
            "message": "Groq无公开额度查询API，请登录Groq Dashboard查看",
        }
    except Exception as e:
        return {"provider": "groq", "error": str(e)}


def main():
    print(f"=== LLM额度查询 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ===\n")

    results = []

    # 阿里云
    print("[1/4] 阿里云百炼...")
    result = check_aliyun_quota()
    results.append(result)
    if "error" in result:
        print(f"    错误: {result.get('error')}")
    else:
        print(f"    OK")

    # OpenRouter
    print("[2/4] OpenRouter...")
    result = check_openrouter_quota()
    results.append(result)
    if "error" in result:
        print(f"    错误: {result.get('error')}")
    elif result.get("status") == "no_api":
        print(f"    {result.get('message')}")
    else:
        print(f"    OK")

    # DeepSeek
    print("[3/4] DeepSeek...")
    result = check_deepseek_quota()
    results.append(result)
    if "error" in result:
        print(f"    错误: {result.get('error')}")
    else:
        print(f"    OK")

    # Groq
    print("[4/4] Groq...")
    result = check_groq_quota()
    results.append(result)
    if "error" in result:
        print(f"    错误: {result.get('error')}")
    else:
        print(f"    {result.get('message', 'OK')}")

    print("\n=== 详细结果 ===")
    print(json.dumps(results, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()