"""
CSV Generator - Automatically generate CSV data from messages.json
"""

import json
import re
from typing import Dict, List, Set, Any, Optional
from pathlib import Path
import csv
from io import StringIO
import os

from .llm_data_planner import LLMDataPlanner
from .rule_based_generator import RuleBasedGenerator


class CSVGenerator:
    """Generate CSV data from AskTable chat messages."""

    def __init__(self, llm_config: Optional[Dict[str, str]] = None):
        self.fields: Set[str] = set()
        self.sample_data: List[Dict[str, Any]] = []

        # LLM 配置
        if llm_config:
            self.llm_base_url = llm_config.get("base_url", "https://aiproxy.asktable.com/v1")
            self.llm_api_key = llm_config.get("api_key", "")
            self.llm_sql_model = llm_config.get("sql_model", "qwen-max-latest")
            self.llm_data_gen_model = llm_config.get("data_gen_model", "qwen-turbo-latest")
        else:
            self.llm_base_url = os.getenv("LLM_BASE_URL", "https://aiproxy.asktable.com/v1")
            self.llm_api_key = os.getenv("LLM_API_KEY", "")
            self.llm_sql_model = os.getenv("LLM_SQL_MODEL", "qwen-max-latest")
            self.llm_data_gen_model = os.getenv("LLM_DATA_GEN_MODEL", "qwen-turbo-latest")

        # 初始化 LLM 规划器和规则生成器
        if self.llm_api_key:
            self.llm_planner = LLMDataPlanner(
                api_key=self.llm_api_key,
                base_url=self.llm_base_url,
                model=self.llm_data_gen_model
            )
        else:
            self.llm_planner = None

        self.rule_generator = RuleBasedGenerator()

    def _analyze_sql_with_llm(self, sql_query: str) -> Dict[str, Any]:
        """
        Use LLM API to analyze SQL query and extract real table fields.

        Args:
            sql_query: SQL query string to analyze

        Returns:
            Dict with extracted tables, fields, and relationships
        """
        if not self.llm_api_key:
            # Fallback to regex-based extraction if no API key
            return self._analyze_sql_with_regex(sql_query)

        try:
            from openai import OpenAI

            client = OpenAI(
                api_key=self.llm_api_key,
                base_url=self.llm_base_url
            )

            prompt = f"""分析以下 SQL 查询，提取真实的数据库表和字段信息，并识别表间关联关系。

SQL 查询:
```sql
{sql_query}
```

请返回 JSON 格式，包含以下信息：
1. tables: 涉及的所有表名（不包括数据库前缀）
2. fields: 每个表的真实字段列表（不包括 SQL 关键字、函数名、表别名）
3. order_by_fields: ORDER BY 子句中的字段
4. join_fields: JOIN 条件中使用的字段
5. foreign_keys: 外键关系列表，每个关系包含 from_table, from_field, to_table, to_field

注意事项：
- 只提取真实的表字段名，不要包含：
  * SQL 关键字（AS, SELECT, FROM, WHERE, GROUP, ORDER, BY, LIMIT, YEAR, MONTH, DAY 等）
  * SQL 函数（SUM, COUNT, AVG, MAX, MIN, DATE_FORMAT 等）
  * 表别名（sp, tb, cc, cps 等单字母或短别名）
  * 临时列名（如 AS 后面定义的别名）
- 从 SELECT 子句中提取字段时，要识别出原始字段名（去掉表前缀和别名）
- 从 WHERE、GROUP BY、ORDER BY 子句中提取字段
- 从 JOIN ON 条件中提取关联字段，并识别外键关系

返回格式示例：
{{
  "tables": ["orders", "customers"],
  "fields": {{
    "orders": ["order_id", "customer_id", "amount"],
    "customers": ["customer_id", "name"]
  }},
  "order_by_fields": ["amount"],
  "join_fields": ["customer_id"],
  "foreign_keys": [
    {{
      "from_table": "orders",
      "from_field": "customer_id",
      "to_table": "customers",
      "to_field": "customer_id"
    }}
  ]
}}

只返回 JSON，不要有其他说明文字。"""

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

            # Extract JSON from response
            content = response.choices[0].message.content.strip()
            # Remove markdown code blocks if present
            if content.startswith("```"):
                content = content.split("```")[1]
                if content.startswith("json"):
                    content = content[4:]
                content = content.strip()

            result = json.loads(content)
            return result

        except Exception as e:
            print(f"Warning: LLM analysis failed ({e}), falling back to regex")
            return self._analyze_sql_with_regex(sql_query)

    def _analyze_sql_with_regex(self, sql_query: str) -> Dict[str, Any]:
        """
        Fallback method: analyze SQL using regex patterns.
        This is less accurate but doesn't require API access.
        """
        result = {
            "tables": [],
            "fields": {},
            "order_by_fields": [],
            "join_fields": [],
            "foreign_keys": []
        }

        # Extract table names
        table_names = self._extract_table_names(sql_query)
        result["tables"] = table_names

        # Extract fields from SELECT clause
        select_fields = self._extract_select_fields(sql_query)

        # Extract ORDER BY fields
        order_fields = self._extract_order_by_fields(sql_query)
        result["order_by_fields"] = order_fields

        # Group fields by table (simplified - assign all to first table)
        if table_names:
            result["fields"][table_names[0]] = list(select_fields)

        return result

    def analyze_messages(self, chat_file: Path) -> Dict[str, Any]:
        """
        Analyze messages.json to extract data structure.

        Returns:
            Dict with analysis results including fields, data types, and sample values
        """
        with open(chat_file, 'r', encoding='utf-8') as f:
            data = json.load(f)

        analysis = {
            'fields': set(),
            'sample_values': {},
            'sql_queries': [],
            'table_name': None,
            'tables': {},  # New: dict of table_name -> table_info
            'questions': [],  # New: list of user questions
        }

        # Extract from messages
        items = data.get('items', [])

        for item in items:
            if item.get('role') == 'human':
                # Extract user questions
                content = item.get('content', {})
                text = content.get('text', '').strip()
                if text:
                    analysis['questions'].append(text)

            elif item.get('role') == 'ai':
                metadata = item.get('metadata', {})
                tool_responses = metadata.get('tool_response', [])

                for response in tool_responses:
                    # Extract SQL query
                    sql = response.get('sql', '')
                    if sql:
                        analysis['sql_queries'].append(sql)

                        # Use LLM to analyze SQL and extract real fields
                        llm_result = self._analyze_sql_with_llm(sql)

                        # Extract table names
                        table_names = llm_result.get('tables', [])

                        # Set primary table name if not set
                        if table_names and not analysis['table_name']:
                            analysis['table_name'] = table_names[0]

                        # Process each table and its fields
                        for table_name in table_names:
                            table_fields = llm_result.get('fields', {}).get(table_name, [])

                            # Initialize table info if not exists
                            if table_name not in analysis['tables']:
                                analysis['tables'][table_name] = {
                                    'fields': set(),
                                    'sample_values': {},
                                    'primary_key': None,
                                    'relationships': [],
                                    'order_fields': []
                                }

                            # Add fields to table
                            for field in table_fields:
                                analysis['fields'].add(field)
                                analysis['tables'][table_name]['fields'].add(field)

                            # Add ORDER BY fields
                            order_fields = llm_result.get('order_by_fields', [])
                            if order_fields:
                                analysis['tables'][table_name]['order_fields'] = order_fields

                    # Extract sample data from result
                    result_str = response.get('result', '')
                    if result_str:
                        # Parse result string to extract field-value pairs
                        self._parse_result_string(result_str, analysis)

        # If no tables were detected, create a single table from all fields
        if not analysis['tables'] and analysis['fields']:
            table_name = analysis['table_name'] or 'data'
            analysis['tables'][table_name] = {
                'fields': analysis['fields'].copy(),
                'sample_values': analysis['sample_values'].copy(),
                'primary_key': None,
                'relationships': []
            }

        return analysis

    def _parse_result_string(self, result_str: str, analysis: Dict[str, Any]):
        """Parse result string to extract sample values."""
        # Split by newline to get individual records
        lines = result_str.strip().split('\n')

        for line in lines:
            # Extract key-value pairs
            matches = re.findall(r"'(\w+)':\s*([^,}]+)", line)
            for field, value in matches:
                analysis['fields'].add(field)

                if field not in analysis['sample_values']:
                    analysis['sample_values'][field] = []

                # Clean up value
                value = value.strip()
                if value not in ['NaT', 'nan', 'None']:
                    analysis['sample_values'][field].append(value)

                # Also add to table-specific sample values if tables exist
                for table_name, table_info in analysis.get('tables', {}).items():
                    if field in table_info['fields']:
                        if field not in table_info['sample_values']:
                            table_info['sample_values'][field] = []
                        if value not in ['NaT', 'nan', 'None']:
                            table_info['sample_values'][field].append(value)

    def _extract_table_names(self, sql: str) -> List[str]:
        """Extract all table names from SQL query (FROM and JOIN clauses)."""
        table_names = []

        # Extract from FROM clause
        from_match = re.search(r'FROM\s+([\w.]+)(?:\s+(\w+))?', sql, re.IGNORECASE)
        if from_match:
            table_full = from_match.group(1)
            # Remove schema prefix if present
            table_name = table_full.split('.')[-1]
            table_names.append(table_name)

        # Extract from JOIN clauses
        join_matches = re.finditer(r'JOIN\s+([\w.]+)(?:\s+(\w+))?', sql, re.IGNORECASE)
        for match in join_matches:
            table_full = match.group(1)
            # Remove schema prefix if present
            table_name = table_full.split('.')[-1]
            if table_name not in table_names:
                table_names.append(table_name)

        return table_names

    def _extract_select_fields(self, sql: str) -> Set[str]:
        """Extract field names from SELECT clause (basic regex fallback)."""
        fields = set()

        # Extract SELECT clause
        select_match = re.search(r'SELECT\s+(.*?)\s+FROM', sql, re.IGNORECASE | re.DOTALL)
        if not select_match:
            return fields

        select_clause = select_match.group(1)

        # Remove common SQL functions and keywords
        sql_keywords = {'DISTINCT', 'AS', 'COUNT', 'SUM', 'AVG', 'MAX', 'MIN', 'YEAR', 'MONTH', 'DAY', 'DATE_FORMAT'}

        # Extract field names (simple pattern)
        field_patterns = re.findall(r'`?(\w+)`?\.`?(\w+)`?|`?(\w+)`?', select_clause)

        for pattern in field_patterns:
            if pattern[1]:  # table.field format
                field_name = pattern[1]
            elif pattern[2]:  # field only
                field_name = pattern[2]
            else:
                continue

            # Skip SQL keywords
            if field_name.upper() not in sql_keywords and len(field_name) > 2:
                fields.add(field_name)

        return fields

    def _resolve_table_alias(self, sql: str, alias: str, table_names: List[str]) -> Optional[str]:
        """Resolve table alias to actual table name."""
        # Look for "FROM table_name alias" or "JOIN table_name alias"
        pattern = rf'(?:FROM|JOIN)\s+([\w.]+)\s+{re.escape(alias)}\b'
        match = re.search(pattern, sql, re.IGNORECASE)
        if match:
            table_full = match.group(1)
            return table_full.split('.')[-1]

        # If alias matches a table name, return it
        if alias in table_names:
            return alias

        # Default to first table
        return table_names[0] if table_names else None

    def generate_csv_from_analysis(
        self,
        analysis: Dict[str, Any],
        num_rows: int = 20,
        case_name: str = "demo",
        max_columns: int = 15
    ) -> str:
        """
        Generate CSV content based on analysis.

        Args:
            analysis: Analysis results from analyze_messages
            num_rows: Number of rows to generate
            case_name: Case name for generating sample data (unused, kept for compatibility)
            max_columns: Maximum number of columns (default: 15)

        Returns:
            CSV content as string
        """
        fields = sorted(analysis['fields'])

        if not fields:
            raise ValueError("No fields found in messages")

        # Limit to max_columns
        if len(fields) > max_columns:
            fields = fields[:max_columns]

        # Generate sample data
        rows = []
        sample_values = analysis.get('sample_values', {})

        for i in range(num_rows):
            row = {}
            for field in fields:
                field_lower = field.lower()

                # For name and score fields, always generate with variation
                if 'name' in field_lower or 'score' in field_lower or 'level' in field_lower:
                    row[field] = self._generate_placeholder_with_variation(field, i, num_rows)
                # Use sample values if available for other fields
                elif field in sample_values and sample_values[field]:
                    # Cycle through sample values
                    value = sample_values[field][i % len(sample_values[field])]
                    row[field] = self._clean_value(value)
                else:
                    # Generate placeholder based on field name with variation
                    row[field] = self._generate_placeholder_with_variation(field, i, num_rows)

            rows.append(row)

        # Convert to CSV
        output = StringIO()
        writer = csv.DictWriter(output, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)

        return output.getvalue()

    def _clean_value(self, value: str) -> str:
        """Clean up value from result string."""
        # Remove quotes and extra whitespace
        value = value.strip().strip('"\'')

        # Handle Timestamp
        if value.startswith('Timestamp('):
            match = re.search(r"'([^']+)'", value)
            if match:
                return match.group(1)

        return value

    def _analyze_questions(self, questions: List[str]) -> Dict[str, Any]:
        """Analyze questions to determine data requirements.

        Args:
            questions: List of user questions

        Returns:
            Dict with question analysis including required rows, operations, etc.
        """
        analysis = {
            'max_rows_needed': 20,  # Default
            'has_top_n': False,
            'top_n_value': 0,
            'has_aggregation': False,
            'has_filter': False,
        }

        for question in questions:
            question_lower = question.lower()

            # Check for "top N" or "前N名" patterns
            top_n_patterns = [
                r'top\s+(\d+)',
                r'前\s*(\d+)\s*名',
                r'最好的\s*(\d+)',
                r'最高的\s*(\d+)',
            ]

            for pattern in top_n_patterns:
                match = re.search(pattern, question_lower)
                if match:
                    n = int(match.group(1))
                    analysis['has_top_n'] = True
                    analysis['top_n_value'] = max(analysis['top_n_value'], n)
                    # Need at least 2x the requested rows for meaningful ranking
                    analysis['max_rows_needed'] = max(analysis['max_rows_needed'], n * 2)

            # Check for aggregation keywords
            agg_keywords = ['平均', '总计', '求和', '最大', '最小', 'average', 'sum', 'total', 'max', 'min', 'count']
            if any(keyword in question_lower for keyword in agg_keywords):
                analysis['has_aggregation'] = True
                # Aggregations need more data
                analysis['max_rows_needed'] = max(analysis['max_rows_needed'], 30)

            # Check for filter keywords
            filter_keywords = ['大于', '小于', '等于', '高于', '低于', '超过', 'greater', 'less', 'equal', 'above', 'below']
            if any(keyword in question_lower for keyword in filter_keywords):
                analysis['has_filter'] = True

        return analysis

    def generate_multiple_csvs_from_analysis(
        self,
        analysis: Dict[str, Any],
        num_rows: int = 20,
        output_dir: Optional[Path] = None,
    ) -> List[tuple]:
        """Generate multiple CSV files based on analysis.

        Args:
            analysis: Analysis results from analyze_messages
            num_rows: Base number of rows to generate
            output_dir: Output directory (if None, return content in memory; if Path, save to files)

        Returns:
            List of (table_name, csv_content) tuples if output_dir is None,
            or List of (table_name, csv_file_path) tuples if output_dir is provided
        """
        # 优先使用 LLM 规划方案
        if self.llm_planner:
            try:
                return self._generate_with_llm_rules(analysis, num_rows, output_dir)
            except Exception as e:
                print(f"Warning: LLM generation failed ({e}), falling back to legacy method")

        # 降级到原有方案
        return self._generate_with_legacy_method(analysis, num_rows)

    def _generate_with_llm_rules(
        self, analysis: Dict[str, Any], num_rows: int, output_dir: Optional[Path] = None
    ) -> List[tuple]:
        """使用 LLM 规则生成数据

        Args:
            analysis: 分析结果
            num_rows: 基础行数
            output_dir: 输出目录（如果为 None，则返回内存中）

        Returns:
            (表名，CSV 内容) 元组列表
        """
        # 1. 用 LLM 规划数据生成规则
        print("🤖 使用 LLM 规划数据生成策略...")
        rules = self.llm_planner.plan_data_generation(analysis, num_rows)

        # 2. 按依赖顺序生成所有表
        print("  按依赖顺序生成表数据...")
        all_tables_data = self.rule_generator.generate_all_tables(rules)

        # 3. 转换为 CSV 并保存
        csv_files = []
        for table_name, table_data in all_tables_data.items():
            print(f"  生成表 {table_name} 的 CSV...")
            csv_content = self._convert_to_csv(table_data)

            if output_dir:
                # 保存到文件
                csv_file = output_dir / f"{table_name}.csv"
                try:
                    output_dir.mkdir(parents=True, exist_ok=True)
                    with open(csv_file, 'w', encoding='utf-8') as f:
                        f.write(csv_content)
                    csv_files.append((table_name, csv_file))
                    print(f"    已保存：{csv_file}")
                except Exception as e:
                    print(f"    保存失败：{e}")
                    csv_files.append((table_name, csv_content))
            else:
                csv_files.append((table_name, csv_content))

        return csv_files

    def _convert_to_csv(self, table_data: List[Dict[str, Any]]) -> str:
        """将表数据转换为 CSV 字符串"""
        if not table_data:
            return ""

        from io import StringIO
        import csv

        output = StringIO()
        fieldnames = list(table_data[0].keys())
        writer = csv.DictWriter(output, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(table_data)

        return output.getvalue()

    def _generate_with_legacy_method(
        self, analysis: Dict[str, Any], num_rows: int
    ) -> List[tuple]:
        """使用原有方法生成数据（降级方案）"""
        tables = analysis.get('tables', {})

        if not tables:
            # Fallback to single CSV generation
            csv_content = self.generate_csv_from_analysis(analysis, num_rows)
            table_name = analysis.get('table_name', 'data')
            return [(table_name, csv_content)]

        # Analyze questions to determine how many rows we need
        questions = analysis.get('questions', [])
        question_analysis = self._analyze_questions(questions)
        actual_rows = max(num_rows, question_analysis['max_rows_needed'])

        # Extract fields used in SQL ORDER BY clauses
        order_by_fields = self._extract_order_by_fields(analysis.get('sql_queries', []))

        csv_contents = []

        for table_name, table_info in tables.items():
            all_fields = table_info['fields']

            if not all_fields:
                continue

            # Prioritize fields in order:
            # 1. Name/ID fields (for identification)
            # 2. Fields used in ORDER BY (for ranking queries)
            # 3. Other fields
            name_id_fields = []
            order_fields = []
            regular_fields = []

            for field in all_fields:
                field_lower = field.lower()
                # High priority: name, id, no fields
                if 'name' in field_lower or field_lower.endswith('_id') or field_lower.endswith('_no'):
                    name_id_fields.append(field)
                # Medium priority: fields used in ORDER BY
                elif field in order_by_fields:
                    order_fields.append(field)
                else:
                    regular_fields.append(field)

            # Sort each group
            name_id_fields.sort()
            order_fields.sort()
            regular_fields.sort()

            # Combine: name/id first, then ORDER BY fields, then regular fields
            fields = name_id_fields + order_fields + regular_fields

            # Limit to 15 columns
            if len(fields) > 15:
                fields = fields[:15]

            # Generate sample data with realistic variation for ranking
            rows = []
            sample_values = table_info.get('sample_values', {})

            for i in range(actual_rows):
                row = {}
                for field in fields:
                    field_lower = field.lower()

                    # For name and score fields, always generate with variation (don't use sample values)
                    # This ensures unique names and proper ranking data
                    if 'name' in field_lower or 'score' in field_lower or 'level' in field_lower:
                        row[field] = self._generate_placeholder_with_variation(field, i, actual_rows)
                    # Use sample values if available for other fields
                    elif field in sample_values and sample_values[field]:
                        # Cycle through sample values
                        value = sample_values[field][i % len(sample_values[field])]
                        row[field] = self._clean_value(value)
                    else:
                        # Generate placeholder based on field name with variation
                        row[field] = self._generate_placeholder_with_variation(field, i, actual_rows)

                rows.append(row)

            # Convert to CSV
            output = StringIO()
            writer = csv.DictWriter(output, fieldnames=fields)
            writer.writeheader()
            writer.writerows(rows)

            csv_contents.append((table_name, output.getvalue()))

        return csv_contents

    def _extract_order_by_fields(self, sql_queries: List[str]) -> Set[str]:
        """Extract field names used in ORDER BY clauses.

        Args:
            sql_queries: List of SQL queries

        Returns:
            Set of field names used in ORDER BY
        """
        order_by_fields = set()

        for sql in sql_queries:
            # Find ORDER BY clause
            order_by_match = re.search(r'ORDER\s+BY\s+(.*?)(?:LIMIT|$)', sql, re.IGNORECASE | re.DOTALL)
            if order_by_match:
                order_clause = order_by_match.group(1)
                # Extract field names (handle table.field and field formats)
                fields = re.findall(r'(\w+)\.(\w+)|(\w+)', order_clause)
                for field_match in fields:
                    if field_match[1]:  # table.field format
                        field_name = field_match[1]
                    elif field_match[2]:  # field only
                        field_name = field_match[2]
                    else:
                        continue

                    # Skip SQL keywords
                    if field_name.upper() not in ['DESC', 'ASC', 'NULLS', 'FIRST', 'LAST']:
                        order_by_fields.add(field_name)

        return order_by_fields

    def _generate_placeholder_with_variation(self, field: str, index: int, total_rows: int) -> str:
        """生成字段占位值（简单默认值）

        Args:
            field: 字段名
            index: 行索引
            total_rows: 总行数

        Returns:
            生成的值
        """
        return f"{field}_{index}"

    def generate_csv_from_messages(
        self,
        chat_file: Path,
        output_file: Optional[Path] = None,
        num_rows: int = 20,
        case_name: str = "demo"
    ) -> str:
        """
        Complete workflow: analyze chat and generate CSV.

        Args:
            chat_file: Path to chat-{number}.json
            output_file: Optional path to save CSV file
            num_rows: Number of rows to generate
            case_name: Case name for generating sample data (unused, kept for compatibility)

        Returns:
            CSV content as string (for backward compatibility, returns first CSV)
        """
        print(f"📖 分析 Chat 文件: {chat_file}")
        analysis = self.analyze_messages(chat_file)

        total_fields = len(analysis['fields'])
        print(f"✓ 发现 {total_fields} 个字段")

        # Limit to 15 columns
        if total_fields > 15:
            print("⚠ 字段数超过 15，将只使用前 15 个字段")

        print(f"✓ 发现 {len(analysis['sql_queries'])} 个 SQL 查询")

        if analysis['table_name']:
            print(f"✓ 表名: {analysis['table_name']}")

        # Check if we have multiple tables
        tables = analysis.get('tables', {})
        if len(tables) > 1:
            print(f"✓ 检测到 {len(tables)} 个数据表")

        print(f"\n🔧 生成 {num_rows} 行虚拟数据...")

        # Use new multi-table generation
        csv_contents = self.generate_multiple_csvs_from_analysis(analysis, num_rows)

        # For backward compatibility, save the first CSV if output_file is provided
        if output_file and csv_contents:
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(csv_contents[0][1])
            print(f"✓ CSV 文件已保存: {output_file}")

        # Return first CSV content for backward compatibility
        return csv_contents[0][1] if csv_contents else ""

    def get_case_description(self, analysis: Dict[str, Any]) -> str:
        """Generate a description for the case based on analysis."""
        table_name = analysis.get('table_name', 'data')
        num_fields = len(analysis['fields'])

        return f"{table_name} 数据分析 ({num_fields} 个字段)"
