ai_mip/import_excel_to_db.py

"""
从Excel文件导入URL数据到数据库
支持开发环境（dev）和生产环境（prod）
"""

import os
import sys
import argparse
from pathlib import Path
from datetime import datetime
from loguru import logger
import pandas as pd
from config import Config
from db_manager import SiteManager

# 配置日志
logger.remove()
logger.add(
    sys.stdout,
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
    level="INFO"
)


class ExcelImporter:
    """Excel数据导入器"""
    
    def __init__(self):
        """初始化导入器"""
        self.config = Config
        self.site_manager = SiteManager()
        
        logger.info("=" * 70)
        logger.info(f"Excel数据导入器已初始化")
        logger.info(f"当前环境: {self.config.ENV}")
        logger.info(f"数据库配置:")
        logger.info(f"  - Host: {self.config.MYSQL_HOST}:{self.config.MYSQL_PORT}")
        logger.info(f"  - Database: {self.config.MYSQL_DATABASE}")
        logger.info(f"  - User: {self.config.MYSQL_USER}")
        logger.info("=" * 70)
        logger.info("提示: 通过设置环境变量 ENV=production 切换到生产环境")
        logger.info("=" * 70)
    
    def read_excel(self, file_path: str) -> pd.DataFrame:
        """
        读取Excel文件
        
        Args:
            file_path: Excel文件路径
            
        Returns:
            DataFrame对象
        """
        try:
            if not Path(file_path).exists():
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            logger.info(f"正在读取Excel文件: {file_path}")
            df = pd.read_excel(file_path)
            
            logger.success(f"成功读取Excel文件，共 {len(df)} 行数据")
            logger.info(f"列名: {df.columns.tolist()}")
            
            return df
            
        except Exception as e:
            logger.error(f"读取Excel文件失败: {str(e)}")
            raise
    
    def validate_dataframe(self, df: pd.DataFrame) -> bool:
        """
        验证DataFrame数据格式
        
        Args:
            df: DataFrame对象
            
        Returns:
            是否验证通过
        """
        # 检查必要的列
        required_columns = ['链接']
        optional_columns = ['序号', '医生', '查询词', '维度']
        
        for col in required_columns:
            if col not in df.columns:
                logger.error(f"缺少必要列: {col}")
                return False
        
        # 检查是否有空链接
        null_count = df['链接'].isna().sum()
        if null_count > 0:
            logger.warning(f"发现 {null_count} 个空链接，将被跳过")
        
        logger.success("数据格式验证通过")
        return True
    
    def import_data(self, df: pd.DataFrame, 
                   query_word: str = None,
                   site_dimension: str = None,
                   frequency: int = 1,
                   time_start: str = '09:00:00',
                   time_end: str = '21:00:00',
                   interval_minutes: int = 30,
                   dry_run: bool = False) -> dict:
        """
        导入数据到数据库
        
        Args:
            df: DataFrame对象
            query_word: 查询词（如果Excel中没有该列，使用此默认值）
            site_dimension: 站点维度（如果Excel中没有该列，使用此默认值）
            frequency: 频次
            time_start: 开始时间
            time_end: 结束时间
            interval_minutes: 执行间隔（分钟）
            dry_run: 是否为试运行模式（不实际插入数据）
            
        Returns:
            导入结果统计
        """
        logger.info("=" * 70)
        logger.info("开始导入数据")
        if dry_run:
            logger.warning("【试运行模式】不会实际插入数据")
        logger.info("=" * 70)
        
        stats = {
            'total': len(df),
            'success': 0,
            'failed': 0,
            'skipped': 0,
            'duplicate': 0
        }
        
        for idx, row in df.iterrows():
            try:
                # 获取链接
                site_url = row.get('链接', None)
                
                # 跳过空链接
                if pd.isna(site_url) or not site_url or str(site_url).strip() == '':
                    logger.warning(f"[{idx + 1}/{stats['total']}] 跳过空链接")
                    stats['skipped'] += 1
                    continue
                
                site_url = str(site_url).strip()
                
                # 获取站点名称（医生名或链接本身）
                site_name = row.get('医生', site_url)
                if pd.isna(site_name):
                    site_name = site_url
                else:
                    site_name = str(site_name).strip()
                
                # 获取查询词
                row_query_word = row.get('查询词', query_word)
                if pd.isna(row_query_word):
                    row_query_word = query_word
                else:
                    row_query_word = str(row_query_word).strip()
                
                # 获取维度
                row_dimension = row.get('维度', site_dimension)
                if pd.isna(row_dimension):
                    row_dimension = site_dimension
                else:
                    row_dimension = str(row_dimension).strip()
                
                logger.info(f"[{idx + 1}/{stats['total']}] 处理: {site_url}")
                logger.info(f"  - 名称: {site_name}")
                logger.info(f"  - 查询词: {row_query_word}")
                logger.info(f"  - 维度: {row_dimension}")
                
                # 试运行模式不实际插入
                if dry_run:
                    logger.info(f"  ✓ [试运行] 跳过插入")
                    stats['success'] += 1
                    continue
                
                # 检查是否已存在
                existing = self.site_manager.get_site_by_url(site_url)
                if existing:
                    logger.warning(f"  × 链接已存在，跳过（ID: {existing['id']}）")
                    stats['duplicate'] += 1
                    continue
                
                # 插入数据库
                site_id = self.site_manager.add_site(
                    site_url=site_url,
                    site_name=site_name,
                    site_dimension=row_dimension,
                    query_word=row_query_word,
                    frequency=frequency,
                    time_start=time_start,
                    time_end=time_end,
                    interval_minutes=interval_minutes
                )
                
                if site_id:
                    logger.success(f"  ✓ 导入成功（ID: {site_id}）")
                    stats['success'] += 1
                else:
                    logger.error(f"  × 导入失败")
                    stats['failed'] += 1
                    
            except Exception as e:
                logger.error(f"[{idx + 1}/{stats['total']}] 处理失败: {str(e)}")
                stats['failed'] += 1
        
        return stats
    
    def print_summary(self, stats: dict):
        """打印导入结果摘要"""
        logger.info("\n" + "=" * 70)
        logger.info("导入完成")
        logger.info("=" * 70)
        logger.info(f"总计: {stats['total']} 条")
        logger.info(f"成功: {stats['success']} 条")
        logger.info(f"失败: {stats['failed']} 条")
        logger.info(f"跳过: {stats['skipped']} 条（空链接）")
        logger.info(f"重复: {stats['duplicate']} 条（已存在）")
        logger.info("=" * 70)
        
        # 计算成功率
        processed = stats['success'] + stats['failed'] + stats['duplicate']
        if processed > 0:
            success_rate = (stats['success'] / processed) * 100
            logger.info(f"成功率: {success_rate:.2f}%")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='从Excel文件导入URL数据到数据库（环境通过ENV环境变量控制）',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 导入到开发环境（默认）
  python import_excel_to_db.py -f "广告链接 1.26（962条）.xlsx"
  
  # 导入到生产环境（设置ENV环境变量）
  $env:ENV="production"; python import_excel_to_db.py -f "广告链接 1.26（962条）.xlsx"
  
  # 指定查询词和维度
  python import_excel_to_db.py -f "广告链接.xlsx" -q "关键词" -d "医疗"
  
  # 试运行模式（不实际插入）
  python import_excel_to_db.py -f "广告链接.xlsx" --dry-run
        """
    )
    
    parser.add_argument('-f', '--file', required=True, help='Excel文件路径')
    parser.add_argument('-q', '--query-word', help='查询词（默认：None）')
    parser.add_argument('-d', '--dimension', help='站点维度（默认：None）')
    parser.add_argument('--frequency', type=int, default=1, help='频次（默认：1）')
    parser.add_argument('--time-start', default='09:00:00', help='开始时间（默认：09:00:00）')
    parser.add_argument('--time-end', default='21:00:00', help='结束时间（默认：21:00:00）')
    parser.add_argument('--interval', type=int, default=30, help='执行间隔分钟数（默认：30）')
    parser.add_argument('--dry-run', action='store_true', help='试运行模式，不实际插入数据')
    
    args = parser.parse_args()
    
    try:
        # 创建导入器
        importer = ExcelImporter()
        
        # 读取Excel
        df = importer.read_excel(args.file)
        
        # 验证数据
        if not importer.validate_dataframe(df):
            logger.error("数据验证失败，终止导入")
            sys.exit(1)
        
        # 显示前几行数据预览
        logger.info("\n数据预览（前5行）:")
        logger.info("\n" + df.head().to_string())
        
        # 确认导入
        if not args.dry_run:
            logger.warning(f"\n即将导入 {len(df)} 条数据到【{importer.config.ENV.upper()}】环境")
            logger.warning(f"数据库: {importer.config.MYSQL_HOST}:{importer.config.MYSQL_PORT}/{importer.config.MYSQL_DATABASE}")
            
            response = input("\n确认继续？[y/N]: ")
            if response.lower() != 'y':
                logger.info("用户取消导入")
                sys.exit(0)
        
        # 导入数据
        stats = importer.import_data(
            df=df,
            query_word=args.query_word,
            site_dimension=args.dimension,
            frequency=args.frequency,
            time_start=args.time_start,
            time_end=args.time_end,
            interval_minutes=args.interval,
            dry_run=args.dry_run
        )
        
        # 打印摘要
        importer.print_summary(stats)
        
        # 根据结果返回退出码
        if stats['failed'] > 0:
            sys.exit(1)
        else:
            sys.exit(0)
            
    except KeyboardInterrupt:
        logger.warning("\n用户中断导入")
        sys.exit(130)
    except Exception as e:
        logger.error(f"导入失败: {str(e)}")
        sys.exit(1)


if __name__ == '__main__':
    main()