#!/usr/bin/env python
# -*- coding: utf-8 -*-

from mcp.server.fastmcp import FastMCP
import requests
import json
import os
import argparse
from datetime import datetime, timedelta
from typing import List, Optional

# 解析命令行参数
parser = argparse.ArgumentParser(description='广告数据查询MCP服务')
parser.add_argument('--token', type=str, required=True, help='API访问token')
args = parser.parse_args()

# 创建MCP服务器
mcp = FastMCP("广告数据查询服务")

# 从命令行获取token
def get_token_from_config():
    # 只从命令行获取token
    if args.token:
        return args.token
    else:
        raise ValueError("必须提供命令行参数--token")

def compress_response(response_data: dict, requested_zhibiao=None, is_deep=False) -> dict:
    """
    压缩API响应数据，减少token占用
    
    策略：
    1. 将list数据转换为表格格式（表头+行数据）
    2. 将原始字段名转换为中文指标名作为headers
    3. 只保留用户请求的指标，过滤掉不需要的数据
    4. 不返回prop_map映射关系，直接使用中文表头
    5. 当is_deep为True时，额外保留所有以UI结尾的字段，确保下探功能正常工作
    """
    if not response_data or response_data.get('code') != 0:
        return response_data
    
    data = response_data.get('data', {})
    if not data:
        return response_data
    
    result = {
        "code": response_data.get('code'),
        "msg": response_data.get('msg')
    }
    
    # 处理列表数据
    items = data.get('list', [])
    if not items:
        result['data'] = data
        return result
    
    # 获取propMap用于转换表头
    prop_map = data.get('propMap', {})
    
    # 如果用户提供了指定的指标列表，准备要保留的字段
    field_indices_to_keep = None
    requested_fields = []
    if requested_zhibiao and prop_map:
        # 找出用户请求指标对应的英文字段名
        for zhibiao in requested_zhibiao:
            if zhibiao in prop_map:
                requested_fields.append(prop_map[zhibiao])
        
        if requested_fields:
            field_indices_to_keep = []
    
    # 构建压缩后的表格式数据
    all_headers = list(items[0].keys())
    
    # 反转propMap，创建从英文字段名到中文指标名的映射
    reverse_prop_map = {}
    if prop_map:
        for cn_name, en_field in prop_map.items():
            reverse_prop_map[en_field] = cn_name
    
    # 转换headers为中文，并过滤字段
    headers = []
    filtered_indices = []  # 要保留的字段索引
    date_index = -1  # 日期字段的索引
    
    for i, field in enumerate(all_headers):
        # 判断是否保留该字段
        keep_field = True
        
        # 如果指定了要保留的字段，检查当前字段是否在列表中
        if field_indices_to_keep is not None:
            # 如果是下探查询(is_deep=True)且字段以UI结尾，则保留
            if is_deep and field.endswith('UI'):
                keep_field = True
            # 否则，只保留用户请求的字段
            elif field not in requested_fields:
                keep_field = False
        
        if not keep_field:
            continue
        
        # 将英文字段名转换为中文指标名
        if field in reverse_prop_map:
            header_name = reverse_prop_map[field]
            headers.append(header_name)
            # 记录日期字段的索引
            if header_name == '日期':
                date_index = len(headers) - 1
        else:
            headers.append(field)  # 如果没有对应的中文名，保留英文名
            # 检查是否为Date字段
            if field == 'Date':
                date_index = len(headers) - 1
        
        filtered_indices.append(i)
    
    # 预先检查是否有总计行
    has_total_row = False
    total_row_index = -1
    for idx, item in enumerate(items):
        if 'Date' in item and item['Date'] == '总计':
            has_total_row = True
            total_row_index = idx
            break
    
    # 按过滤后的索引提取数据
    rows = []
    for idx, item in enumerate(items):
        item_values = list(item.values())
        # 只保留需要的字段
        filtered_row = [item_values[i] for i in filtered_indices]
        
        # 如果是总计行且日期字段存在，确保其值为"总计"
        if has_total_row and idx == total_row_index and date_index >= 0 and len(filtered_row) > date_index:
            filtered_row[date_index] = "总计"
        
        rows.append(filtered_row)
    
    # 如果没有找到明确的总计行，但有多行数据，则假设最后一行是总计行
    if not has_total_row and len(rows) > 1 and date_index >= 0:
        # 检查最后一行的日期字段是否为空或非日期格式
        last_row = rows[-1]
        if len(last_row) > date_index:
            if not last_row[date_index] or last_row[date_index] == "":
                last_row[date_index] = "总计"
    
    # 构建结果 - 使用中文字段名，不返回prop_map
    compressed_data = {
        "headers": headers,  # 已转换为中文的表头
        "rows": rows
    }
    
    result['data'] = compressed_data
    
    return result

