"""
销售数据生成与处理脚本

功能：生成测试销售数据并写入目标表 test_sales_data
目标表：test_sales_data（销售数据测试表）
更新模式：append（追加）

作者：cursor
创建时间：2025-12-31
"""

import argparse
import logging
import os
import random
import sys
from datetime import datetime, timedelta
from decimal import Decimal
from typing import Any, Dict, List, Optional

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker

# 添加项目根目录到路径
sys.path.insert(
    0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
)

try:
    from app.config.config import config, current_env  # type: ignore

    # 根据当前环境获取配置类
    Config = config.get(current_env, config["default"])
except ImportError:

    class Config:  # type: ignore
        SQLALCHEMY_DATABASE_URI = os.environ.get(
            "DATABASE_URI", "postgresql://user:password@localhost:5432/database"
        )


# 配置日志
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


# 测试数据配置
PRODUCTS = [
    {"id": "P001", "name": "笔记本电脑", "category": "电子产品", "base_price": 5999.00},
    {"id": "P002", "name": "无线鼠标", "category": "电子产品", "base_price": 129.00},
    {"id": "P003", "name": "机械键盘", "category": "电子产品", "base_price": 399.00},
    {"id": "P004", "name": "显示器", "category": "电子产品", "base_price": 1599.00},
    {"id": "P005", "name": "办公椅", "category": "办公用品", "base_price": 899.00},
    {"id": "P006", "name": "办公桌", "category": "办公用品", "base_price": 1299.00},
    {"id": "P007", "name": "文件柜", "category": "办公用品", "base_price": 599.00},
    {"id": "P008", "name": "打印机", "category": "办公设备", "base_price": 1899.00},
    {"id": "P009", "name": "投影仪", "category": "办公设备", "base_price": 3999.00},
    {"id": "P010", "name": "白板", "category": "办公用品", "base_price": 299.00},
]

CUSTOMERS = [
    {"id": "C001", "name": "张三"},
    {"id": "C002", "name": "李四"},
    {"id": "C003", "name": "王五"},
    {"id": "C004", "name": "赵六"},
    {"id": "C005", "name": "钱七"},
    {"id": "C006", "name": "孙八"},
    {"id": "C007", "name": "周九"},
    {"id": "C008", "name": "吴十"},
    {"id": "C009", "name": "郑十一"},
    {"id": "C010", "name": "王十二"},
]

REGIONS = {
    "华东": ["上海", "杭州", "南京", "苏州", "无锡"],
    "华北": ["北京", "天津", "石家庄", "太原", "济南"],
    "华南": ["广州", "深圳", "珠海", "东莞", "佛山"],
    "华中": ["武汉", "长沙", "郑州", "南昌", "合肥"],
    "西南": ["成都", "重庆", "昆明", "贵阳", "南宁"],
}

PAYMENT_METHODS = ["现金", "信用卡", "支付宝", "微信支付", "银行转账"]

ORDER_STATUSES = ["已完成", "已发货", "处理中", "待付款", "已取消"]


