#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""AIGU Data API · MCP Server（零依赖，纯 Python 标准库）

让支持 MCP（Model Context Protocol）的 AI 软件直接查询 A 股全量数据：
Claude Desktop / Claude Code / Cursor / Cline / Windsurf / ZCode 等。

配置（客户端 mcpServers 配置段）：
  {
    "aigu": {
      "command": "python3",
      "args": ["/path/to/aigu_mcp.py"],
      "env": { "AIGU_KEY": "ak-你的Key" }
    }
  }

Key 获取：https://aiguq.com/console.html 注册即得（体验版每天5次，标准版100次/天）。
协议：stdin/stdout 上的 newline 分隔 JSON-RPC 2.0；日志只写 stderr。
"""
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request

BASE_URL = os.environ.get("AIGU_BASE", "https://aiguq.com")
API_KEY = os.environ.get("AIGU_KEY", "")
PROTOCOL_VERSION = "2024-11-05"

TOOLS = [
    ("get_quote", "获取单只A股实时快照：最新价/涨跌幅/五档/量额", {
        "type": "object", "required": ["code"],
        "properties": {"code": {"type": "string", "description": "证券代码，如 600519 或 600519.SH"}}}),
    ("get_batch", "批量实时快照，最多50只，计1次调用", {
        "type": "object", "required": ["codes"],
        "properties": {"codes": {"type": "string", "description": "逗号分隔代码，如 600519,000001"}}}),
    ("get_kline", "K线序列：8周期(1m~1M)×5种复权×任意区间", {
        "type": "object", "required": ["code"],
        "properties": {
            "code": {"type": "string"},
            "period": {"type": "string", "enum": ["1m", "5m", "15m", "30m", "60m", "1d", "1w", "1M"], "description": "默认1d"},
            "count": {"type": "integer", "description": "根数，默认120最大1000"},
            "dividend_type": {"type": "string", "enum": ["none", "front", "back", "front_ratio", "back_ratio"], "description": "复权，默认none；回测推荐back"},
            "start_time": {"type": "string", "description": "YYYYMMDD"},
            "end_time": {"type": "string", "description": "YYYYMMDD"}}}),
    ("get_tick_history", "历史分笔回放（单标的单日，计5次调用）", {
        "type": "object", "required": ["code", "date"],
        "properties": {"code": {"type": "string"}, "date": {"type": "string", "description": "YYYYMMDD"},
                       "count": {"type": "integer", "description": "笔数，默认1000最大5000"}}}),
    ("get_rank", "全市场涨幅榜 Top100", {"type": "object", "properties": {}}),
    ("get_zt_pool", "当日涨停池", {"type": "object", "properties": {}}),
    ("get_index", "五大指数：上证/深成/创业板/沪深300/科创50", {"type": "object", "properties": {}}),
    ("get_stock_list", "全市场证券列表", {
        "type": "object",
        "properties": {"market": {"type": "string", "enum": ["SH", "SZ", "BJ", "ALL"]},
                       "keyword": {"type": "string", "description": "代码包含过滤，如 688"}}}),
    ("get_instrument", "合约信息：涨跌停价/股本/上市日/最小变动价位", {
        "type": "object", "required": ["code"], "properties": {"code": {"type": "string"}}}),
    ("get_sector_list", "板块分类列表（行业/概念 7000+，申万行业含SW前缀）", {
        "type": "object",
        "properties": {"keyword": {"type": "string", "description": "如 银行"}, "limit": {"type": "integer"}}}),
    ("get_sector_members", "板块成分股（板块名来自 get_sector_list）", {
        "type": "object", "required": ["name"], "properties": {"name": {"type": "string"}}}),
    ("get_index_weight", "指数成分权重（降序，如 000300.SH 沪深300）", {
        "type": "object",
        "properties": {"index": {"type": "string", "description": "默认000300.SH；常用 000905.SH中证500 000852.SH中证1000"}}}),
    ("get_financial", "财务报表（8张表，计2次调用）：Balance/Income/CashFlow/Pershareindex/Capital/Holdernum/Top10holder/Top10flowholder", {
        "type": "object", "required": ["code"],
        "properties": {"code": {"type": "string"}, "table": {"type": "string", "description": "默认 Income"},
                       "count": {"type": "integer", "description": "期数1-12，默认4"}}}),
    ("get_dividend", "除权除息全部历史与dr复权因子", {
        "type": "object", "required": ["code"], "properties": {"code": {"type": "string"}}}),
    ("get_trading_dates", "交易日历与节假日", {
        "type": "object",
        "properties": {"market": {"type": "string"},
                       "start_time": {"type": "string", "description": "YYYYMMDD"},
                       "end_time": {"type": "string", "description": "YYYYMMDD"}}}),
    ("get_cb_info", "可转债要素：票息/转股/规模/期限", {
        "type": "object", "required": ["code"], "properties": {"code": {"type": "string"}}}),
    ("get_ipo", "新股/新转债申购信息", {
        "type": "object",
        "properties": {"start_time": {"type": "string"}, "end_time": {"type": "string"}}}),
    ("get_etf_info", "ETF申赎清单（不带code=全部1700+只清单）", {
        "type": "object", "properties": {"code": {"type": "string"}}}),
    ("get_usage", "查询当前Key今日用量与配额", {"type": "object", "properties": {}}),
]

# 工具名 -> (路径, 参数映射)
_ROUTES = {
    "get_quote":         ("/v1/quote", lambda a: {"code": a["code"]}),
    "get_batch":         ("/v1/quote/batch", lambda a: {"codes": a["codes"]}),
    "get_kline":         ("/v1/kline", lambda a: {k: str(v) for k, v in {
        "code": a["code"], "period": a.get("period", "1d"),
        "count": a.get("count", 120), "dividend_type": a.get("dividend_type", "none"),
        "start_time": a.get("start_time", ""), "end_time": a.get("end_time", "")}.items() if v != ""}),
    "get_tick_history":  ("/v1/tick_history", lambda a: {"code": a["code"], "date": a["date"],
                                                         "count": a.get("count", 1000)}),
    "get_rank":          ("/v1/rank", lambda a: {}),
    "get_zt_pool":       ("/v1/zt_pool", lambda a: {}),
    "get_index":         ("/v1/index", lambda a: {}),
    "get_stock_list":    ("/v1/stock_list", lambda a: {k: v for k, v in {
        "market": a.get("market", "ALL"), "keyword": a.get("keyword", "")}.items() if v}),
    "get_instrument":    ("/v1/instrument", lambda a: {"code": a["code"]}),
    "get_sector_list":   ("/v1/sector_list", lambda a: {k: str(v) for k, v in {
        "keyword": a.get("keyword", ""), "limit": a.get("limit", 50)}.items() if v}),
    "get_sector_members": ("/v1/sector", lambda a: {"name": a["name"], "limit": a.get("limit", 1000)}),
    "get_index_weight":  ("/v1/index_weight", lambda a: {"index": a.get("index", "000300.SH")}),
    "get_financial":     ("/v1/financial", lambda a: {k: str(v) for k, v in {
        "code": a["code"], "table": a.get("table", "Income"), "count": a.get("count", 4)}.items()}),
    "get_dividend":      ("/v1/dividend", lambda a: {"code": a["code"]}),
    "get_trading_dates": ("/v1/trading_dates", lambda a: {k: v for k, v in {
        "market": a.get("market", "SH"), "start_time": a.get("start_time", ""),
        "end_time": a.get("end_time", "")}.items() if v}),
    "get_cb_info":       ("/v1/cb_info", lambda a: {"code": a["code"]}),
    "get_ipo":           ("/v1/ipo", lambda a: {k: v for k, v in {
        "start_time": a.get("start_time", ""), "end_time": a.get("end_time", "")}.items() if v}),
    "get_etf_info":      ("/v1/etf_info", lambda a: {"code": a["code"]} if a.get("code") else {"limit": 300}),
    "get_usage":         ("/v1/usage", lambda a: {}),
}


def api_get(path, params):
    qs = ("?" + urllib.parse.urlencode(params)) if params else ""
    req = urllib.request.Request(BASE_URL + path + qs, headers={
        "X-API-Key": API_KEY, "User-Agent": "aigu-mcp/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            return 0, json.loads(r.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        reason = e.headers.get("x-rate-reason", "") if e.headers else ""
        try:
            msg = json.loads(e.read().decode("utf-8")).get("detail", "")
        except Exception:
            msg = ""
        return e.code, {"detail": msg or ("HTTP %d" % e.code), "reason": reason}
    except Exception as e:
        return -1, {"detail": str(e)}


def tools_schema():
    return [{"name": n, "description": d, "inputSchema": s} for n, d, s in TOOLS]


def call_tool(name, args):
    if name not in _ROUTES:
        return {"isError": True, "content": [{"type": "text", "text": f"未知工具 {name}"}]}
    path, build = _ROUTES[name]
    try:
        params = build(args or {})
    except Exception as e:
        return {"isError": True, "content": [{"type": "text", "text": f"参数错误: {e}"}]}
    if not API_KEY:
        return {"isError": True, "content": [{"type": "text",
                "text": "未配置 AIGU_KEY 环境变量。到 https://aiguq.com/console.html 注册获取 API Key。"}]}
    code, body = api_get(path, params)
    if code != 0 and "data" not in body:
        text = json.dumps(body, ensure_ascii=False)
        if code == 401:
            text = "API Key 无效或已停用（401）。请检查 AIGU_KEY。"
        elif code == 429:
            text = "超过配额（429）。体验版每天5次；可用 get_usage 查询，或到 https://aiguq.com 升级套餐。"
        return {"isError": True, "content": [{"type": "text", "text": text}]}
    data = body.get("data", body)
    return {"content": [{"type": "text",
                         "text": json.dumps(data, ensure_ascii=False, default=str)}]}


def reply(rid, result):
    sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": rid, "result": result},
                                ensure_ascii=False, default=str) + "\n")
    sys.stdout.flush()


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except Exception:
            continue
        method = msg.get("method", "")
        rid = msg.get("id")
        if method == "initialize":
            reply(rid, {
                "protocolVersion": PROTOCOL_VERSION,
                "capabilities": {"tools": {}},
                "serverInfo": {"name": "aigu-data-api", "version": "1.0.0"},
            })
        elif method.startswith("notifications/"):
            continue   # 通知不回包
        elif method == "ping":
            reply(rid, {})
        elif method == "tools/list":
            reply(rid, {"tools": tools_schema()})
        elif method == "tools/call":
            params = msg.get("params") or {}
            reply(rid, call_tool(params.get("name", ""), params.get("arguments") or {}))
        elif rid is not None:
            reply(rid, {"error": {"code": -32601, "message": "method not found: " + method}})


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass