import os
import json
import pymysql
from pymysql.cursors import DictCursor
import pymongo
import requests
from dotenv import load_dotenv
from typing import List, Dict, Generator

# 加载环境变量
load_dotenv()

class StudentDataProcessorLMStudio:
    def __init__(self):
        # 初始化 MySQL 连接（支持批量遍历）
        self.mysql_conn = pymysql.connect(
            host=os.getenv("MYSQL_HOST"),
            port=int(os.getenv("MYSQL_PORT")),
            user=os.getenv("MYSQL_USER"),
            password=os.getenv("MYSQL_PASSWORD"),
            database=os.getenv("MYSQL_DB"),
            charset="utf8mb4"
        )
        self.mysql_cursor = self.mysql_conn.cursor()

        # 初始化 MongoDB 连接
        self.mongo_client = pymongo.MongoClient(os.getenv("MONGO_URI"))
        self.mongo_db = self.mongo_client[os.getenv("MONGO_DB")]
        self.mongo_collection = self.mongo_db[os.getenv("MONGO_COLLECTION")]

        # LM Studio 配置
        self.lm_studio_api_url = os.getenv("LM_STUDIO_API_URL")
        self.lm_studio_model = os.getenv("LM_STUDIO_MODEL_NAME")
        self.lm_studio_api_token = os.getenv("LM_STUDIO_API_TOKEN")
        self.batch_size = int(os.getenv("BATCH_SIZE"))

    def batch_read_student_data(self) -> Generator[List[Dict], None, None]:
        """分批遍历 MySQL 所有学生数据（避免一次性加载过多数据）"""
        try:
            # 先获取总条数，用于分批
            count_query = "SELECT COUNT(*) as total FROM student WHERE is_deleted = 0"
            self.mysql_cursor.execute(count_query)
            total = self.mysql_cursor.fetchone()["total"]
            print(f"MySQL 中共计 {total} 条学生数据，将按每批 {self.batch_size} 条处理")

            # 分批读取
            offset = 0
            while offset < total:
                query = f"""
                    SELECT id, name, age, gender, class, phone, email, address 
                    FROM student 
                    WHERE is_deleted = 0
                    LIMIT {self.batch_size} OFFSET {offset}
                """
                self.mysql_cursor.execute(query)
                batch_data = self.mysql_cursor.fetchall()
                if not batch_data:
                    break
                yield batch_data
                offset += self.batch_size
                print(f"已读取第 {offset//self.batch_size} 批数据（累计 {offset} 条）")
        except Exception as e:
            print(f"分批读取 MySQL 数据失败：{str(e)}")
            raise

    def process_batch_with_lmstudio(self, batch_data: List[Dict]) -> List[Dict]:
        """调用 LM Studio 本地模型整理单批数据"""
        if not batch_data:
            return []
        
        # 构造适配本地模型的提示词（更简洁，避免本地模型处理复杂指令出错）
        prompt = f"""
        请严格按照以下要求整理学生数据，仅返回JSON数组，不要添加任何额外文字、解释或备注：
        整理规则：
        1. 标准化字段：
           - 姓名：去除首尾空格，仅保留中文（如" 李四 "→"李四"）
           - 年龄：转为整数，非数字则设为null
           - 手机号：仅保留11位纯数字，不符合则设为null
           - 邮箱：验证格式，无效则设为null
           - 性别：非"男"/"女"则设为"未知"
           - 班级：为空则设为"未分配"
        2. 新增字段：
           - student_id：格式为"2024"+补零后的id（如id=5→"202400005"，id=123→"202400123"）
           - data_quality："完整"（所有字段非null）/"部分缺失"（1-2个字段null）/"严重缺失"（≥3个字段null）
        原始数据：{json.dumps(batch_data, ensure_ascii=False, indent=2)}
        """

        # 构造 LM Studio API 请求体（兼容 OpenAI 格式）
        payload = {
            "model": self.lm_studio_model,
            "messages": [
                {"role": "system", "content": "你是数据整理专家，仅返回JSON格式的处理结果，无其他内容"},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.1,  # 低随机性保证结果稳定
            "max_tokens": 2000   # 足够容纳批量数据的JSON
        }

        try:
            # 调用 LM Studio 本地 API
            response = requests.post(
                self.lm_studio_api_url,
                headers={"Authorization": f"Bearer {self.lm_studio_api_token}"},
                json=payload,
                timeout=60  # 本地模型处理稍慢，延长超时时间
            )
            response.raise_for_status()  # 抛出HTTP错误
            result = response.json()
            
            # 解析本地模型返回的结果
            processed_text = result["choices"][0]["message"]["content"].strip()
            # 清理可能的多余字符（本地模型可能返回```json开头/结尾）
            if processed_text.startswith("```json"):
                processed_text = processed_text.replace("```json", "").replace("```", "").strip()
            processed_data = json.loads(processed_text)
            return processed_data
        except json.JSONDecodeError:
            print(f"本地模型返回非标准JSON，跳过该批：{processed_text[:200]}")
            return []
        except Exception as e:
            print(f"调用LM Studio失败：{str(e)}")
            return []

    def save_to_mongodb(self, processed_data: List[Dict]):
        """将整理后的数据存入MongoDB（支持增量插入）"""
        if not processed_data:
            return
        try:
            # 批量插入，跳过重复id（避免重复存储）
            for data in processed_data:
                self.mongo_collection.update_one(
                    {"id": data["id"]},  # 以原始id为唯一键
                    {"$set": data},      # 存在则更新，不存在则插入
                    upsert=True
                )
            print(f"成功存入/更新 {len(processed_data)} 条数据到MongoDB")
        except Exception as e:
            print(f"存入MongoDB失败：{str(e)}")

    def close_connections(self):
        """关闭数据库连接"""
        if self.mysql_cursor:
            self.mysql_cursor.close()
        if self.mysql_conn:
            self.mysql_conn.close()
        if self.mongo_client:
            self.mongo_client.close()
        print("所有连接已关闭")

    def run(self):
        """执行完整流程：分批读取→本地模型处理→存入MongoDB"""
        try:
            # 遍历所有批次数据
            for batch in self.batch_read_student_data():
                # 处理单批数据
                processed_batch = self.process_batch_with_lmstudio(batch)
                # 存入MongoDB
                self.save_to_mongodb(processed_batch)
            print("✅ 全量数据处理完成！")
        except Exception as e:
            print(f"程序执行异常：{str(e)}")
        finally:
            self.close_connections()

# 程序入口
if __name__ == "__main__":
    processor = StudentDataProcessorLMStudio()
    processor.run()