@mcp.tool()
def get_ad_count_list(
    version: str = "0.1.85", 
    appid: str = "59",
    start_time: Optional[str] = None,
    end_time: Optional[str] = None,
    zhibiao_list: Optional[List[str]] = ["日期", "创角成本", "新增创角", "广告计划名称", "创意名称", "项目名称", "广告状态", "备注", "新增注册", "创角率", "点击率", "激活率", "点击成本", "活跃用户", "曝光次数", "千次展现均价", "点击数", "一阶段花费", "二阶段花费", "当日充值", "当日付费次数", "当日充值人数", "新增付费人数", "首充付费人数", "首充付费次数", "老用户付费人数", "新增付费金额", "首充付费金额", "老用户付费金额", "新增付费率", "活跃付费率", "活跃arppu", "新增arppu", "小游戏注册首日广告变现金额", "小游戏注册首日广告变现ROI", "当月注册用户充值金额", "消耗", "新增付费成本", "付费成本", "注册成本", "首日ROI", "累计ROI", "分成后首日ROI", "分成后累计ROI", "付费首日ROI", "付费累计ROI", "付费分成后首日ROI", "付费分成后累计ROI", "计算累计ROI所用金额", "计算累计ROI所用消耗", "24小时ROI"],
    media: Optional[List[str]] = None,
    group_key: Optional[str] = None,
    toushou: Optional[List[str]] = None,
    self_cid: Optional[List[str]] = None,
    ji_hua_id: Optional[List[str]] = None,
    ji_hua_name: Optional[str] = None,
    ad_status: Optional[List[str]] = None,
    creative_id: Optional[List[str]] = None,
    vp_adgroup_id: Optional[List[str]] = None,
    is_deep: Optional[bool] = False,
) -> dict:
    """
    当用户需要分析广告数据时，你需要使用MCP工具`get_ad_count_list`来查询广告数据。该工具可以获取广告投放的各类指标数据，支持多种筛选和分组方式。
    """
    

    token = get_token_from_config()
    
    # 设置默认值
    if start_time is None:
        # 默认查询昨天的数据
        yesterday = datetime.now() - timedelta(days=1)
        start_time = yesterday.strftime("%Y-%m-%d")
    
    if end_time is None:
        # 默认查询到今天
        end_time = datetime.now().strftime("%Y-%m-%d")
    if zhibiao_list is None:
        zhibiao_list = ["日期", "创角成本", "新增创角", "广告计划名称", "创意名称", "项目名称", "广告状态", "备注", "新增注册", "创角率", "点击率", "激活率", "点击成本", "活跃用户", "曝光次数", "千次展现均价", "点击数", "一阶段花费", "二阶段花费", "当日充值", "当日付费次数", "当日充值人数", "新增付费人数", "首充付费人数", "首充付费次数", "老用户付费人数", "新增付费金额", "首充付费金额", "老用户付费金额", "新增付费率", "活跃付费率", "活跃arppu", "新增arppu", "小游戏注册首日广告变现金额", "小游戏注册首日广告变现ROI", "当月注册用户充值金额", "消耗", "新增付费成本", "付费成本", "注册成本", "首日ROI", "累计ROI", "分成后首日ROI", "分成后累计ROI", "付费首日ROI", "付费累计ROI", "付费分成后首日ROI", "付费分成后累计ROI", "计算累计ROI所用金额", "计算累计ROI所用消耗", "24小时ROI"]

    # API接口地址
    url = "https://bi.dartou.com/testapi/ad/GetAdCountList"
    
    # 设置请求头
    headers = {
        "X-Token": token,
        "X-Ver": version,
        "Content-Type": "application/json"
    }
    
    # 构建请求体
    payload = {
        "appid": appid,
        "start_time": start_time,
        "end_time": end_time,
        "zhibiao_list": zhibiao_list,
        "media": media,
        "group_key": group_key,
        "toushou": toushou,
        "self_cid": self_cid,
        "ji_hua_id": ji_hua_id,
        "ji_hua_name": ji_hua_name,
        "ad_status": ad_status,
        "creative_id": creative_id,
        "vp_adgroup_id": vp_adgroup_id,
        "is_deep": is_deep
    }
    
    
    try:
        # 发送POST请求
        response = requests.post(url, headers=headers, data=json.dumps(payload))
        
        # 解析响应
        result = response.json()
        
        # 检查响应状态
        if result.get("code") == 0:
            print("请求成功!")
            return compress_response(result, zhibiao_list, is_deep)
        else:
            print(f"请求失败: {result.get('msg')}")
            return result
    
    except Exception as e:
        print(f"发生错误: {str(e)}")
        return {"code": -1, "msg": str(e)}
    

def main() -> None:
    mcp.run(transport="stdio")
