"""基于规则的数据生成器 - 执行 LLM 生成的规则"""

import random
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional
import csv
from io import StringIO


class RuleBasedGenerator:
    """基于规则的数据生成器"""

    # 中文姓名库（扩展版）
    CHINESE_NAMES = [
        '张明', '李华', '王芳', '刘强', '陈静', '赵伟', '孙丽', '周杰', '吴敏', '郑涛',
        '客服帅哥', '石蕲', '郝建军', '王加加', '张积海', '闫旭啸', '胡太斯', '耿奇轩',
        '小梦', '小雪', '巴兰', '李坤', '楼越', '马云', '刘备', '关羽', '张飞', '赵云',
        '诸葛亮', '孙权', '周瑜', '鲁肃', '黄盖', '甘宁', '太史慈', '吕蒙', '陆逊', '姜维',
        '魏延', '马超', '黄忠', '严颜', '张辽', '徐晃', '张郃', '曹仁', '夏侯惇', '夏侯渊',
        '许褚', '典韦', '庞德', '于禁', '李典', '乐进', '文聘', '曹洪', '曹真', '司马懿',
        '王莽', '刘秀', '曹操', '孙策', '袁绍', '吕布', '董卓', '貂蝉', '大乔', '小乔'
    ]

    def __init__(self, seed: int = 42):
        self.seed = seed
        random.seed(seed)
        # 存储已生成的表数据，用于外键引用
        self.generated_tables: Dict[str, List[Dict[str, Any]]] = {}

    def generate_all_tables(self, rules: Dict[str, Any]) -> Dict[str, List[Dict[str, Any]]]:
        """按依赖顺序生成所有表的数据

        Args:
            rules: 所有表的生成规则

        Returns:
            表名到数据行列表的映射
        """
        tables_rules = rules.get('tables', {})

        # 拓扑排序：确定表的生成顺序
        table_order = self._topological_sort(tables_rules)

        # 按顺序生成每个表
        result = {}
        for table_name in table_order:
            table_rules = tables_rules[table_name]
            table_data = self.generate_table_data(table_name, table_rules)
            result[table_name] = table_data
            self.generated_tables[table_name] = table_data

        return result

    def _topological_sort(self, tables_rules: Dict[str, Dict[str, Any]]) -> List[str]:
        """拓扑排序：根据外键依赖确定表的生成顺序

        Args:
            tables_rules: 所有表的规则

        Returns:
            排序后的表名列表
        """
        # 构建依赖图
        dependencies = {}  # table -> [依赖的表列表]
        for table_name, table_rule in tables_rules.items():
            deps = set()
            for field_name, field_rule in table_rule.get('fields', {}).items():
                if field_rule.get('type') == 'foreign_key':
                    ref_table = field_rule.get('reference_table')
                    if ref_table and ref_table in tables_rules:
                        deps.add(ref_table)
            dependencies[table_name] = list(deps)

        # Kahn 算法拓扑排序
        in_degree = {table: 0 for table in tables_rules}
        for table, deps in dependencies.items():
            for dep in deps:
                in_degree[dep] = in_degree.get(dep, 0)

        for table, deps in dependencies.items():
            for dep in deps:
                in_degree[table] += 1

        queue = [table for table, degree in in_degree.items() if degree == 0]
        result = []

        while queue:
            table = queue.pop(0)
            result.append(table)

            # 减少依赖此表的其他表的入度
            for other_table, deps in dependencies.items():
                if table in deps:
                    in_degree[other_table] -= 1
                    if in_degree[other_table] == 0:
                        queue.append(other_table)

        # 检查是否有循环依赖
        if len(result) != len(tables_rules):
            # 有循环依赖，返回原始顺序
            print("Warning: Circular dependency detected, using original order")
            return list(tables_rules.keys())

        return result

    def generate_table_data(self, table_name: str, rules: Dict[str, Any]) -> List[Dict[str, Any]]:
        """根据规则生成表数据

        Args:
            table_name: 表名
            rules: 表的生成规则

        Returns:
            数据行列表
        """
        num_rows = rules.get('num_rows', 20)
        field_rules = rules.get('fields', {})

        rows = []
        for i in range(num_rows):
            row = {}
            for field_name, field_rule in field_rules.items():
                row[field_name] = self._generate_field_value(
                    field_name, field_rule, i, num_rows
                )
            rows.append(row)

        return rows

    def _generate_field_value(
        self, field_name: str, rule: Dict[str, Any], index: int, total_rows: int
    ) -> str:
        """根据规则生成字段值"""
        field_type = rule.get('type', 'default')

        if field_type == 'chinese_name':
            return self._generate_chinese_name(index, rule.get('unique', False))

        elif field_type == 'normal_distribution':
            return self._generate_normal_distribution(rule, index)

        elif field_type == 'uniform':
            return self._generate_uniform(rule)

        elif field_type == 'date_range':
            return self._generate_date_range(rule, index, total_rows)

        elif field_type == 'choice':
            return self._generate_choice(rule)

        elif field_type == 'sequential':
            return self._generate_sequential(rule, index)

        elif field_type == 'foreign_key':
            return self._generate_foreign_key(rule)

        else:
            return f"value_{index}"

    def _generate_chinese_name(self, index: int, unique: bool) -> str:
        """生成中文姓名"""
        if unique and index < len(self.CHINESE_NAMES):
            return self.CHINESE_NAMES[index]
        elif unique:
            return f"学生{index + 1:03d}"
        else:
            return random.choice(self.CHINESE_NAMES)

    def _generate_normal_distribution(self, rule: Dict[str, Any], index: int) -> str:
        """生成正态分布数值"""
        mean = rule.get('mean', 50)
        std = rule.get('std', 10)
        min_val = rule.get('min', 0)
        max_val = rule.get('max', 100)

        # 使用确定性随机（基于 index）
        random.seed(self.seed + index)
        value = random.gauss(mean, std)
        value = max(min_val, min(max_val, value))

        # 如果是整数范围，返回整数
        if isinstance(min_val, int) and isinstance(max_val, int):
            return str(int(value))
        else:
            return f"{value:.2f}"

    def _generate_uniform(self, rule: Dict[str, Any]) -> str:
        """生成均匀分布数值"""
        min_val = rule.get('min', 0)
        max_val = rule.get('max', 100)

        if isinstance(min_val, int) and isinstance(max_val, int):
            return str(random.randint(min_val, max_val))
        else:
            return f"{random.uniform(min_val, max_val):.2f}"

    def _generate_date_range(
        self, rule: Dict[str, Any], index: int, total_rows: int
    ) -> str:
        """生成日期范围内的日期"""
        start_str = rule.get('start', '2025-01-01')
        end_str = rule.get('end', '2025-12-31')
        date_format = rule.get('format', '%Y-%m-%d %H:%M:%S')

        start_date = datetime.strptime(start_str, '%Y-%m-%d')
        end_date = datetime.strptime(end_str, '%Y-%m-%d')

        # 均匀分布日期
        delta = end_date - start_date
        random_days = (delta.days * index) / total_rows
        random_date = start_date + timedelta(days=random_days)

        return random_date.strftime(date_format)

    def _generate_choice(self, rule: Dict[str, Any]) -> str:
        """从选项中随机选择"""
        choices = rule.get('choices', ['value1', 'value2'])
        return str(random.choice(choices))

    def _generate_sequential(self, rule: Dict[str, Any], index: int) -> str:
        """生成顺序编号"""
        start = rule.get('start', 1)
        prefix = rule.get('prefix', '')
        return f"{prefix}{start + index:03d}"

    def _generate_foreign_key(self, rule: Dict[str, Any]) -> str:
        """生成外键值：从引用表中采样

        Args:
            rule: 外键规则，包含 reference_table 和 reference_field

        Returns:
            从引用表采样的字段值
        """
        ref_table = rule.get('reference_table')
        ref_field = rule.get('reference_field')

        if not ref_table or not ref_field:
            return "NULL"

        # 从已生成的表中获取数据
        ref_data = self.generated_tables.get(ref_table)
        if not ref_data:
            print(f"Warning: Reference table '{ref_table}' not found, using NULL")
            return "NULL"

        # 随机采样一个值
        sampled_row = random.choice(ref_data)
        value = sampled_row.get(ref_field)

        if value is None:
            print(f"Warning: Field '{ref_field}' not found in table '{ref_table}', using NULL")
            return "NULL"

        return str(value)

    def generate_csv(
        self, table_name: str, rules: Dict[str, Any]
    ) -> str:
        """生成 CSV 内容

        Args:
            table_name: 表名
            rules: 表的生成规则

        Returns:
            CSV 字符串
        """
        rows = self.generate_table_data(table_name, rules)

        if not rows:
            return ""

        # 转换为 CSV
        output = StringIO()
        fieldnames = list(rows[0].keys())
        writer = csv.DictWriter(output, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

        return output.getvalue()
