"""market_sync MCP Server。

提供 5 类工具（tool），让 LLM 能查 market_sync 的同步情况：

1. list_sync_tasks              — 列出所有同步任务
2. get_task_status              — 查单个任务的状态机 + 最近 run
3. list_datasets                — 列出所有数据集（含数据规模）
4. get_dataset_info             — 单个数据集的详细范围（日期范围/数据量/列）
5. get_today_sync_summary       — 今日同步汇总

传输: stdio（本地启动 Claude Desktop / Claude Code 直连）。

启动:
  python -m app.mcp_server
  # 或作为 MCP server 注册到客户端

架构:
  - MCP 协议层: mcp.server.Server + stdio_server
  - 业务层: 复用 app.core.db.ops + app.core.sync.registry 的 helper
  - 不直接操作 PG,通过 ORM session（一致的事务/连接管理）
"""
from __future__ import annotations

import asyncio
import json
import sys
from datetime import datetime
from pathlib import Path

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))


# ── MCP server bootstrap ────────────────────────────────────────────
from mcp.server import Server
from mcp.server.sse import SseServerTransport
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent

server = Server("market_sync")


# ── helpers ─────────────────────────────────────────────────────────

def _serialize(obj):
    """JSON-friendly: datetime/date/Decimal → str."""
    if isinstance(obj, dict):
        return {k: _serialize(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_serialize(v) for v in obj]
    if isinstance(obj, (datetime,)):
        return obj.isoformat()
    try:
        # handle date
        import datetime as _dt
        if isinstance(obj, _dt.date):
            return obj.isoformat()
    except Exception:
        pass
    return obj


def _ok(data) -> list[TextContent]:
    """Standard success response: JSON dump."""
    return [TextContent(
        type="text",
        text=json.dumps(_serialize(data), ensure_ascii=False, indent=2, default=str),
    )]


def _err(msg: str) -> list[TextContent]:
    return [TextContent(type="text", text=f"❌ {msg}")]


# ── Tool implementations ─────────────────────────────────────────────

def list_sync_tasks() -> list[dict]:
    from app.tasks import TASKS
    from app.core.db import ops as db_ops
    from app.core.sync import get_registry_status

    registry = {r["dataset_id"]: r for r in get_registry_status()}
    history = db_ops.list_sync_history(days=7, limit=1000)
    last_by_task: dict[str, dict] = {}
    for h in history:
        if h["dataset_id"] not in last_by_task:
            last_by_task[h["dataset_id"]] = h

    out = []
    for ds_id, cls in TASKS.items():
        reg = registry.get(ds_id, {})
        last = last_by_task.get(ds_id)
        out.append({
            "dataset_id": ds_id,
            "name": reg.get("name", ""),
            "description": reg.get("description", ""),
            "source": reg.get("source", ""),
            "sort_order": reg.get("sort_order", 0),
            "current_status": reg.get("status", "unknown"),
            "last_run_at": (last or {}).get("started_at"),
            "last_run_status": (last or {}).get("status"),
            "last_run_rows": (last or {}).get("rows_written"),
            "history_count_7d": sum(1 for h in history if h["dataset_id"] == ds_id),
        })
    return sorted(out, key=lambda x: x["sort_order"])


def get_task_status(dataset_id: str, days: int = 7) -> dict:
    from app.core.db import ops as db_ops
    from app.core.sync import get_registry_status

    reg = next((r for r in get_registry_status() if r["dataset_id"] == dataset_id), None)
    if not reg:
        return {"error": f"unknown dataset_id: {dataset_id}"}
    history = db_ops.list_sync_history(dataset_id=dataset_id, days=days, limit=50)
    return {
        "dataset_id": dataset_id,
        "current_state": reg,
        "recent_runs": history,
    }


def list_datasets() -> list[dict]:
    """每个数据集的当前规模（行数 + 日期范围 + 来源）。"""
    from sqlalchemy import func, select
    from app.core.db.models import (
        KlineStock, KlineIndex, Kline5Min, Moneyflow, Share, TickTrade,
        MarketRegimeDaily,
        LonghubangDaily, LonghubangSeat, Stock,
        NodeCategory, Node, StockNodeMap, KlineStockMADaily,
        KlineStockMACDDaily, KlineStockKDJDaily, KlineStockBOLLDaily,
    )
    from app.core.db.orm import SessionLocal

    targets = [
        ("stocks", Stock),
        ("kline_stock", KlineStock),
        ("kline_index", KlineIndex),
        ("kline_5min", Kline5Min),
        ("kline_stock_ma_daily", KlineStockMADaily),
        ("moneyflow", Moneyflow),
        ("share", Share),
        ("tick_trade", TickTrade),
        ("market_regime_daily", MarketRegimeDaily),
        ("longhubang_daily", LonghubangDaily),
        ("longhubang_seat", LonghubangSeat),
        ("node_categories", NodeCategory),
        ("nodes", Node),
        ("stock_node_map", StockNodeMap),
        ("kline_stock_macd_daily", KlineStockMACDDaily),
        ("kline_stock_kdj_daily", KlineStockKDJDaily),
        ("kline_stock_boll_daily", KlineStockBOLLDaily),
    ]
    out = []
    with SessionLocal() as s:
        for name, model in targets:
            try:
                count = int(s.execute(select(func.count()).select_from(model)).scalar() or 0)
                dr = None
                if hasattr(model, "trade_date"):
                    mn = s.execute(select(func.min(model.trade_date))).scalar()
                    mx = s.execute(select(func.max(model.trade_date))).scalar()
                    dr = {"min": str(mn), "max": str(mx)} if mn else None
                out.append({"table": name, "rows": count, "date_range": dr})
            except Exception as e:
                out.append({"table": name, "rows": -1, "error": str(e)[:100]})
    return out


def get_dataset_info(table_name: str) -> dict:
    """单表 schema + 数据范围 + 抽样行。"""
    from sqlalchemy import inspect
    from app.core.db.orm import SessionLocal

    # 表名白名单（避免任意 SQL 注入）
    allowed = {
        "stocks", "kline_stock", "kline_index", "kline_5min",
        "kline_stock_ma_daily", "moneyflow", "share", "tick_trade",
        "market_regime_daily",
        "longhubang_daily", "longhubang_seat",
        "node_categories", "nodes", "stock_node_map",
        "kline_stock_macd_daily", "kline_stock_kdj_daily", "kline_stock_boll_daily",
        "dataset_registry", "sync_history", "config",
    }
    if table_name not in allowed:
        return {"error": f"unknown table: {table_name}", "allowed": sorted(allowed)}

    with SessionLocal() as s:
        insp = inspect(s.get_bind())
        cols = [{"name": c["name"], "type": str(c["type"]), "nullable": c.get("nullable", True)}
                for c in insp.get_columns(table_name, schema="market_data")]
        pk = insp.get_pk_constraint(table_name, schema="market_data")
        idx = [i["name"] for i in insp.get_indexes(table_name, schema="market_data")]
        sample = s.execute(
            __import__("sqlalchemy").text(
                f'SELECT * FROM market_data."{table_name}" LIMIT 3'
            )
        ).mappings().all()
        return {
            "table": table_name,
            "schema": "market_data",
            "primary_key": pk.get("constrained_columns", []),
            "indexes": idx,
            "columns": cols,
            "sample_rows": [dict(r) for r in sample],
        }


def get_today_sync_summary() -> dict:
    from datetime import date
    from app.core.db import ops as db_ops

    today = date.today().isoformat()
    summary_rows = db_ops.daily_sync_summary(days=1)
    today_summary = next((r for r in summary_rows if r["date"] == today), None) or {
        "date": today, "ok": 0, "warning": 0, "error": 0, "blocked": 0, "total": 0,
    }
    history = db_ops.list_sync_history(days=1, limit=500)
    return {
        "date": today,
        "summary": today_summary,
        "tasks_today": history,
    }


# ── MCP tool 注册 ──────────────────────────────────────────────────

@server.list_tools()
async def list_tools():
    return [
        Tool(
            name="list_sync_tasks",
            description="列出 market_sync 项目的所有同步任务（dataset_id + 名称 + 当前状态 + 最近 7 天运行统计）。",
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="get_task_status",
            description="查询单个同步任务的详细状态机 + 最近 N 天所有 run 历史。",
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_id": {"type": "string", "description": "例如 kline_daily / moneyflow"},
                    "days": {"type": "integer", "description": "查最近 N 天, 默认 7", "default": 7},
                },
                "required": ["dataset_id"],
            },
        ),
        Tool(
            name="list_datasets",
            description="列出所有 PG 数据集（含 market_data schema 下 17 张业务表）的行数 + 日期范围。",
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="get_dataset_info",
            description="查询单个数据集的表结构（列名 + 类型 + 主键 + 索引）+ 3 行抽样数据。",
            inputSchema={
                "type": "object",
                "properties": {
                    "table_name": {"type": "string", "description": "如 kline_stock / moneyflow / tick_trade"},
                },
                "required": ["table_name"],
            },
        ),
        Tool(
            name="get_today_sync_summary",
            description="今日同步情况汇总（OK / Warning / Error / Blocked 计数 + 每个任务的最新 run）。",
            inputSchema={"type": "object", "properties": {}},
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict):
    try:
        if name == "list_sync_tasks":
            return _ok(list_sync_tasks())
        elif name == "get_task_status":
            ds_id = arguments.get("dataset_id", "")
            days = int(arguments.get("days", 7))
            return _ok(get_task_status(ds_id, days))
        elif name == "list_datasets":
            return _ok(list_datasets())
        elif name == "get_dataset_info":
            return _ok(get_dataset_info(arguments.get("table_name", "")))
        elif name == "get_today_sync_summary":
            return _ok(get_today_sync_summary())
        else:
            return _err(f"unknown tool: {name}")
    except Exception as e:
        import traceback
        return _err(f"{type(e).__name__}: {e}\n{traceback.format_exc(limit=3)}")


# ── main ────────────────────────────────────────────────────────────

async def main_sse(host: str = "127.0.0.1", port: int = 8101):
    """启动 MCP 的 SSE (HTTP) 服务，供持久运行。
    
    客户端连接: http://{host}:{port}/sse
    """
    sse_transport = SseServerTransport("/messages")

    async def app(scope, receive, send):
        if scope["type"] != "http":
            return
        path = scope.get("path", "")
        if path == "/sse":
            async with sse_transport.connect_sse(scope, receive, send) as (read, write):
                await server.run(read, write, server.create_initialization_options())
        elif path == "/messages" and scope.get("method") == "POST":
            await sse_transport.handle_post_message(scope, receive, send)

    import uvicorn
    from app.core.utils.logging import get_logger
    logger = get_logger("mcp")
    logger.info(f"MCP SSE server listening on http://{host}:{port}/sse")
    config = uvicorn.Config(app, host=host, port=port, log_level="info")
    srv = uvicorn.Server(config)
    await srv.serve()


def main():
    from app.core.datasource.registry import build_default_registry
    build_default_registry()

    import sys
    if "--sse" in sys.argv:
        host = "127.0.0.1"
        port = 8101
        for i, arg in enumerate(sys.argv):
            if arg == "--host" and i + 1 < len(sys.argv):
                host = sys.argv[i + 1]
            if arg == "--port" and i + 1 < len(sys.argv):
                port = int(sys.argv[i + 1])
        asyncio.run(main_sse(host, port))
    else:
        asyncio.run(_stdio_main())


async def _stdio_main():
    async with stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            server.create_initialization_options(),
        )


if __name__ == "__main__":
    main()