class SalesDataGenerator:
    """销售数据生成器"""

    TARGET_TABLE = "test_sales_data"
    TARGET_SCHEMA = "public"

    def __init__(
        self,
        target_schema: str = "public",
        update_mode: str = "append",
        batch_size: int = 100,
    ):
        """
        初始化生成器

        Args:
            target_schema: 目标表schema
            update_mode: 更新模式，'append'（追加）或 'full'（全量更新）
            batch_size: 批量处理大小
        """
        self.target_schema = target_schema
        self.update_mode = update_mode.lower()
        self.batch_size = batch_size

        self.engine: Optional[Engine] = None
        self.session: Optional[Session] = None

        self.generated_count = 0
        self.inserted_count = 0
        self.error_count = 0

        if self.update_mode not in ["append", "full"]:
            raise ValueError(
                f"不支持的更新模式: {update_mode}，仅支持 'append' 或 'full'"
            )

        logger.info(
            f"初始化销售数据生成器: "
            f"目标表={target_schema}.{self.TARGET_TABLE}, "
            f"更新模式={update_mode}"
        )

    def connect_database(self) -> bool:
        """连接数据库"""
        try:
            db_uri = Config.SQLALCHEMY_DATABASE_URI

            if not db_uri:
                logger.error("未找到数据库配置（SQLALCHEMY_DATABASE_URI）")
                return False

            self.engine = create_engine(db_uri)
            SessionLocal = sessionmaker(bind=self.engine)
            self.session = SessionLocal()

            with self.engine.connect() as conn:
                conn.execute(text("SELECT 1"))

            logger.info(f"成功连接数据库: {db_uri.split('@')[-1]}")
            return True

        except Exception as e:
            logger.error(f"连接数据库失败: {str(e)}")
            return False

    def check_table_exists(self) -> bool:
        """检查目标表是否存在"""
        try:
            if not self.engine:
                return False

            inspector = inspect(self.engine)
            tables = inspector.get_table_names(schema=self.target_schema)
            return self.TARGET_TABLE in tables

        except Exception as e:
            logger.error(f"检查表是否存在失败: {str(e)}")
            return False

    def create_target_table(self) -> bool:
        """创建目标表"""
        try:
            if self.check_table_exists():
                logger.info(f"目标表 {self.target_schema}.{self.TARGET_TABLE} 已存在")
                return True

            if not self.session:
                logger.error("数据库会话未初始化")
                return False

            create_sql = text(f"""
                CREATE TABLE IF NOT EXISTS {self.target_schema}.{self.TARGET_TABLE} (
                    order_id VARCHAR(50),
                    order_date DATE,
                    customer_id VARCHAR(50),
                    customer_name VARCHAR(100),
                    product_id VARCHAR(50),
                    product_name VARCHAR(200),
                    category VARCHAR(100),
                    quantity INTEGER,
                    unit_price NUMERIC(10, 2),
                    total_amount NUMERIC(12, 2),
                    discount_rate NUMERIC(5, 2),
                    payment_method VARCHAR(50),
                    region VARCHAR(100),
                    city VARCHAR(100),
                    status VARCHAR(50),
                    created_at TIMESTAMP,
                    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            self.session.execute(create_sql)

            # 添加表注释
            comment_sql = text(
                f"COMMENT ON TABLE {self.target_schema}.{self.TARGET_TABLE} "
                f"IS 'Sales data table - test data'"
            )
            self.session.execute(comment_sql)

            self.session.commit()
            logger.info(f"成功创建目标表 {self.target_schema}.{self.TARGET_TABLE}")
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"创建目标表失败: {str(e)}")
            return False

    def generate_order_id(self, index: int) -> str:
        """生成订单ID"""
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        return f"ORD{timestamp}{index:04d}"

    def generate_sales_data(
        self,
        count: int = 100,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        生成测试销售数据

        Args:
            count: 生成数据条数
            start_date: 订单开始日期
            end_date: 订单结束日期

        Returns:
            生成的销售数据列表
        """
        # 解析日期范围
        if start_date:
            start = datetime.strptime(start_date, "%Y-%m-%d")
        else:
            start = datetime.now() - timedelta(days=30)

        if end_date:
            end = datetime.strptime(end_date, "%Y-%m-%d")
        else:
            end = datetime.now()

        date_range = (end - start).days
        if date_range <= 0:
            date_range = 1

        data_rows = []

        for i in range(count):
            # 随机选择产品、客户、地区
            product = random.choice(PRODUCTS)
            customer = random.choice(CUSTOMERS)
            region = random.choice(list(REGIONS.keys()))
            city = random.choice(REGIONS[region])

            # 生成订单数据
            order_date = start + timedelta(days=random.randint(0, date_range))
            quantity = random.randint(1, 10)
            unit_price = Decimal(str(product["base_price"]))

            # 随机折扣 0-20%
            discount_rate = Decimal(str(random.randint(0, 20))) / 100
            total_amount = unit_price * quantity * (1 - discount_rate)

            row = {
                "order_id": self.generate_order_id(i),
                "order_date": order_date.date(),
                "customer_id": customer["id"],
                "customer_name": customer["name"],
                "product_id": product["id"],
                "product_name": product["name"],
                "category": product["category"],
                "quantity": quantity,
                "unit_price": float(unit_price),
                "total_amount": float(total_amount.quantize(Decimal("0.01"))),
                "discount_rate": float(discount_rate),
                "payment_method": random.choice(PAYMENT_METHODS),
                "region": region,
                "city": city,
                "status": random.choice(ORDER_STATUSES),
                "created_at": datetime.now(),
            }
            data_rows.append(row)

        self.generated_count = len(data_rows)
        logger.info(f"成功生成 {self.generated_count} 条销售数据")
        return data_rows

    def clear_target_table(self) -> bool:
        """清空目标表"""
        try:
            if not self.session:
                logger.error("数据库会话未初始化")
                return False

            delete_sql = text(f"DELETE FROM {self.target_schema}.{self.TARGET_TABLE}")
            self.session.execute(delete_sql)
            self.session.commit()

            logger.info(f"目标表 {self.target_schema}.{self.TARGET_TABLE} 已清空")
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"清空目标表失败: {str(e)}")
            return False

    def insert_data(self, data_rows: List[Dict[str, Any]]) -> bool:
        """插入数据到目标表"""
        try:
            if not data_rows:
                logger.warning("没有数据需要插入")
                return True

            if not self.session:
                logger.error("数据库会话未初始化")
                return False

            # 全量更新模式：先清空目标表
            if self.update_mode == "full":
                if not self.clear_target_table():
                    return False

            # 构建插入 SQL（适应现有表结构，不包含 create_time）
            columns = [
                "order_id",
                "order_date",
                "customer_id",
                "customer_name",
                "product_id",
                "product_name",
                "category",
                "quantity",
                "unit_price",
                "total_amount",
                "discount_rate",
                "payment_method",
                "region",
                "city",
                "status",
                "created_at",
            ]

            columns_str = ", ".join(columns)
            placeholders = ", ".join([f":{col}" for col in columns])

            insert_sql = text(f"""
                INSERT INTO {self.target_schema}.{self.TARGET_TABLE} ({columns_str})
                VALUES ({placeholders})
            """)

            # 批量插入
            success_count = 0
            for i, row in enumerate(data_rows):
                try:
                    self.session.execute(insert_sql, row)
                    success_count += 1

                    if success_count % self.batch_size == 0:
                        self.session.commit()
                        logger.info(f"已插入 {success_count} 条数据...")

                except Exception as e:
                    self.error_count += 1
                    logger.error(f"插入数据失败 (行 {i}): {str(e)}")

            self.session.commit()
            self.inserted_count = success_count

            logger.info(
                f"数据插入完成: 成功 {self.inserted_count} 条, 失败 {self.error_count} 条"
            )
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"批量插入数据失败: {str(e)}")
            return False

    def close_connection(self):
        """关闭数据库连接"""
        if self.session:
            try:
                self.session.close()
                logger.info("数据库会话已关闭")
            except Exception as e:
                logger.error(f"关闭数据库会话失败: {str(e)}")

        if self.engine:
            try:
                self.engine.dispose()
                logger.info("数据库引擎已释放")
            except Exception as e:
                logger.error(f"释放数据库引擎失败: {str(e)}")

    def run(
        self,
        count: int = 100,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        执行数据生成流程

        Args:
            count: 生成数据条数
            start_date: 订单开始日期
            end_date: 订单结束日期

        Returns:
            执行结果
        """
        result = {
            "success": False,
            "generated_count": 0,
            "inserted_count": 0,
            "error_count": 0,
            "update_mode": self.update_mode,
            "message": "",
            "start_time": datetime.now().isoformat(),
            "end_time": None,
        }

        try:
            logger.info("=" * 60)
            logger.info("开始销售数据生成")
            logger.info(f"目标表: {self.target_schema}.{self.TARGET_TABLE}")
            logger.info(f"生成数量: {count}")
            logger.info(f"更新模式: {self.update_mode}")
            logger.info("=" * 60)

            # 1. 连接数据库
            if not self.connect_database():
                result["message"] = "连接数据库失败"
                return result

            # 2. 创建目标表
            if not self.create_target_table():
                result["message"] = "创建目标表失败"
                return result

            # 3. 生成测试数据
            data_rows = self.generate_sales_data(
                count=count,
                start_date=start_date,
                end_date=end_date,
            )

            # 4. 插入数据
            if self.insert_data(data_rows):
                result["success"] = True
                result["generated_count"] = self.generated_count
                result["inserted_count"] = self.inserted_count
                result["error_count"] = self.error_count
                result["message"] = (
                    f"处理完成: 生成 {self.generated_count} 条, "
                    f"插入 {self.inserted_count} 条, "
                    f"失败 {self.error_count} 条"
                )
            else:
                result["message"] = "插入数据失败"

        except Exception as e:
            logger.error(f"处理过程发生异常: {str(e)}")
            result["message"] = f"处理失败: {str(e)}"

        finally:
            result["end_time"] = datetime.now().isoformat()
            self.close_connection()

        logger.info("=" * 60)
        logger.info(f"处理结果: {result['message']}")
        logger.info("=" * 60)

        return result


def generate_sales_data(
    target_schema: str = "public",
    update_mode: str = "append",
    count: int = 100,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
) -> Dict[str, Any]:
    """
    生成销售数据（入口函数）

    Args:
        target_schema: 目标表schema
        update_mode: 更新模式
        count: 生成数据条数
        start_date: 订单开始日期
        end_date: 订单结束日期

    Returns:
        处理结果
    """
    generator = SalesDataGenerator(
        target_schema=target_schema,
        update_mode=update_mode,
    )
    return generator.run(
        count=count,
        start_date=start_date,
        end_date=end_date,
    )


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="销售数据生成工具")

    parser.add_argument(
        "--target-schema",
        type=str,
        default="public",
        help="目标表schema（默认：public）",
    )

    parser.add_argument(
        "--update-mode",
        type=str,
        choices=["append", "full"],
        default="append",
        help="更新模式：append（追加）或 full（全量更新），默认：append",
    )

    parser.add_argument(
        "--count",
        type=int,
        default=100,
        help="生成数据条数（默认：100）",
    )

    parser.add_argument(
        "--start-date",
        type=str,
        default=None,
        help="订单开始日期，格式：YYYY-MM-DD",
    )

    parser.add_argument(
        "--end-date",
        type=str,
        default=None,
        help="订单结束日期，格式：YYYY-MM-DD",
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    result = generate_sales_data(
        target_schema=args.target_schema,
        update_mode=args.update_mode,
        count=args.count,
        start_date=args.start_date,
        end_date=args.end_date,
    )

    # 输出结果
    print("\n" + "=" * 60)
    print(f"处理结果: {'成功' if result['success'] else '失败'}")
    print(f"消息: {result['message']}")
    print(f"生成: {result['generated_count']} 条")
    print(f"插入: {result['inserted_count']} 条")
    print(f"失败: {result['error_count']} 条")
    print(f"更新模式: {result['update_mode']}")
    print(f"开始时间: {result['start_time']}")
    print(f"结束时间: {result['end_time']}")
    print("=" * 60)

    exit(0 if result["success"] else 1)
