import pandas as pd
import os
import sys
from datetime import datetime
from sqlalchemy import create_engine, Table, Column, String, DECIMAL, MetaData, text

# 数据库连接配置
DB_CONFIG = {
    'host': '192.168.18.199',
    'port': 3306,
    'user': 'root',
    'password': 'Chlry#$.8',
    'database': 'db_gp_cj'
}

# 获取项目根目录的绝对路径
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

def connect_to_db():
    """连接到数据库"""
    try:
        connection_string = f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
        engine = create_engine(connection_string)
        return engine
    except Exception as e:
        print(f"数据库连接失败: {str(e)}")
        sys.exit(1)

def create_table_if_not_exists(engine):
    """如果不存在则创建表"""
    metadata = MetaData()
    
    # 定义表结构
    main_business = Table(
        'gp_main_business',
        metadata,
        Column('stock_code', String(20), primary_key=True, comment='股票代码'),
        Column('stock_name', String(50), comment='股票简称'),
        Column('report_date', String(8), primary_key=True, comment='报告期(YYYYMMDD)'),
        Column('product_rank', String(1), primary_key=True, comment='项目排名(1-5)'),
        Column('product_name', String(100), comment='项目名称'),
        Column('revenue', DECIMAL(20, 2), comment='项目收入'),
        Column('cost', DECIMAL(20, 2), comment='项目成本'),
        Column('profit', DECIMAL(20, 2), comment='项目毛利'),
        Column('profit_margin', DECIMAL(10, 4), comment='项目毛利率')
    )
    
    # 创建表
    try:
        metadata.create_all(engine)
        print("表创建成功或已存在")
    except Exception as e:
        print(f"创建表失败: {str(e)}")
        sys.exit(1)
        
    return main_business

def clean_numeric_value(value):
    """清理数字字符串，移除逗号并转换为浮点数"""
    if pd.isna(value) or value == "0" or value == "":
        return None
    try:
        # 移除逗号并转换为浮点数
        return float(str(value).replace(',', ''))
    except:
        return None

def import_csv_to_db(csv_file, engine, table, report_date=None):
    """导入CSV文件到数据库"""
    try:
        # 确保使用绝对路径
        if not os.path.isabs(csv_file):
            csv_file = os.path.join(ROOT_DIR, csv_file)
        
        # 检查文件是否存在
        if not os.path.exists(csv_file):
            print(f"文件不存在: {csv_file}")
            return
        print(f"正在处理文件: {csv_file}")
        # 读取CSV文件
        df = pd.read_csv(csv_file, encoding='gbk')
        # 如果没有指定报告期，从文件名或使用当前年份的半年报日期
        if report_date is None:
            filename = os.path.basename(csv_file)
            if len(filename.split('.')[0]) == 8 and filename.split('.')[0].isdigit():
                report_date = filename.split('.')[0]
            else:
                # 使用当前年份的半年报日期（20240630）
                current_year = datetime.now().year
                report_date = f"{current_year}0630"
        
        print(f"处理报告期: {report_date}")
        
        # 连接到数据库
        conn = engine.connect()
        
        # 开始事务
        trans = conn.begin()
        
        try:
            # 先删除该报告期的数据
            delete_query = text(f"DELETE FROM gp_main_business WHERE report_date = '{report_date}'")
            conn.execute(delete_query)
            
            # 列名映射，方便访问列
            column_prefixes = [
                '主营构成(按产品)-第一名-',
                '主营构成(按产品)-第二名-',
                '主营构成(按产品)-第三名-',
                '主营构成(按产品)-第四名-',
                '主营构成(按产品)-第五名-'
            ]
            
            column_suffixes = [
                '项目名称', '项目收入', '项目成本', '项目毛利', '项目毛利率'
            ]
            
            # 插入新数据
            for _, row in df.iterrows():
                stock_code = row['代码']
                stock_name = row['简称']
                
                # 处理前5个主营业务项目
                for i, prefix in enumerate(column_prefixes, 1):
                    # 获取列名
                    name_col = f"{prefix}{column_suffixes[0]}"
                    revenue_col = f"{prefix}{column_suffixes[1]}"
                    cost_col = f"{prefix}{column_suffixes[2]}"
                    profit_col = f"{prefix}{column_suffixes[3]}"
                    margin_col = f"{prefix}{column_suffixes[4]}"
                    
                    # 获取项目名称
                    product_name = row[name_col] if name_col in row else None
                    
                    # 如果没有项目名称或项目名称为"0"，跳过
                    if pd.isna(product_name) or product_name == "0" or product_name == "":
                        continue
                    
                    # 清理并转换数值
                    revenue = clean_numeric_value(row[revenue_col] if revenue_col in row else None)
                    cost = clean_numeric_value(row[cost_col] if cost_col in row else None)
                    profit = clean_numeric_value(row[profit_col] if profit_col in row else None)
                    margin = clean_numeric_value(row[margin_col] if margin_col in row else None)
                    
                    # 准备插入的数据
                    data = {
                        'stock_code': stock_code,
                        'stock_name': stock_name,
                        'report_date': report_date,
                        'product_rank': str(i),
                        'product_name': product_name,
                        'revenue': revenue,
                        'cost': cost,
                        'profit': profit,
                        'profit_margin': margin
                    }
                    
                    # 插入数据
                    insert_query = table.insert().values(**data)
                    conn.execute(insert_query)
            
            # 提交事务
            trans.commit()
            print(f"成功导入 {csv_file} 数据")
            
        except Exception as e:
            # 回滚事务
            trans.rollback()
            print(f"导入数据失败: {str(e)}")
        finally:
            # 关闭连接
            conn.close()
            
    except Exception as e:
        print(f"处理CSV文件 {csv_file} 失败: {str(e)}")

def main():
    """主函数"""
    # 检查命令行参数
    if len(sys.argv) < 2:
        print("用法: python import_main_business.py <csv_file> [report_date]")
        print("例如: python import_main_business.py src/data/111.csv 20240630")
        print("如果不提供report_date，将使用文件名或当前年份的半年报日期")
        sys.exit(1)
        
    csv_file = sys.argv[1]
    
    # 如果提供了报告期参数
    report_date = sys.argv[2] if len(sys.argv) > 2 else None
    
    # 连接数据库
    engine = connect_to_db()
    
    # 创建表
    table = create_table_if_not_exists(engine)
    
    # 导入数据
    import_csv_to_db(csv_file, engine, table, report_date)
    
    print("数据导入完成")

if __name__ == "__main__":
    main() 