#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
任务二: 数据导入脚本
将清洗后的数据导入MySQL数据库
"""

import pandas as pd
import pymysql
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

print("=" * 80)
print("跨境电商数据分析竞赛 - 任务二:数据导入")
print("=" * 80)

# 数据库连接配置(请根据实际情况修改)
DB_CONFIG = {
    'host': 'localhost',
    'user': 'root',
    'password': 'your_password',  # 请修改为实际密码
    'database': 'ecommerce_analysis',
    'charset': 'utf8mb4'
}

def create_connection():
    """创建数据库连接"""
    try:
        connection = pymysql.connect(**DB_CONFIG)
        print(f"✓ 成功连接到MySQL数据库")
        return connection
    except Exception as e:
        print(f"✗ 数据库连接失败: {e}")
        return None

def load_cleaned_data():
    """加载清洗后的数据"""
    print(f"\n【步骤1】加载清洗后的数据")
    print("-" * 80)
    
    df = pd.read_csv('../data/ecommerce_sales_cleaned.csv', encoding='utf-8-sig')
    df['order_date'] = pd.to_datetime(df['order_date'])
    df['product_launch_date'] = pd.to_datetime(df['product_launch_date'])
    
    print(f"✓ 数据加载完成: {len(df)} 条记录")
    return df

def populate_dim_product(connection, df):
    """填充产品维度表"""
    print(f"\n【步骤2】填充产品维度表 (dim_product)")
    print("-" * 80)
    
    cursor = connection.cursor()
    
    # 提取唯一产品
    products = df[['product_id', 'product_name', 'category', 'sub_category', 
                   'product_launch_date', 'supplier_id']].drop_duplicates()
    
    insert_sql = """
    INSERT INTO dim_product 
    (product_id, product_name, category, sub_category, product_launch_date, supplier_id)
    VALUES (%s, %s, %s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE 
        product_name = VALUES(product_name),
        category = VALUES(category),
        sub_category = VALUES(sub_category)
    """
    
    data = []
    for _, row in products.iterrows():
        data.append((
            row['product_id'],
            row['product_name'],
            row['category'],
            row['sub_category'],
            row['product_launch_date'].strftime('%Y-%m-%d'),
            row['supplier_id']
        ))
    
    cursor.executemany(insert_sql, data)
    connection.commit()
    print(f"✓ 插入产品记录: {len(data)} 条")
    cursor.close()

def populate_dim_region(connection, df):
    """填充区域维度表"""
    print(f"\n【步骤3】填充区域维度表 (dim_region)")
    print("-" * 80)
    
    cursor = connection.cursor()
    
    # 提取唯一区域和国家组合
    regions = df[['region', 'country']].drop_duplicates()
    
    insert_sql = """
    INSERT INTO dim_region (region, country)
    VALUES (%s, %s)
    ON DUPLICATE KEY UPDATE region = VALUES(region)
    """
    
    data = [(row['region'], row['country']) for _, row in regions.iterrows()]
    
    cursor.executemany(insert_sql, data)
    connection.commit()
    print(f"✓ 插入区域记录: {len(data)} 条")
    cursor.close()

def populate_dim_time(connection, df):
    """填充时间维度表"""
    print(f"\n【步骤4】填充时间维度表 (dim_time)")
    print("-" * 80)
    
    cursor = connection.cursor()
    
    # 提取所有唯一日期
    dates = pd.concat([df['order_date'], df['product_launch_date']]).unique()
    dates = pd.to_datetime(dates).sort_values()
    
    insert_sql = """
    INSERT INTO dim_time 
    (date_id, date, year, quarter, month, week, day_of_week, is_weekend)
    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE date = VALUES(date)
    """
    
    data = []
    for date in dates:
        date_id = int(date.strftime('%Y%m%d'))
        year = date.year
        quarter = (date.month - 1) // 3 + 1
        month = date.month
        week = date.isocalendar()[1]
        day_of_week = date.weekday() + 1  # 1-7
        is_weekend = day_of_week in [6, 7]
        
        data.append((date_id, date.strftime('%Y-%m-%d'), year, quarter, 
                    month, week, day_of_week, is_weekend))
    
    cursor.executemany(insert_sql, data)
    connection.commit()
    print(f"✓ 插入时间记录: {len(data)} 条")
    cursor.close()

def populate_dim_customer(connection, df):
    """填充客户维度表"""
    print(f"\n【步骤5】填充客户维度表 (dim_customer)")
    print("-" * 80)
    
    cursor = connection.cursor()
    
    # 计算每个客户的统计信息
    customer_stats = df.groupby('customer_id').agg({
        'order_date': 'min',  # 首次购买日期
        'order_id': 'count',  # 总订单数
        'total_amount': 'sum'  # 总消费金额
    }).reset_index()
    
    customer_stats.columns = ['customer_id', 'first_purchase_date', 'total_orders', 'total_amount']
    
    insert_sql = """
    INSERT INTO dim_customer 
    (customer_id, first_purchase_date, total_orders, total_amount)
    VALUES (%s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE 
        first_purchase_date = VALUES(first_purchase_date),
        total_orders = VALUES(total_orders),
        total_amount = VALUES(total_amount)
    """
    
    data = []
    for _, row in customer_stats.iterrows():
        data.append((
            row['customer_id'],
            row['first_purchase_date'].strftime('%Y-%m-%d'),
            int(row['total_orders']),
            float(row['total_amount'])
        ))
    
    cursor.executemany(insert_sql, data)
    connection.commit()
    print(f"✓ 插入客户记录: {len(data)} 条")
    cursor.close()

def populate_fact_sales(connection, df):
    """填充销售事实表"""
    print(f"\n【步骤6】填充销售事实表 (fact_sales)")
    print("-" * 80)
    
    cursor = connection.cursor()
    
    # 获取region_id映射
    cursor.execute("SELECT region_id, region, country FROM dim_region")
    region_map = {}
    for region_id, region, country in cursor.fetchall():
        region_map[(region, country)] = region_id
    
    # 计算利润
    df['profit'] = (df['unit_price'] - df['cost_price']) * df['quantity'] * (1 - df['discount_rate'])
    
    insert_sql = """
    INSERT INTO fact_sales 
    (order_id, product_id, customer_id, region_id, order_date_id, order_date,
     quantity, unit_price, cost_price, discount_rate, total_amount, 
     shipping_cost, delivery_days, is_returned, stock_quantity, profit)
    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE order_id = VALUES(order_id)
    """
    
    data = []
    for _, row in df.iterrows():
        region_id = region_map.get((row['region'], row['country']))
        if region_id is None:
            continue
        
        order_date_id = int(row['order_date'].strftime('%Y%m%d'))
        
        data.append((
            row['order_id'],
            row['product_id'],
            row['customer_id'],
            region_id,
            order_date_id,
            row['order_date'].strftime('%Y-%m-%d'),
            int(row['quantity']),
            float(row['unit_price']),
            float(row['cost_price']),
            float(row['discount_rate']),
            float(row['total_amount']),
            float(row['shipping_cost']),
            int(row['delivery_days']),
            bool(row['is_returned']),
            int(row['stock_quantity']),
            float(row['profit'])
        ))
    
    # 分批插入
    batch_size = 1000
    for i in range(0, len(data), batch_size):
        batch = data[i:i+batch_size]
        cursor.executemany(insert_sql, batch)
        connection.commit()
        print(f"  已插入: {min(i+batch_size, len(data))}/{len(data)} 条", end='\r')
    
    print(f"\n✓ 插入销售记录: {len(data)} 条")
    cursor.close()

def verify_data(connection):
    """验证数据导入"""
    print(f"\n【步骤7】验证数据导入")
    print("-" * 80)
    
    cursor = connection.cursor()
    
    tables = ['dim_product', 'dim_region', 'dim_time', 'dim_customer', 'fact_sales']
    
    for table in tables:
        cursor.execute(f"SELECT COUNT(*) FROM {table}")
        count = cursor.fetchone()[0]
        print(f"✓ {table:20s}: {count:8d} 条记录")
    
    cursor.close()

def main():
    """主函数"""
    # 加载数据
    df = load_cleaned_data()
    
    # 连接数据库
    connection = create_connection()
    if connection is None:
        print("\n✗ 无法连接到数据库,请检查配置")
        print("  提示: 请先安装MySQL并修改DB_CONFIG中的连接信息")
        print("  提示: 需要先执行 sql/01_create_tables.sql 创建数据库和表")
        return
    
    try:
        # 依次填充各表
        populate_dim_product(connection, df)
        populate_dim_region(connection, df)
        populate_dim_time(connection, df)
        populate_dim_customer(connection, df)
        populate_fact_sales(connection, df)
        
        # 验证数据
        verify_data(connection)
        
        print("\n" + "=" * 80)
        print("数据导入完成!")
        print("=" * 80)
        print("\n提示: 可以使用 sql/03_business_queries.sql 中的查询语句进行数据分析")
        
    except Exception as e:
        print(f"\n✗ 数据导入失败: {e}")
    finally:
        if connection:
            connection.close()
            print("\n✓ 数据库连接已关闭")

if __name__ == '__main__':
    main()
