#!/usr/bin/env python3
"""重构截图识别: 遵循 LLM 执行协议

协议要求:
1. LLM 输出节标题格式(不用 JSON)
2. parse_response 按节标题精确匹配
3. 多层门禁校验
4. 写入前快照
5. 确认反馈
"""
import os, sys, json, re, sqlite3
os.chdir("/home/hmo/MoFin/deploy/profile-scripts")
os.environ["GIT_ALLOW_COMMIT"] = "1"

DB = "/home/hmo/MoFin/data/mofin.db"

# === 1. 定义节标题 schema ===
TRADE_SECTIONS = [
    {"name": "交易动作", "field": "action", "type": "enum", "values": ["买入", "卖出"]},
    {"name": "股票代码", "field": "code", "type": "string", "constraint": "6位数字"},
    {"name": "股票名称", "field": "name", "type": "string"},
    {"name": "交易数量", "field": "shares", "type": "integer", "constraint": "必须>0"},
    {"name": "交易价格", "field": "price", "type": "number", "constraint": "必须>0"},
]

# === 2. 节标题解析器(复用已有) ===
def _section_line(text, name):
    for line in text.split("\n"):
        if re.match(r'^\s*【' + name + r'】', line):
            return line
    return ""

# === 3. 解析 LLM 输出 ===
def parse_trade_response(text):
    result = {}
    for sec in TRADE_SECTIONS:
        line = _section_line(text, sec["name"])
        if not line:
            continue
        val = re.sub(r'^\s*【' + sec["name"] + r'】\s*', '', line).strip()
        
        if sec["type"] == "enum":
            if val in sec["values"]:
                result[sec["field"]] = val
        elif sec["type"] == "integer":
            nums = re.findall(r'\d+', val)
            if nums:
                result[sec["field"]] = int(nums[0])
        elif sec["type"] == "number":
            nums = re.findall(r'\d+\.?\d*', val)
            if nums:
                result[sec["field"]] = float(nums[0])
        elif sec["type"] == "string":
            result[sec["field"]] = val
    return result

# === 4. 门禁校验 ===
def validate_trade(parsed, current_holdings=None):
    errors = []
    
    # 动作
    if parsed.get("action") not in ("买入", "卖出"):
        errors.append("交易动作必须是'买入'或'卖出'")
    
    # 代码
    code = parsed.get("code", "")
    if not re.match(r'^\d{6}$', code):
        errors.append(f"股票代码'{code}'不是6位数字")
    else:
        # 检查代码是否存在
        conn = sqlite3.connect(DB, timeout=30)
        exists = conn.execute("SELECT 1 FROM stock_daily WHERE code=? LIMIT 1", (code,)).fetchone()
        conn.close()
        if not exists:
            errors.append(f"股票代码'{code}'在 stock_daily 中不存在")
    
    # 数量
    shares = parsed.get("shares", 0)
    if not isinstance(shares, int) or shares <= 0:
        errors.append(f"交易数量'{shares}'必须是正整数")
    
    # 价格
    price = parsed.get("price", 0)
    if not isinstance(price, (int, float)) or price <= 0:
        errors.append(f"交易价格'{price}'必须是正数")
    elif price > 0:
        # 与现价偏差检查
        conn = sqlite3.connect(DB, timeout=30)
        row = conn.execute("SELECT close FROM stock_daily WHERE code=? ORDER BY date DESC LIMIT 1", (code,)).fetchone()
        conn.close()
        if row and row[0]:
            current_price = float(row[0])
            if current_price > 0:
                deviation = abs(price - current_price) / current_price
                if deviation > 0.1:
                    errors.append(f"交易价格{price}与现价{current_price}偏差{deviation:.1%}超过10%")
    
    # 卖出数量检查
    if parsed.get("action") == "卖出" and current_holdings:
        current_shares = current_holdings.get(code, {}).get("shares", 0)
        if shares > current_shares:
            errors.append(f"卖出数量{shares}超过当前持仓{current_shares}")
    
    return errors

# === 5. 写入 holdings ===
def execute_trade(parsed):
    code = parsed["code"]
    action = parsed["action"]
    shares = parsed["shares"]
    price = parsed["price"]
    name = parsed.get("name", "")
    
    conn = sqlite3.connect(DB, timeout=30)
    conn.execute("PRAGMA busy_timeout=30000")
    
    # 快照(简单版: 记录当前状态)
    old = conn.execute("SELECT shares, cost FROM holdings WHERE code=?", (code,)).fetchone()
    old_shares = old[0] if old else 0
    old_cost = old[1] if old else 0
    
    if action == "买入":
        new_shares = old_shares + shares
        new_cost = ((old_cost * old_shares) + (price * shares)) / new_shares if new_shares > 0 else price
        conn.execute("""
            INSERT INTO holdings (code, name, shares, cost, price, is_active)
            VALUES (?, ?, ?, ?, ?, 1)
            ON CONFLICT(code) DO UPDATE SET
                shares=?, cost=?, price=?, is_active=1
        """, (code, name, new_shares, round(new_cost, 4), price, new_shares, round(new_cost, 4), price))
    elif action == "卖出":
        new_shares = old_shares - shares
        if new_shares <= 0:
            new_shares = 0
            conn.execute("UPDATE holdings SET shares=0, is_active=0 WHERE code=?", (code,))
        else:
            conn.execute("UPDATE holdings SET shares=?, price=? WHERE code=?", (new_shares, price, code))
    
    conn.commit()
    conn.close()
    
    print(f"  ✅ {action} {name}({code}) {shares}股 @{price} → 持仓{new_shares}股")
    return {"ok": True, "new_shares": new_shares}

# === 6. 完整流程 ===
def process_screenshot(ocr_text):
    """OCR文本 → 结构化提取 → 校验 → 写入"""
    from llm_client import call_llm, REASSESS_MODEL
    
    prompt = f"""请从以下OCR文字中提取交易信息，严格按节标题格式输出：

{ocr_text}

请严格按以下格式输出（每个字段一个节标题）：

【交易动作】买入 或 卖出
【股票代码】6位数字
【股票名称】中文名称
【交易数量】正整数
【交易价格】数字
"""
    
    result = call_llm(prompt, model=REASSESS_MODEL, max_tokens=None)
    if not result["ok"] or not result.get("content"):
        return {"ok": False, "error": "LLM调用失败"}
    
    parsed = parse_trade_response(result["content"])
    
    # 门禁校验
    errors = validate_trade(parsed)
    if errors:
        return {"ok": False, "errors": errors, "parsed": parsed}
    
    # 获取当前持仓(卖出检查用)
    conn = sqlite3.connect(DB, timeout=30)
    current = {}
    for row in conn.execute("SELECT code, shares, cost FROM holdings WHERE is_active=1"):
        current[row[0]] = {"shares": row[1], "cost": row[2]}
    conn.close()
    
    # 二次校验(卖出数量)
    errors = validate_trade(parsed, current)
    if errors:
        return {"ok": False, "errors": errors, "parsed": parsed}
    
    # 执行
    exec_result = execute_trade(parsed)
    return {"ok": True, "parsed": parsed, "exec": exec_result}

if __name__ == "__main__":
    print("trade_capture 模块已加载(协议版)")
    print(f"节标题: {[s['name'] for s in TRADE_SECTIONS]}")