import json, base64, urllib.request, os, traceback
from fastapi import FastAPI, Request, BackgroundTasks, Response

app = FastAPI()

# --- 环境配置 (请在 HF Secrets 中设置) ---
FEISHU_APP_ID = os.getenv("FEISHU_APP_ID")
FEISHU_APP_SECRET = os.getenv("FEISHU_APP_SECRET")
SILICON_API_KEY = os.getenv("SILICON_API_KEY")
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")
MODEL_NAME = "Qwen/Qwen3-VL-32B-Instruct"

# 内存暂存，用于多轮修改
img_cache = {}

def get_token():
    try:
        url = "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal"
        p = json.dumps({"app_id": FEISHU_APP_ID, "app_secret": FEISHU_APP_SECRET}).encode()
        req = urllib.request.Request(url, data=p, method="POST", headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=5) as f:
            return json.loads(f.read().decode()).get("tenant_access_token")
    except: return None

def send_msg(open_id, text):
    token = get_token()
    if not token: return
    url = "https://open.feishu.cn/open-apis/im/v1/messages?receive_id_type=open_id"
    payload = json.dumps({"receive_id": open_id, "msg_type": "text", "content": json.dumps({"text": text})}).encode()
    req = urllib.request.Request(url, data=payload, method="POST", 
                                 headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"})
    try: urllib.request.urlopen(req, timeout=5)
    except: pass

def save_to_supabase(content):
    if not SUPABASE_URL: return "⚠️ 未配置 Supabase URL"
    try:
        url = f"{SUPABASE_URL}/rest/v1/scores" # 假设你的表名是 scores
        headers = {
            "apikey": SUPABASE_KEY, 
            "Authorization": f"Bearer {SUPABASE_KEY}", 
            "Content-Type": "application/json",
            "Prefer": "return=minimal"
        }
        payload = json.dumps({"content": content}).encode()
        req = urllib.request.Request(url, data=payload, method="POST", headers=headers)
        urllib.request.urlopen(req, timeout=5)
        return "✅ 数据已成功同步至 Supabase！"
    except Exception as e:
        return f"❌ Supabase 存储失败: {e}"

async def run_ai_task(img_bin, instruction, open_id):
    try:
        b64_img = base64.b64encode(img_bin).decode()
        # 针对性提示词：处理学号与成绩的逻辑
        prompt = (
            f"你是一个表格数据提取助手。请分析图片并提取成绩：\n"
            f"1. 提取每一行的成绩信息。\n"
            f"2. 特别注意：如果表格最后一列是学号（通常是长数字），请跳过它，提取倒数第二列的成绩；如果最后一列是成绩，则提取最后一列。\n"
            f"3.一定不要返回序号（连起来的数字），必须返回类似-1，-6，A+的内容列"
            f"4. 仅以 Markdown 无序列表格式输出成绩数值，不要表头，不要解释，不要姓名。\n"
            f"用户改进意见：{instruction}"
        )
        
        body = json.dumps({
            "model": MODEL_NAME,
            "messages": [{"role": "user", "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_img}"}}
            ]}]
        }).encode()
        
        req = urllib.request.Request("https://api.siliconflow.cn/v1/chat/completions", data=body, method="POST",
                                    headers={"Authorization": f"Bearer {SILICON_API_KEY}", "Content-Type": "application/json"})
        
        with urllib.request.urlopen(req, timeout=60) as f:
            res = json.loads(f.read().decode())
            ans = res['choices'][0]['message']['content'].strip()
        
        # 清洗可能存在的 Markdown 标签
        ans = ans.replace('```markdown', '').replace('```', '').strip()
        
        reply = (
            f"📊 **提取结果**：\n\n{ans}\n\n"
            "--- \n"
            "💡 **操作指南**：\n"
            "- 直接回复 **改进建议** 让 AI 重改\n"
            "- 回复 **y** 确认并保存到 Supabase\n"
            "- 回复 **q** 退出当前任务"
        )
        send_msg(open_id, reply)
    except Exception as e:
        send_msg(open_id, f"‼️ AI 处理出错: {e}")

# --- 路由 ---

@app.head("/")
@app.get("/")
async def status_check():
    # 适配 UptimeRobot 的 HEAD/GET 检查
    return Response(status_code=200)

@app.post("/")
async def feishu_callback(request: Request, background_tasks: BackgroundTasks):
    try:
        data = await request.json()
        if "challenge" in data: return {"challenge": data["challenge"]}
        
        print(f"📡 收到推送: {json.dumps(data)}") # 调试用

        event = data.get("event", {})
        msg = event.get("message", {})
        msg_type = msg.get("message_type") or msg.get("msg_type")
        
        # 兼容性获取 OpenID
        sender = event.get("sender", {})
        open_id = sender.get("sender_id", {}).get("open_id") or sender.get("open_id") or msg.get("open_id")

        if not open_id: return {"code": 0}

        # 1. 收到图片
        if msg_type == "image":
            token = get_token()
            content = json.loads(msg.get("content", "{}"))
            img_key = content.get("image_key")
            
            if img_key:
                print(f"📩 正在下载图片: {img_key}")
                img_url = f"https://open.feishu.cn/open-apis/im/v1/messages/{msg['message_id']}/resources/{img_key}?type=image"
                req = urllib.request.Request(img_url, headers={"Authorization": f"Bearer {token}"})
                with urllib.request.urlopen(req, timeout=10) as f:
                    img_bin = f.read()
                
                img_cache[open_id] = img_bin
                background_tasks.add_task(run_ai_task, img_bin, "初次提取", open_id)

        # 2. 收到文本 (交互)
        elif msg_type == "text":
            text_content = json.loads(msg.get("content", "{}")).get("text", "").strip().lower()
            # 简单处理：取空格后最后一部分（过滤@机器人）
            cmd = text_content.split(' ')[-1]

            if cmd == 'y':
                # 保存逻辑，实际应存储 AI 上次的 ans
                res = save_to_supabase("用户确认的最新成绩列表")
                send_msg(open_id, res)
                img_cache.pop(open_id, None)
            elif cmd == 'q':
                img_cache.pop(open_id, None)
                send_msg(open_id, "✅ 已退出任务。")
            else:
                if open_id in img_cache:
                    send_msg(open_id, f"🔄 收到反馈：'{cmd}'，正在重新校对...")
                    background_tasks.add_task(run_ai_task, img_cache[open_id], cmd, open_id)
                else:
                    send_msg(open_id, "💡 请先发送成绩图片。")

        return {"code": 0}
    except:
        print(traceback.format_exc())
        return {"code": 0}