import json
import os
import sys
import requests

BASE_URL = 'https://api.deepseek.com'


# ============================================================
# 内部辅助
# ============================================================

def resolve_api_key(arguments):
    """解析 API Key，优先级：参数 > DEEPSEEK_API_KEY 环境变量"""
    if arguments.get('api_key'):
        return arguments['api_key']

    key = os.environ.get('DEEPSEEK_API_KEY', '')
    if key:
        return key

    return None


# ============================================================
# 工具函数
# ============================================================

def get_user_balance(api_key=None):
    """查询 DeepSeek 账户余额"""
    key = resolve_api_key({'api_key': api_key} if api_key else {})
    if not key:
        raise Exception(
            '未找到 API Key，请通过以下任一方式提供：\n'
            '1. 调用工具时传入 api_key 参数\n'
            '2. 设置环境变量 DEEPSEEK_API_KEY'
        )

    headers = {
        'Authorization': f'Bearer {key}',
        'Accept': 'application/json',
    }
    resp = requests.get(f'{BASE_URL}/user/balance', headers=headers, timeout=15)
    if resp.status_code != 200:
        raise Exception(f'API 请求失败 (HTTP {resp.status_code}): {resp.text}')
    return resp.json()


# ============================================================
# MCP JSON-RPC 处理
# ============================================================

TOOLS = [
    {
        "name": "get_user_balance",
        "description": "查询 DeepSeek API 账户余额，返回各币种的可用余额、赠送余额和充值余额。支持从参数或环境变量 DEEPSEEK_API_KEY 读取 API Key。",
        "inputSchema": {
            "type": "object",
            "properties": {
                "api_key": {
                    "type": "string",
                    "description": "DeepSeek API Key（可选，不传则从环境变量 DEEPSEEK_API_KEY 读取）"
                }
            },
            "required": []
        }
    }
]

TOOL_MAP = {
    "get_user_balance": get_user_balance
}


def handle_request(req):
    method = req.get("method", "")
    req_id = req.get("id")

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": "ds-balance-inquiry",
                    "version": "1.0.0"
                }
            }
        }

    if method == "ping":
        return {"jsonrpc": "2.0", "id": req_id, "result": {}}

    if method == "notifications/initialized":
        return None

    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}}

    if method == "tools/call":
        tool_name = req["params"]["name"]
        arguments = req["params"].get("arguments", {})
        tool_func = TOOL_MAP.get(tool_name)
        if not tool_func:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32601, "message": f"未知工具: {tool_name}"}}
        try:
            result = tool_func(**arguments)
            text = json.dumps(result, ensure_ascii=False, indent=2)
            return {"jsonrpc": "2.0", "id": req_id,
                    "result": {"content": [{"type": "text", "text": text}]}}
        except Exception as e:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32000, "message": str(e)}}

    return {"jsonrpc": "2.0", "id": req_id,
            "error": {"code": -32601, "message": f"未知方法: {method}"}}


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            resp = handle_request(req)
            if resp is None:
                continue
            sys.stdout.write(json.dumps(resp, ensure_ascii=False) + '\n')
            sys.stdout.flush()
        except json.JSONDecodeError:
            pass


if __name__ == '__main__':
    main()