"""
智能体使用统计工具
记录和统计智能体的调用次数（PostgreSQL数据库存储）
"""
from datetime import datetime
from typing import Dict, Any

from config import POSTGRES_CONNECTION_STRING


ROUTE_FLAG_DISPLAY_NAMES = {
    "故障排查与修理": "故障排查与修理",
    "操作使用": "操作使用",
    "舰船百科": "舰船百科"
}

ROUTE_FLAG_ICONS = {
    "故障排查与修理": "detection",
    "操作使用": "repair",
    "舰船百科": "doc"
}

ROUTE_FLAG_COLORS = {
    "故障排查与修理": "#FFB800",
    "操作使用": "#00C2FF",
    "舰船百科": "#9153FF"
}


async def init_agent_usage_table():
    """
    初始化 agent_usage_records 表
    """
    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError:
        print("[agent_usage_statistics] psycopg_pool 未安装，无法初始化表")
        return

    async with AsyncConnectionPool(
        POSTGRES_CONNECTION_STRING,
        kwargs={"autocommit": True}
    ) as pool:
        async with pool.connection() as conn:
            async with conn.cursor() as cur:
                await cur.execute("""
                    CREATE TABLE IF NOT EXISTS agent_usage_records (
                        id SERIAL PRIMARY KEY,
                        chat_id TEXT NOT NULL DEFAULT '',
                        route_flag TEXT NOT NULL DEFAULT '',
                        message_id TEXT NOT NULL DEFAULT '',
                        created_at TIMESTAMP NOT NULL DEFAULT NOW(),
                        UNIQUE(chat_id, route_flag)
                    )
                """)
                await cur.execute("""
                    CREATE INDEX IF NOT EXISTS agent_usage_records_route_flag_idx ON agent_usage_records(route_flag)
                """)
                await cur.execute("""
                    CREATE INDEX IF NOT EXISTS agent_usage_records_created_at_idx ON agent_usage_records(created_at)
                """)
                print("[agent_usage_statistics] agent_usage_records 表初始化完成")


async def record_agent_usage(chat_id: str, route_flag: str, message_id: str = "") -> bool:
    """
    记录智能体使用情况到数据库

    Args:
        chat_id: 会话ID
        route_flag: 路由标志（智能体类型）
        message_id: 消息ID

    Returns:
        是否记录成功
    """
    if not chat_id or not route_flag:
        return False

    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError:
        print("[agent_usage_statistics] psycopg_pool 未安装，无法记录")
        return False

    try:
        async with AsyncConnectionPool(
            POSTGRES_CONNECTION_STRING,
            kwargs={"autocommit": True}
        ) as pool:
            async with pool.connection() as conn:
                async with conn.cursor() as cur:
                    await cur.execute(
                        """
                        INSERT INTO agent_usage_records (chat_id, route_flag, message_id, created_at)
                        VALUES (%s, %s, %s, %s)
                        ON CONFLICT (chat_id, route_flag) DO UPDATE SET
                            message_id = EXCLUDED.message_id,
                            created_at = EXCLUDED.created_at
                        """,
                        (chat_id, route_flag, message_id or "", datetime.now())
                    )
        return True
    except Exception as e:
        print(f"[agent_usage_statistics] 记录智能体使用失败: {e}")
        return False


def _format_count(count: int) -> str:
    if count >= 10000:
        return f"{count / 10000:.1f}万"
    else:
        return f"{count:,}"


async def get_agent_usage_statistics() -> Dict[str, Any]:
    """
    从数据库查询智能体使用统计

    统计逻辑：
    - 同一 chat_id + route_flag 只保留一条记录
    - 统计各类型的总使用次数
    - 统计近6个月的月度使用次数

    Returns:
        统计结果字典
    """
    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError:
        print("[agent_usage_statistics] psycopg_pool 未安装，无法查询")
        return {"success": False, "data": None}

    required_routes = ["操作使用", "故障排查与修理", "舰船百科"]

    now = datetime.now()
    last_6_months = []
    for i in range(5, -1, -1):
        year = now.year
        month = now.month - i
        while month <= 0:
            month += 12
            year -= 1
        last_6_months.append({
            "year": year,
            "month": month,
            "key": f"{year}-{month:02d}",
            "label": f"{month}月"
        })

    try:
        async with AsyncConnectionPool(
            POSTGRES_CONNECTION_STRING,
            kwargs={"autocommit": True}
        ) as pool:
            async with pool.connection() as conn:
                async with conn.cursor() as cur:
                    # 查询各类型总使用次数
                    route_counts = {}
                    for route_flag in required_routes:
                        await cur.execute(
                            """
                            SELECT COUNT(*) FROM agent_usage_records
                            WHERE route_flag = %s
                            """,
                            (route_flag,)
                        )
                        row = await cur.fetchone()
                        route_counts[route_flag] = row[0] if row else 0

                    # 查询近6个月各类型月度使用次数
                    monthly_counts = {route: {m["key"]: 0 for m in last_6_months} for route in required_routes}

                    for route_flag in required_routes:
                        for m in last_6_months:
                            month_start = datetime(m["year"], m["month"], 1)
                            if m["month"] == 12:
                                month_end = datetime(m["year"] + 1, 1, 1)
                            else:
                                month_end = datetime(m["year"], m["month"] + 1, 1)

                            await cur.execute(
                                """
                                SELECT COUNT(*) FROM agent_usage_records
                                WHERE route_flag = %s AND created_at >= %s AND created_at < %s
                                """,
                                (route_flag, month_start, month_end)
                            )
                            row = await cur.fetchone()
                            monthly_counts[route_flag][m["key"]] = row[0] if row else 0

    except Exception as e:
        print(f"[agent_usage_statistics] 查询统计数据失败: {e}")
        import traceback
        traceback.print_exc()
        return {"success": False, "data": None}

    total = sum(route_counts.get(route, 0) for route in required_routes)

    types_data = []
    monthly_stats = []
    for route_flag in required_routes:
        count = route_counts.get(route_flag, 0)
        display_name = ROUTE_FLAG_DISPLAY_NAMES.get(route_flag, route_flag)
        icon = ROUTE_FLAG_ICONS.get(route_flag, "info")
        color = ROUTE_FLAG_COLORS.get(route_flag, "#8C8C8C")

        types_data.append({
            "name": display_name,
            "count": _format_count(count),
            "icon": icon,
            "color": color
        })

        monthly_data = []
        for m in last_6_months:
            monthly_data.append({
                "month": m["label"],
                "count": monthly_counts[route_flag][m["key"]]
            })

        monthly_stats.append({
            "name": display_name,
            "color": color,
            "monthlyCounts": monthly_data
        })

    return {
        "success": True,
        "data": {
            "total": _format_count(total),
            "types": types_data,
            "monthlyStats": monthly_stats,
            "months": [m["label"] for m in last_6_months]
        }
    }