import logging
from typing import List, Dict, Optional
import pandas as pd
from sqlalchemy import create_engine, text
import os
from openai import OpenAI

# 设置日志记录
logger = logging.getLogger(__name__)

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

# API配置
API_KEY = "28cfe71a-c6fa-4c5d-9b4e-d8474f0d3b93"  # 请填入火山引擎的 API Key
MODEL = "doubao-1-5-pro-32k-250115"    # 请填入火山引擎的模型名称

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

# 初始化OpenAI客户端
client = OpenAI(
    base_url="https://ark.cn-beijing.volces.com/api/v3/",
    api_key=API_KEY
)

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:
        logger.error(f"数据库连接失败: {str(e)}")
        raise

def load_category_list():
    """加载类别列表"""
    try:
        # 使用项目根目录构建文件路径
        file_path = os.path.join(ROOT_DIR, 'data', '字典表.txt')
        with open(file_path, 'r', encoding='utf-8') as f:
            categories = [line.strip() for line in f if line.strip()]
        return categories
    except Exception as e:
        logger.error(f"加载类别列表失败: {str(e)}")
        raise

def create_category_table(engine):
    """创建产品类别表"""
    try:
        with engine.connect() as conn:
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS gp_product_category (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    category_name VARCHAR(100) NOT NULL,
                    product_name VARCHAR(200) NOT NULL,
                    stock_code VARCHAR(20) NOT NULL,
                    stock_name VARCHAR(50) NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    INDEX idx_stock_code (stock_code),
                    INDEX idx_product_name (product_name)
                )
            """))
            conn.commit()
    except Exception as e:
        logger.error(f"创建表失败: {str(e)}")
        raise

def get_company_categorization_prompt(stock_name: str, products: List[str], categories: List[str]) -> str:
    """生成企业产品分类的提示词"""
    prompt = f"""你是一个专业的产品分类专家。请将以下企业的产品名称分类到最合适的类别中。

企业名称：{stock_name}

产品列表：
{chr(10).join(f"- {product}" for product in products)}

可选的类别列表：
{chr(10).join(f"- {category}" for category in categories)}

请为每个产品选择最合适的类别。如果产品名称与多个类别相关，选择最匹配的一个。
如果确实无法确定类别，请返回"其他"。

