"""
独立脚本：从 t_sys_file 导入系统图向量到 Milvus
可以单独运行，不依赖整个项目启动

使用方法：
    python -m app.service.recommendation_system.import_sys_sketch_to_milvus
    或
    python app/service/recommendation_system/import_sys_sketch_to_milvus.py
"""
import sys
import os
import logging
import argparse
from pathlib import Path

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(project_root))

import numpy as np
import pymysql
from tqdm import tqdm

from app.service.recommendation_system.config import (
    MYSQL_CONFIG, TABLE_SYS_FILE,
    RECOMMENDATION_CONFIG, MILVUS_COLLECTION_SKETCH_VECTORS
)
from app.service.recommendation_system.vector_utils import extract_feature_vector, normalize_vector
from app.service.recommendation_system.milvus_client import create_collection, insert_vectors

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('import_sys_sketch.log', encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


def get_sys_file_records(conn, limit=None, offset=0):
    """
    从 t_sys_file 表获取系统图记录
    
    Args:
        conn: 数据库连接
        limit: 限制数量（None 表示不限制）
        offset: 偏移量
    
    Returns:
        记录列表，每个元素为 (id, url, style, level3_type, level2_type, deprecated)
    """
    cursor = conn.cursor()
    
    query = f"""
        SELECT id, url, style, level3_type, level2_type, deprecated 
        FROM {TABLE_SYS_FILE} 
        WHERE level1_type = 'Images' 
          AND style IS NOT NULL 
          AND style != '' 
          AND deprecated != 1
        ORDER BY id
    """
    
    if limit:
        query += f" LIMIT {limit} OFFSET {offset}"
    
    cursor.execute(query)
    records = cursor.fetchall()
    cursor.close()
    
    return records


def get_total_count(conn):
    """获取总记录数"""
    cursor = conn.cursor()
    cursor.execute(f"""
        SELECT COUNT(*) 
        FROM {TABLE_SYS_FILE} 
        WHERE level1_type = 'Images' 
          AND style IS NOT NULL 
          AND style != '' 
          AND deprecated != 1
    """)
    count = cursor.fetchone()[0]
    cursor.close()
    return count


def process_and_insert_batch(records, batch_size=1000, retry_times=3):
    """
    处理并批量插入向量
    
    Args:
        records: 记录列表
        batch_size: 批量大小
        retry_times: 失败重试次数
    
    Returns:
        (成功数量, 失败数量)
    """
    success_count = 0
    failed_count = 0
    failed_records = []
    batch_data = []
    
    # 使用 tqdm 显示进度
    with tqdm(total=len(records), desc="处理记录", unit="条") as pbar:
        for idx, (sys_file_id, url, style, level3_type, level2_type, deprecated) in enumerate(records):
            try:
                # 计算 category
                category = f"{level3_type.lower()}_{level2_type.lower()}"
                
                # 提取特征向量
                feature_vector = extract_feature_vector(url)
                # 归一化，便于 IP≈cosine 度量
                feature_vector = normalize_vector(feature_vector)
                
                # 检查向量是否有效
                if np.all(feature_vector == 0):
                    logger.warning(f"向量提取失败，跳过: {url} (id={sys_file_id})")
                    failed_count += 1
                    failed_records.append((sys_file_id, url))
                    pbar.update(1)
                    continue
                
                # 准备数据
                data_item = {
                    "path": url,
                    "sys_file_id": sys_file_id,
                    "style": style,
                    "category": category,
                    "is_system_sketch": 1,
                    "deprecated": deprecated if deprecated else 0,
                    "feature_vector": feature_vector.tolist()
                }
                
                batch_data.append(data_item)
                
                # 批量写入
                if len(batch_data) >= batch_size:
                    try:
                        insert_vectors(batch_data)
                        success_count += len(batch_data)
                        batch_data = []
                        logger.info(f"已成功插入 {success_count} 条记录")
                    except Exception as e:
                        logger.error(f"批量写入失败: {e}")
                        failed_count += len(batch_data)
                        failed_records.extend([(item["sys_file_id"], item["path"]) for item in batch_data])
                        batch_data = []
                
                pbar.update(1)
                
            except Exception as e:
                logger.error(f"处理记录失败 [id={sys_file_id}, url={url}]: {e}")
                failed_count += 1
                failed_records.append((sys_file_id, url))
                pbar.update(1)
        
        # 写入剩余数据
        if batch_data:
            try:
                insert_vectors(batch_data)
                success_count += len(batch_data)
                logger.info(f"写入剩余 {len(batch_data)} 条记录")
            except Exception as e:
                logger.error(f"写入剩余数据失败: {e}")
                failed_count += len(batch_data)
                failed_records.extend([(item["sys_file_id"], item["path"]) for item in batch_data])
    
    # 重试失败记录
    if failed_records and retry_times > 0:
        logger.info(f"开始重试 {len(failed_records)} 条失败记录，最多重试 {retry_times} 次...")
        
        for retry in range(retry_times):
            if not failed_records:
                break
            
            retry_failed = []
            with tqdm(total=len(failed_records), desc=f"重试第 {retry + 1} 次", unit="条") as pbar:
                for sys_file_id, url in failed_records:
                    try:
                        # 重新查询记录信息
                        conn = pymysql.connect(**MYSQL_CONFIG)
                        cursor = conn.cursor()
                        cursor.execute(f"""
                            SELECT id, url, style, level3_type, level2_type, deprecated 
                            FROM {TABLE_SYS_FILE} 
                            WHERE id = %s
                        """, (sys_file_id,))
                        record = cursor.fetchone()
                        cursor.close()
                        conn.close()
                        
                        if not record:
                            retry_failed.append((sys_file_id, url))
                            pbar.update(1)
                            continue
                        
                        sys_file_id, url, style, level3_type, level2_type, deprecated = record
                        category = f"{level3_type.lower()}_{level2_type.lower()}"
                        
                        feature_vector = extract_feature_vector(url)
                        feature_vector = normalize_vector(feature_vector)
                        if np.all(feature_vector == 0):
                            retry_failed.append((sys_file_id, url))
                            pbar.update(1)
                            continue
                        
                        data_item = {
                            "path": url,
                            "sys_file_id": sys_file_id,
                            "style": style,
                            "category": category,
                            "is_system_sketch": 1,
                            "deprecated": deprecated if deprecated else 0,
                            "feature_vector": feature_vector.tolist()
                        }
                        
                        insert_vectors([data_item])
                        success_count += 1
                        failed_count -= 1
                        pbar.update(1)
                        
                    except Exception as e:
                        logger.error(f"重试失败 [id={sys_file_id}, url={url}]: {e}")
                        retry_failed.append((sys_file_id, url))
                        pbar.update(1)
            
            failed_records = retry_failed
            if failed_records:
                logger.warning(f"第 {retry + 1} 次重试后仍有 {len(failed_records)} 条记录失败")
    
    return success_count, failed_count, failed_records


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='从 t_sys_file 导入系统图向量到 Milvus')
    parser.add_argument('--batch-size', type=int, default=1000, help='批量处理大小（默认：1000）')
    parser.add_argument('--retry-times', type=int, default=3, help='失败重试次数（默认：3）')
    parser.add_argument('--limit', type=int, default=None, help='限制处理数量（用于测试，默认：不限制）')
    parser.add_argument('--offset', type=int, default=0, help='起始偏移量（默认：0）')
    parser.add_argument('--skip-create-collection', action='store_true', help='跳过创建集合（如果集合已存在）')
    
    args = parser.parse_args()
    
    logger.info("=" * 60)
    logger.info("开始从 t_sys_file 导入系统图向量到 Milvus")
    logger.info("=" * 60)
    logger.info(f"配置参数：")
    logger.info(f"  - 批量大小: {args.batch_size}")
    logger.info(f"  - 重试次数: {args.retry_times}")
    logger.info(f"  - 限制数量: {args.limit if args.limit else '不限制'}")
    logger.info(f"  - 起始偏移: {args.offset}")
    logger.info("=" * 60)
    
    # 1. 创建 Milvus 集合
    if not args.skip_create_collection:
        logger.info("创建 Milvus 集合...")
        try:
            create_collection()
            logger.info("Milvus 集合创建成功（或已存在）")
        except Exception as e:
            logger.error(f"创建 Milvus 集合失败: {e}")
            return
    else:
        logger.info("跳过创建集合")
    
    # 2. 连接数据库
    logger.info("连接数据库...")
    try:
        conn = pymysql.connect(**MYSQL_CONFIG)
        logger.info("数据库连接成功")
    except Exception as e:
        logger.error(f"数据库连接失败: {e}")
        return
    
    try:
        # 3. 获取总记录数
        total_count = get_total_count(conn)
        logger.info(f"找到 {total_count} 条系统图记录")
        
        if total_count == 0:
            logger.warning("没有找到系统图数据")
            return
        
        # 4. 获取记录
        logger.info("获取记录...")
        records = get_sys_file_records(conn, limit=args.limit, offset=args.offset)
        logger.info(f"获取到 {len(records)} 条记录")
        
        if not records:
            logger.warning("没有获取到记录")
            return
        
        # 5. 处理并插入
        logger.info("开始处理记录...")
        success_count, failed_count, failed_records = process_and_insert_batch(
            records,
            batch_size=args.batch_size,
            retry_times=args.retry_times
        )
        
        # 6. 输出结果
        logger.info("=" * 60)
        logger.info("导入完成！")
        logger.info(f"  - 成功: {success_count} 条")
        logger.info(f"  - 失败: {failed_count} 条")
        if failed_records:
            logger.warning(f"  - 失败记录列表（前10条）:")
            for sys_file_id, url in failed_records[:10]:
                logger.warning(f"    ID={sys_file_id}, URL={url}")
            if len(failed_records) > 10:
                logger.warning(f"    ... 还有 {len(failed_records) - 10} 条失败记录")
        logger.info("=" * 60)
        
    except Exception as e:
        logger.error(f"处理过程中发生错误: {e}", exc_info=True)
    finally:
        conn.close()
        logger.info("数据库连接已关闭")


if __name__ == "__main__":
    main()

