"""
Prefect MCP 销售数据管道
替代 mcp_sales_pipeline.py — 自动重试/日志/状态追踪
用法: python3 mcp_prefect_pipeline.py
"""
import json, subprocess, time, os
from datetime import datetime, timedelta
from collections import defaultdict
from prefect import flow, task, get_run_logger

TK = "m-abfb29e8-3104-434f-9944-8d0bb592f8cd"
PK = "6593cd71471290e3cc6be6e6"
REL_ID = "6593cd71471290e3cc6be6e6:work_item_relation:relation_1713785263022"
BASE = "/Users/liuxinyuan/Desktop/Hermes输出-工作类/数据"
TODAY = datetime(2026, 7, 9)

@task(retries=3, retry_delay_seconds=5, timeout_seconds=60)
def mcp_call(method: str, args: dict) -> dict:
    """调用飞书MCP，自动重试3次"""
    r = subprocess.run(['curl','-s','--max-time','60','-X','POST',
        'https://project.feishu.cn/mcp_server/v1',
        '-H',f'X-Mcp-Token: {TK}','-H','Content-Type: application/json',
        '-d', json.dumps({"jsonrpc":"2.0","method":"tools/call","params":{"name":method,"arguments":args},"id":1})],
        capture_output=True, text=True, timeout=65)
    raw = json.loads(r.stdout)
    for c in raw.get('result',{}).get('content',[]):
        t = c.get('text','')
        if 'log_id' in t: continue
        if t.startswith('{'): return json.loads(t)
    return {}

@task(retries=2, retry_delay_seconds=3)
def pull_july_orders() -> list:
    """拉取7月全部销售订单"""
    logger = get_run_logger()
    orders = []
    page = 0
    while True:
        mql = f'SELECT `单据编号`, `创建时间`, `金额`, `货号#`, `单价`, work_item_id FROM `销售管理`.`销售订单` WHERE `创建时间` >= "2026-07-01" LIMIT 50 OFFSET {page*50}'
        result = mcp_call("search_by_mql", {"project_key": PK, "mql": mql})
        if not result: break
        
        for gid, gitems in result.get('data',{}).items():
            for item in gitems:
                fields = {}
                for f in item.get('moql_field_list',[]):
                    k = f['name']
                    vdict = f.get('value', {})
                    vals = list(vdict.values()) if vdict else ['']
                    v = vals[0] if vals else ''
                    if isinstance(v, dict): v = list(v.values())[0] if v else ''
                    fields[k] = v
                orders.append(fields)
        
        page += 1
        if len(result.get('data',{}).get('1',[])) < 50: break
        time.sleep(0.3)
    
    logger.info(f"拉取7月订单: {len(orders)}条")
    return orders

@task(retries=1, retry_delay_seconds=2)
def enrich_customers(orders: list) -> tuple:
    """补全客户名"""
    logger = get_run_logger()
    customer_cache = {}
    
    for i, o in enumerate(orders):
        wid = str(o.get('工作项id',o.get('work_item_id','')))
        if not wid or wid in customer_cache: continue
        
        result = mcp_call("list_related_workitem", {
            "project_key": PK, "work_item_id": wid, "relation_id": REL_ID
        })
        
        if result and result.get('list'):
            customer_cache[wid] = result['list'][0].get('name', '无')
        else:
            customer_cache[wid] = '未关联'
        
        if (i+1) % 50 == 0:
            logger.info(f"  客户补全: {i+1}/{len(orders)}")
        time.sleep(0.15)
    
    logger.info(f"客户补全完成: {len(customer_cache)}个")
    return orders, customer_cache

@task
def build_and_save(orders: list, customer_cache: dict) -> dict:
    """构建每日销售JSON并分类"""
    logger = get_run_logger()
    
    daily = defaultdict(list)
    daily_amt = defaultdict(float)
    life_sci_amt = 0
    diag_amt = 0
    
    for o in orders:
        wid = str(o.get('工作项id',o.get('work_item_id','')))
        cust = customer_cache.get(wid, '未查询')
        date = (o.get('创建时间','') or '')[:10]
        amt = float(o.get('金额',0) or 0)
        
        is_ls = '课题组' in cust
        
        entry = {
            'order_no': o.get('单据编号',''),
            'customer': cust,
            'product_code': o.get('货号#',''),
            'unit_price': float(o.get('单价',0) or 0),
            'qty': float(o.get('销售数量',0) or 0),
            'amount': amt,
            'date': date,
            'url': f'https://project.feishu.cn/xsguanli/xsdd/detail/{wid}' if wid else ''
        }
        
        daily[date].append(entry)
        daily_amt[date] += amt
        if is_ls: life_sci_amt += amt
        else: diag_amt += amt
    
    july_total = life_sci_amt + diag_amt
    
    output = {
        'days': sorted(daily.keys(), reverse=True),
        'daily': dict(daily),
        'daily_summary': {d:{'amount':daily_amt[d],'count':len(daily[d])} for d in daily},
        'total_orders': len(orders),
        'total_amount': july_total,
        'life_sci_amount': life_sci_amt,
        'diag_amount': diag_amt,
        'generated': TODAY.strftime('%Y-%m-%d %H:%M'),
        'source': '飞书项目MCP · Prefect管道'
    }
    
    outpath = f'{BASE}/daily_sales_mcp.json'
    with open(outpath, 'w', encoding='utf-8') as f:
        json.dump(output, f, ensure_ascii=False, default=str)
    
    logger.info(f"✅ 保存: {len(orders)}条, ¥{july_total:,.0f} (生命科学¥{life_sci_amt:,.0f} + 诊断¥{diag_amt:,.0f})")
    return output

@flow(name="MCP销售管道", log_prints=True)
def mcp_sales_pipeline():
    """每日从飞书MCP拉取销售订单，补全客户名，生成JSON"""
    orders = pull_july_orders()
    orders, customer_cache = enrich_customers(orders)
    result = build_and_save(orders, customer_cache)
    print(f"\n📊 管道完成: {result['total_orders']}条 ¥{result['total_amount']:,.0f}")
    return result

if __name__ == '__main__':
    mcp_sales_pipeline()