请按照以下格式返回结果（每行一个产品，用制表符分隔产品名和类别）：
产品名称1    类别1
产品名称2    类别2
...
"""
    return prompt

def categorize_company_products(stock_name: str, products: List[str], categories: List[str]) -> Dict[str, str]:
    """调用API对企业产品进行分类"""
    try:
        # 构建系统提示语
        system_message = {
            "role": "system",
            "content": """你是一个专业的产品分类专家。你的任务是准确地将产品名称分类到预定义的类别中。
            请仔细分析产品名称，考虑其功能、用途和行业特点，选择最匹配的类别。
            如果产品名称与多个类别相关，选择最匹配的一个。
            如果确实无法确定类别，请返回"其他"。
            请按照指定格式返回结果，每行一个产品，用制表符分隔产品名和类别。
            注意：请确保返回的类别名称完全匹配预定义类别列表中的名称。"""
        }
        
        # 构建用户消息
        user_message = {
            "role": "user",
            "content": get_company_categorization_prompt(stock_name, products, categories)
        }
        
        # 调用OpenAI API
        completion = client.chat.completions.create(
            model=MODEL,
            messages=[system_message, user_message],
            temperature=0.3
        )
        
        content = completion.choices[0].message.content.strip()
        
        # 打印原始返回内容以便调试
        print("\nAPI返回的原始内容：")
        print(content)
        
        # 解析返回结果
        category_mapping = {}
        for line in content.split('\n'):
            line = line.strip()
            if not line:  # 跳过空行
                continue
                
            # 尝试不同的分隔符
            if '\t' in line:
                parts = line.split('\t')
            elif ' ' in line:
                # 使用第一个空格分割，保留产品名中的其他空格
                parts = line.split(' ', 1)
            else:
                print(f"警告：无法解析的行: {line}")
                continue
            
            if len(parts) != 2:
                print(f"警告：格式不正确的行: {line}")
                continue
                
            product, category = parts
            product = product.strip()
            category = category.strip()
            
            # 验证类别是否在预定义类别中
            if category not in categories:
                print(f"警告：未知类别 '{category}'，将使用'其他'")
                category = "其他"
            
            # 确保产品名称完整
            if product not in products:
                print(f"警告：产品名称不匹配 '{product}'，将使用原始名称")
                product = next((p for p in products if p in product or product in p), product)
            
            category_mapping[product] = category
        
        if not category_mapping:
            print("警告：没有成功解析任何分类结果")
            return {}
            
        return category_mapping
            
    except Exception as e:
        logger.error(f"产品分类失败: {str(e)}")
        return {}

def process_company_products():
    """处理每个企业的产品并保存分类结果"""
    try:
        # 连接数据库
        engine = connect_to_db()
        
        # 创建类别表
        create_category_table(engine)
        
        # 加载类别列表
        categories = load_category_list()
        
        while True:
            # 获取一个未处理企业的所有产品
            with engine.connect() as conn:
                query = """
                    SELECT DISTINCT m.stock_code, m.stock_name, m.product_name
                    FROM gp_main_business m
                    LEFT JOIN gp_product_category c 
                    ON m.stock_code = c.stock_code 
                    AND m.product_name = c.product_name
                    WHERE c.id IS NULL
                    AND m.product_name IS NOT NULL 
                    AND m.product_name != ''
                    GROUP BY m.stock_code, m.stock_name
                    LIMIT 1
                """
                company = pd.read_sql(query, conn)
                
                if company.empty:
                    print("所有产品都已处理完成！")
                    break
                
                stock_code = company['stock_code'].iloc[0]
                stock_name = company['stock_name'].iloc[0]
                
                # 获取该企业的所有未处理产品
                query = f"""
                    SELECT DISTINCT product_name
                    FROM gp_main_business
                    WHERE stock_code = '{stock_code}'
                    AND product_name NOT IN (
                        SELECT product_name 
                        FROM gp_product_category 
                        WHERE stock_code = '{stock_code}'
                    )
                    AND product_name IS NOT NULL 
                    AND product_name != ''
                """
                products = pd.read_sql(query, conn)
                
                if products.empty:
                    continue
                
                product_list = products['product_name'].tolist()
                
                # 调用API进行分类
                print(f"\n处理企业：{stock_name} ({stock_code})")
                print(f"产品数量：{len(product_list)}")
                print("\n产品列表：")
                for product in product_list:
                    print(f"- {product}")
                
                category_mapping = categorize_company_products(stock_name, product_list, categories)
                
                if not category_mapping:
                    print("\n警告：未能获取分类结果，将使用'其他'类别")
                    category_mapping = {product: "其他" for product in product_list}
                
                # 显示分类结果
                print("\n分类结果：")
                for product, category in category_mapping.items():
                    print(f"{product} -> {category}")
                
                # 自动保存分类结果
                with engine.connect() as conn:
                    for product, category in category_mapping.items():
                        insert_query = """
                            INSERT INTO gp_product_category 
                            (category_name, product_name, stock_code, stock_name) 
                            VALUES (:category_name, :product_name, :stock_code, :stock_name)
                        """
                        conn.execute(text(insert_query), 
                                   {"category_name": category, 
                                    "product_name": product, 
                                    "stock_code": stock_code,
                                    "stock_name": stock_name})
                    conn.commit()
                print(f"已保存 {stock_name} 的分类结果")
            
    except Exception as e:
        logger.error(f"处理产品失败: {str(e)}")
        raise

def update_product_categories_by_rules(category_rules: List[Dict[str, str]]):
    """
    根据给定的规则批量更新产品分类
    
    Args:
        category_rules: 分类规则列表，每个规则是一个字典，格式为：
            {
                "category": "类别名称",
                "products": "产品1,产品2,产品3,..."
            }
    """
    try:
        # 连接数据库
        engine = connect_to_db()
        
        with engine.connect() as conn:
            for rule in category_rules:
                category = rule["category"]
                products = [p.strip() for p in rule["products"].split(",")]
                
                # 构建 IN 查询条件
                products_str = "','".join(products)
                update_query = f"""
                    UPDATE gp_product_category 
                    SET category_name = '{category}'
                    WHERE product_name IN ('{products_str}')
                """
                
                # 执行更新
                result = conn.execute(text(update_query))
                conn.commit()
                
                print(f"已更新 {result.rowcount} 条记录到类别 '{category}'")
                
    except Exception as e:
        logger.error(f"批量更新产品分类失败: {str(e)}")
        raise

if __name__ == "__main__":
    # 设置日志级别
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 示例：使用规则更新分类
    category_rules = [
            {
                "category": "其他汽车零部件",
                "num": 25,
                "products": "其他汽车零部件,汽车零部件制造,汽车检具,工装模具及汽车零部件,汽车应用类,缓速器,汽车件,汽车产品,汽车类零部件,汽车及工业产品,汽车部品,汽车、轨道交通构件,功能件,冲压业务,汽车零部件收入"
            }
    ]
    
    # 更新分类
    # update_product_categories_by_rules(category_rules)
    
    # 处理产品
    process_company_products()