"""LLM 数据规划器 - 用大模型生成数据生成规则"""

import json
from typing import Dict, List, Any, Optional


class LLMDataPlanner:
    """用 LLM 规划数据生成策略"""

    def __init__(self, api_key: str, base_url: str = "https://aiproxy.dminfra.cn/v1", model: str = "qwen-turbo"):
        self.api_key = api_key
        self.base_url = base_url
        self.model = model

        # 延迟导入 openai
        try:
            from openai import OpenAI
            self.client = OpenAI(
                api_key=api_key,
                base_url=base_url
            )
        except ImportError:
            raise ImportError(
                "openai package is required for LLM data planning. "
                "Install it with: pip install openai"
            )

    def generate_case_name(self, analysis: Dict[str, Any]) -> str:
        """用 LLM 生成案例名称

        Args:
            analysis: 分析结果（包含 table_name, fields, questions）

        Returns:
            中文案例名称
        """
        table_name = analysis.get('table_name', '')
        fields = list(analysis.get('fields', set()))[:10]
        questions = analysis.get('questions', [])[:3]

        prompt = f"""根据以下数据库信息，生成一个简洁的中文案例名称（3-6 个字）：

表名: {table_name}
字段: {', '.join(fields)}
用户问题: {', '.join(questions)}

要求：
1. 名称要体现数据的业务领域
2. 以"分析"结尾
3. 只返回名称，不要解释

示例：
- 表名 student_performance → "学生成绩分析"
- 表名 sales_order → "销售订单分析"
- 表名 sport_data → "体育数据分析"
"""

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0
            )

            case_name = response.choices[0].message.content.strip()
            # 移除可能的引号
            case_name = case_name.strip('"\'')
            return case_name

        except Exception as e:
            print(f"Warning: LLM case name generation failed ({e}), using fallback")
            # 降级方案
            if table_name:
                return f"{table_name}分析"
            return f"数据分析案例（{len(fields)}个字段）"

    def plan_data_generation(self, analysis: Dict[str, Any], num_rows: int = 20) -> Dict[str, Any]:
        """用 LLM 规划数据生成规则

        Args:
            analysis: 分析结果（包含 tables, fields, questions, sql_queries）
            num_rows: 基础行数

        Returns:
            数据生成规则
        """
        tables = analysis.get('tables', {})
        questions = analysis.get('questions', [])
        sql_queries = analysis.get('sql_queries', [])
        foreign_keys = analysis.get('foreign_keys', [])

        # 构建表信息摘要
        table_summaries = []
        for table_name, table_info in tables.items():
            fields = list(table_info.get('fields', set()))
            table_summaries.append(f"表 {table_name}: {', '.join(fields)}")

        # 构建外键关系摘要
        fk_summaries = []
        if foreign_keys:
            for fk in foreign_keys:
                fk_summaries.append(
                    f"- {fk['from_table']}.{fk['from_field']} → {fk['to_table']}.{fk['to_field']}"
                )

        prompt = f"""为以下数据库表设计虚拟数据生成规则。

表结构:
{chr(10).join(table_summaries)}

{"外键关系:" + chr(10) + chr(10).join(fk_summaries) if fk_summaries else ""}

用户问题:
{chr(10).join(f"- {q}" for q in questions[:5])}

SQL 查询示例:
{chr(10).join(f"```sql{chr(10)}{sql}{chr(10)}```" for sql in sql_queries[:2])}

请返回 JSON 格式的生成规则：
{{
  "tables": {{
    "table_name": {{
      "num_rows": 50,  // 根据问题推断需要的行数（如"前10名"需要至少20行）
      "fields": {{
        "name": {{
          "type": "chinese_name",
          "unique": true,
          "description": "中文姓名"
        }},
        "customer_id": {{
          "type": "foreign_key",
          "reference_table": "customers",
          "reference_field": "customer_id",
          "description": "外键，引用 customers 表的 customer_id"
        }},
        "score": {{
          "type": "normal_distribution",
          "mean": 75,
          "std": 15,
          "min": 0,
          "max": 100,
          "description": "分数，符合正态分布"
        }},
        "age": {{
          "type": "uniform",
          "min": 18,
          "max": 22,
          "description": "年龄，均匀分布"
        }},
        "date": {{
          "type": "date_range",
          "start": "2025-01-01",
          "end": "2025-12-31",
          "format": "%Y-%m-%d %H:%M:%S",
          "description": "日期范围"
        }},
        "gender": {{
          "type": "choice",
          "choices": ["1", "2"],
          "description": "性别，1=男，2=女"
        }}
      }}
    }}
  }}
}}

支持的字段类型：
- chinese_name: 中文姓名
- foreign_key: 外键，从引用表中采样（reference_table, reference_field）
- normal_distribution: 正态分布数值（mean, std, min, max）
- uniform: 均匀分布数值（min, max）
- date_range: 日期范围（start, end, format）
- choice: 从列表中选择（choices）
- sequential: 顺序编号（start, prefix）

注意：
1. 根据字段名推断合理的数据类型和范围
2. 如果问题涉及排序（如"前N名"），确保相关字段有足够的数据分布
3. num_rows 要根据问题需求调整（如"前10名"至少需要20行）
4. 对于外键字段，必须使用 foreign_key 类型，并指定 reference_table 和 reference_field
5. 被引用的主表（如 customers）应该先生成，从表（如 orders）后生成

只返回 JSON，不要解释。
"""

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0
            )

            content = response.choices[0].message.content.strip()

            # 移除可能的 markdown 代码块
            if content.startswith("```"):
                content = content.split("```")[1]
                if content.startswith("json"):
                    content = content[4:]
                content = content.strip()

            rules = json.loads(content)
            return rules

        except Exception as e:
            print(f"Warning: LLM data planning failed ({e}), using fallback")
            # 降级方案：返回基础规则
            return self._generate_fallback_rules(analysis, num_rows)

    def _generate_fallback_rules(self, analysis: Dict[str, Any], num_rows: int) -> Dict[str, Any]:
        """降级方案：生成基础规则"""
        tables = analysis.get('tables', {})
        rules = {"tables": {}}

        for table_name, table_info in tables.items():
            fields = table_info.get('fields', set())
            field_rules = {}

            for field in fields:
                field_lower = field.lower()

                # 简单的字段类型推断
                if 'name' in field_lower:
                    field_rules[field] = {"type": "chinese_name", "unique": True}
                elif 'score' in field_lower:
                    field_rules[field] = {
                        "type": "normal_distribution",
                        "mean": 75, "std": 15, "min": 0, "max": 100
                    }
                elif 'age' in field_lower:
                    field_rules[field] = {"type": "uniform", "min": 18, "max": 25}
                elif 'date' in field_lower or 'time' in field_lower:
                    field_rules[field] = {
                        "type": "date_range",
                        "start": "2025-01-01",
                        "end": "2025-12-31",
                        "format": "%Y-%m-%d %H:%M:%S"
                    }
                elif 'id' in field_lower or 'no' in field_lower:
                    field_rules[field] = {"type": "sequential", "start": 1, "prefix": "ID"}
                else:
                    field_rules[field] = {"type": "default"}

            rules["tables"][table_name] = {
                "num_rows": num_rows,
                "fields": field_rules
            }

        return rules

