"""
Text-to-SQL 节点

将自然语言转换为 SQL 查询并执行，支持流式输出和图表推荐
"""
import json
import logging
import re
import time
from dataclasses import dataclass
from typing import Any, Dict, Generator, List, Optional, Tuple

import sqlparse

from ..base import BaseNode, NodeContext, NodeResult
from ..registry import NodeRegistry

logger = logging.getLogger(__name__)


@dataclass
class TextToSqlStreamEvent:
    """Text-to-SQL 流式输出事件"""
    event_type: str = ''  # thought, sql_chunk, sql_complete, executing, result
    content: str = ''
    is_finished: bool = False
    data: Any = None


# Text-to-SQL Function Calling 工具定义
TEXT_TO_SQL_TOOL = {
    'name': 'generate_sql',
    'description': '根据用户的自然语言问题生成 SQL 查询语句',
    'parameters': {
        'type': 'object',
        'properties': {
            'thought': {
                'type': 'string',
                'description': '分析思路，简要描述如何理解用户问题并设计 SQL',
            },
            'sql': {
                'type': 'string',
                'description': '生成的 SQL 查询语句，必须是有效的 SELECT 语句，必须格式化，如果表结构中包含 schema 信息，SQL 中的表名必须使用完整格式：schema.table_name',
            },
        },
        'required': ['thought', 'sql'],
    },
}

# Text-to-SQL Function Calling 版 System Prompt（简化版）
TEXT_TO_SQL_SYSTEM_PROMPT_FC = """你是一个专业的数据分析助手，擅长将自然语言转换为 SQL 查询。

## 数据库信息
数据库类型: {db_type}
当前日期: {current_date}

## 可用的表结构
{schema_context}

## 任务要求
1. 根据用户问题生成准确的 SQL 查询
2. 只生成 SELECT 语句，禁止生成 INSERT/UPDATE/DELETE/DROP 等修改语句
3. 使用表的注释理解业务含义
4. 添加必要的 WHERE 条件和 ORDER BY
5. 如果表结构中包含 schema 信息，SQL 中的表名必须使用完整格式：schema.table_name

## 注意事项
- SQL 必须是有效的 {db_type} 语法
- 避免使用 SELECT *，明确指定需要的字段
- 对于大数据量查询，添加 LIMIT 限制
- 使用 DISTINCT 时，ORDER BY 的列必须出现在 SELECT 列表中

请使用 generate_sql 函数返回结果。
"""

# Text-to-SQL 专用 System Prompt（原有版本，作为 fallback）
TEXT_TO_SQL_SYSTEM_PROMPT = """你是一个专业的数据分析助手，擅长将自然语言转换为 SQL 查询。

## 数据库信息
数据库类型: {db_type}
当前日期: {current_date}

## 可用的表结构
{schema_context}

## 任务要求
1. 根据用户问题生成准确的 SQL 查询
2. 只生成 SELECT 语句，禁止生成 INSERT/UPDATE/DELETE/DROP 等修改语句
3. 使用表的注释理解业务含义
4. 添加必要的 WHERE 条件和 ORDER BY
5. 如果表结构中包含 schema 信息，SQL 中的表名必须使用完整格式：schema.table_name
6. 如果需要，推荐最合适的图表类型

## 输出格式
严格按照以下 JSON 格式输出，不要输出其他内容：
{{
  "thought": "你的分析思路（简洁描述）",
  "sql": "生成的 SQL 语句",
  "chart_type": "bar|line|pie|scatter|radar|table|none",
  "chart_config": {{
    "x_field": "X轴字段名（用于 bar/line/scatter）",
    "y_field": "Y轴字段名（用于 scatter）",
    "series_fields": ["系列字段1", "系列字段2"],
    "name_field": "名称字段（用于 pie）",
    "value_field": "数值字段（用于 pie/gauge）",
    "title": "图表标题"
  }}
}}

## 图表类型选择指南
- bar（柱状图）: 分类对比，如各部门销售额
- line（折线图）: 时间趋势，如每日销售额变化
- pie（饼图）: 占比分析，如各类别占比
- scatter（散点图）: 分布分析，如价格与销量关系
- radar（雷达图）: 多维对比，如产品多指标评分
- table（表格）: 详细数据展示
- none（无图表）: 不适合可视化的数据

## 注意事项
- SQL 必须是有效的 {db_type} 语法
- 避免使用 SELECT *，明确指定需要的字段
- 对于大数据量查询，添加 LIMIT 限制
- 时间字段使用标准格式
- 使用 DISTINCT 时，ORDER BY 的列必须出现在 SELECT 列表中
- 如需去重并排序，考虑使用子查询或窗口函数
"""


def serialize_db_value(value: Any) -> Any:
    """将数据库值转换为可 JSON 序列化的格式"""
    from datetime import date, datetime
    from decimal import Decimal
    from uuid import UUID
    
    if value is None:
        return None
    if isinstance(value, datetime):
        return value.isoformat()
    if isinstance(value, date):
        return value.isoformat()
    if isinstance(value, Decimal):
        return float(value)
    if isinstance(value, UUID):
        return str(value)
    if isinstance(value, bytes):
        return value.decode('utf-8', errors='replace')
    if isinstance(value, (list, tuple)):
        return [serialize_db_value(v) for v in value]
    if isinstance(value, dict):
        return {k: serialize_db_value(v) for k, v in value.items()}
    return value


@NodeRegistry.register
class TextToSqlNode(BaseNode):
    """
    Text-to-SQL 节点
    
    将自然语言转换为 SQL 查询并执行
    """
    
    node_type = 'text_to_sql'
    node_name = 'Text-to-SQL'
    node_category = 'data'
    node_icon = 'database-zap'
    node_description = '将自然语言转换为 SQL 查询并执行，支持图表推荐'
    
    inputs = [
        {
            'name': 'user_question',
            'type': 'string',
            'description': '用户的自然语言问题',
        },
    ]
    
    outputs = [
        {
            'name': 'text_to_sql_result_sql',
            'type': 'string',
            'description': '生成的 SQL 语句',
        },
        {
            'name': 'text_to_sql_result_thought',
            'type': 'string',
            'description': 'SQL 生成思路',
        },
    ]
    
    def execute(self, context: NodeContext) -> NodeResult:
        """执行 Text-to-SQL（同步方法）"""
        import asyncio
        
        loop = asyncio.get_event_loop()
        if loop.is_running():
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as executor:
                future = executor.submit(asyncio.run, self.execute_async(context))
                return future.result()
        else:
            return loop.run_until_complete(self.execute_async(context))
    
    async def execute_async(self, context: NodeContext) -> NodeResult:
        """异步执行 Text-to-SQL"""
        start_time = time.time()
        
        try:
            # 获取配置
            user_question = self.config.get('user_question', '')
            db_config = self.config.get('db_config') or {}
            db_connection = db_config.get('dbName', 'default')
            schema_name = db_config.get('schema', 'public')
            model_id = self.config.get('model_id', '')
            selected_tables = self.config.get('selected_tables', [])
            table_relations = self.config.get('table_relations', [])  # 手动指定的表关系
            output_variable = self.config.get('output_variable', 'text_to_sql_result')
            include_relations = self.config.get('include_table_relations', True)
            
            # 解析变量
            if user_question:
                user_question = context.resolve_template(user_question)
            
            if not user_question:
                return NodeResult(
                    success=False,
                    error='请输入要查询的问题',
                )
            
            if not model_id:
                return NodeResult(
                    success=False,
                    error='请选择 LLM 模型',
                )
            
            # Step 1: 获取数据库 Schema
            
            schema_context = await self._get_schema_context(
                db_connection, 
                schema_name,
                selected_tables,
                include_relations,
                table_relations
            )
            if not schema_context:
                return NodeResult(
                    success=False,
                    error=f'无法获取数据库 {db_connection} 的表结构信息',
                )
            
            # Step 2: 调用 LLM 生成 SQL
            from datetime import datetime
            from ai_platform.services.llm_service import LLMService
            
            db_type = await self._get_db_type(db_connection)
            llm_service = LLMService(context.db_session)
            llm_result = None
            use_function_calling = self.config.get('use_function_calling', True)
            
            # 优先尝试 Function Calling
            if use_function_calling:
                try:
                    system_prompt = TEXT_TO_SQL_SYSTEM_PROMPT_FC.format(
                        db_type=db_type,
                        current_date=datetime.now().strftime('%Y-%m-%d'),
                        schema_context=schema_context,
                    )
                    
                    messages = [
                        {'role': 'system', 'content': system_prompt},
                        {'role': 'user', 'content': f'用户问题: {user_question}'},
                    ]
                    
                    response = await llm_service.chat_async(
                        model_id=model_id,
                        messages=messages,
                        temperature=0.1,
                        max_tokens=2048,
                        tools=[TEXT_TO_SQL_TOOL],
                        tool_choice='required',
                    )
                    
                    # 解析 Function Calling 响应
                    if response.tool_calls and len(response.tool_calls) > 0:
                        tool_call = response.tool_calls[0]
                        if tool_call.name == 'generate_sql':
                            import json
                            arguments = tool_call.arguments
                            if isinstance(arguments, str):
                                llm_result = json.loads(arguments)
                            else:
                                llm_result = arguments
                            logger.info(f'Function Calling 成功生成 SQL')
                except Exception as fc_error:
                    logger.warning(f'Function Calling 失败，回退到 JSON 解析模式: {fc_error}')
                    llm_result = None
            
            # Fallback: 使用 JSON 解析方式
            if not llm_result:
                system_prompt = TEXT_TO_SQL_SYSTEM_PROMPT.format(
                    db_type=db_type,
                    current_date=datetime.now().strftime('%Y-%m-%d'),
                    schema_context=schema_context,
                )
                
                messages = [
                    {'role': 'system', 'content': system_prompt},
                    {'role': 'user', 'content': f'用户问题: {user_question}'},
                ]
                
                response = await llm_service.chat_async(
                    model_id=model_id,
                    messages=messages,
                    temperature=0.1,
                    max_tokens=2048,
                )
                
                # 解析 LLM 响应
                llm_result = self._parse_llm_response(response.content)
                if not llm_result:
                    return NodeResult(
                        success=False,
                        error='LLM 响应解析失败，请重试',
                        metadata={'raw_response': response.content},
                    )
            
            sql = llm_result.get('sql', '')
            thought = llm_result.get('thought', '')
            
            # Step 3: 验证 SQL 安全性
            if not self._validate_sql(sql):
                return NodeResult(
                    success=False,
                    error='生成的 SQL 不安全或不是 SELECT 语句',
                    metadata={'sql': sql},
                )
            
            # Step 4: 格式化 SQL
            sql = self._format_sql(sql)
            
            elapsed_time = int((time.time() - start_time) * 1000)
            
            # 构建输出
            result_data = {
                'sql': sql,
                'thought': thought,
            }
            
            return NodeResult(
                success=True,
                output=result_data,
                output_variables={
                    f'{output_variable}_sql': sql,
                    f'{output_variable}_thought': thought,
                },
                tokens_used=response.total_tokens,
                elapsed_time=elapsed_time,
                metadata={
                    'model': response.model,
                    'thought': thought,
                    'suggested_next_node': 'db_sql',
                },
            )
            
        except Exception as e:
            logger.exception(f'Text-to-SQL 节点执行失败: {e}')
            elapsed_time = int((time.time() - start_time) * 1000)
            return NodeResult(
                success=False,
                error=str(e),
                elapsed_time=elapsed_time,
            )
    
    def execute_stream(self, context: NodeContext) -> Generator[TextToSqlStreamEvent, None, NodeResult]:
        """
        流式执行 Text-to-SQL
        
        Yields:
            TextToSqlStreamEvent: 流式输出事件
            
        Returns:
            NodeResult: 最终执行结果
        """
        start_time = time.time()
        
        try:
            # 获取配置
            user_question = self.config.get('user_question', '')
            db_config = self.config.get('db_config') or {}
            db_connection = db_config.get('dbName', 'default')
            schema_name = db_config.get('schema', 'public')
            model_id = self.config.get('model_id', '')
            selected_tables = self.config.get('selected_tables', [])
            output_variable = self.config.get('output_variable', 'text_to_sql_result')
            include_relations = self.config.get('include_table_relations', True)
            
            # 解析变量
            if user_question:
                user_question = context.resolve_template(user_question)
            
            if not user_question:
                yield TextToSqlStreamEvent(
                    event_type='error',
                    content='请输入要查询的问题',
                    is_finished=True,
                )
                return NodeResult(success=False, error='请输入要查询的问题')
            
            if not model_id:
                yield TextToSqlStreamEvent(
                    event_type='error',
                    content='请选择 LLM 模型',
                    is_finished=True,
                )
                return NodeResult(success=False, error='请选择 LLM 模型')
            
            # Step 1: 获取数据库 Schema
            yield TextToSqlStreamEvent(
                event_type='thought',
                content='正在获取数据库表结构...',
            )
            
            table_relations = self.config.get('table_relations', [])  # 手动指定的表关系
            
            import asyncio
            schema_context = asyncio.get_event_loop().run_until_complete(
                self._get_schema_context(
                    db_connection,
                    schema_name,
                    selected_tables,
                    include_relations,
                    table_relations
                )
            )
            
            if not schema_context:
                yield TextToSqlStreamEvent(
                    event_type='error',
                    content=f'无法获取数据库 {db_connection} 的表结构信息',
                    is_finished=True,
                )
                return NodeResult(
                    success=False,
                    error=f'无法获取数据库 {db_connection} 的表结构信息',
                )
            
            # Step 2: 调用 LLM 生成 SQL（流式）
            yield TextToSqlStreamEvent(
                event_type='thought',
                content='正在分析问题并生成 SQL...',
            )
            
            from datetime import datetime
            from ai_platform.services.llm_service import LLMService
            
            db_type = asyncio.get_event_loop().run_until_complete(
                self._get_db_type(db_connection)
            )
            
            llm_service = LLMService()
            accumulated_content = ''
            total_tokens = 0
            llm_result = None
            use_function_calling = self.config.get('use_function_calling', True)
            
            # 优先尝试 Function Calling
            if use_function_calling:
                try:
                    llm_result, total_tokens = yield from self._call_llm_with_function_calling(
                        llm_service, model_id, db_type, schema_context, user_question
                    )
                    logger.info(f'Function Calling 成功: {llm_result}')
                except Exception as fc_error:
                    logger.warning(f'Function Calling 失败，回退到 JSON 解析模式: {fc_error}')
                    llm_result = None
            
            # Fallback: 使用原有的 JSON 解析方式
            if not llm_result:
                system_prompt = TEXT_TO_SQL_SYSTEM_PROMPT.format(
                    db_type=db_type,
                    current_date=datetime.now().strftime('%Y-%m-%d'),
                    schema_context=schema_context,
                )
                
                messages = [
                    {'role': 'system', 'content': system_prompt},
                    {'role': 'user', 'content': f'用户问题: {user_question}'},
                ]
                
                for chunk in llm_service.chat_stream_sync(
                    model_id=model_id,
                    messages=messages,
                    temperature=0.1,
                    max_tokens=2048,
                ):
                    accumulated_content += chunk.content
                    
                    if chunk.is_finished and chunk.total_tokens > 0:
                        total_tokens = chunk.total_tokens
                    
                    yield TextToSqlStreamEvent(
                        event_type='sql_chunk',
                        content=chunk.content,
                        is_finished=chunk.is_finished,
                    )
                
                # 解析 LLM 响应
                llm_result = self._parse_llm_response(accumulated_content)
                if not llm_result:
                    yield TextToSqlStreamEvent(
                        event_type='error',
                        content='LLM 响应解析失败',
                        is_finished=True,
                    )
                    return NodeResult(
                        success=False,
                        error='LLM 响应解析失败',
                        metadata={'raw_response': accumulated_content},
                    )
            
            sql = llm_result.get('sql', '')
            thought = llm_result.get('thought', '')
            
            yield TextToSqlStreamEvent(
                event_type='sql_complete',
                content=sql,
                data={'thought': thought},
            )
            
            # Step 3: 验证 SQL 安全性
            if not self._validate_sql(sql):
                yield TextToSqlStreamEvent(
                    event_type='error',
                    content='生成的 SQL 不安全或不是 SELECT 语句',
                    is_finished=True,
                )
                return NodeResult(
                    success=False,
                    error='生成的 SQL 不安全或不是 SELECT 语句',
                    metadata={'sql': sql},
                )
            
            # Step 4: 格式化 SQL
            sql = self._format_sql(sql)
            
            elapsed_time = int((time.time() - start_time) * 1000)
            
            # 构建输出
            result_data = {
                'sql': sql,
                'thought': thought,
            }
            
            yield TextToSqlStreamEvent(
                event_type='result',
                content='SQL 生成完成',
                is_finished=True,
                data=result_data,
            )
            
            return NodeResult(
                success=True,
                output=result_data,
                output_variables={
                    f'{output_variable}_sql': sql,
                    f'{output_variable}_thought': thought,
                },
                tokens_used=total_tokens,
                elapsed_time=elapsed_time,
                metadata={
                    'thought': thought,
                    'suggested_next_node': 'db_sql',
                },
            )
            
        except Exception as e:
            logger.exception(f'Text-to-SQL 流式执行失败: {e}')
            yield TextToSqlStreamEvent(
                event_type='error',
                content=str(e),
                is_finished=True,
            )
            return NodeResult(
                success=False,
                error=str(e),
                elapsed_time=int((time.time() - start_time) * 1000),
            )
    
    async def _get_schema_context(
        self, 
        db_connection: str, 
        schema_name: str = 'public',
        selected_tables: List[str] = None,
        include_relations: bool = True,
        manual_relations: List[Dict[str, Any]] = None
    ) -> str:
        """获取数据库 Schema 上下文"""
        try:
            from core.database_manager.service import AsyncDatabaseManagerService
            
            # 创建数据库服务实例
            db_service = await AsyncDatabaseManagerService.create(db_connection)
            schema_name = self._resolve_schema_name(db_service, schema_name)
            
            # 获取表列表
            if selected_tables and len(selected_tables) > 0:
                tables = selected_tables
            else:
                # 获取所有表
                tables_info = await db_service.get_tables(schema_name=schema_name)
                tables = [
                    t.get('table_name') or t.get('name')
                    for t in tables_info[:20]
                    if t.get('table_name') or t.get('name')
                ]  # 限制最多 20 个表
            
            schema_parts = []
            table_relations = []
            
            for table_name in tables:
                try:
                    columns = await db_service.get_table_columns(table_name, schema_name)
                    
                    col_desc = []
                    for col in columns:
                        col_name = col.get('column_name') or col.get('name', '')
                        col_type = col.get('data_type') or col.get('type', '')
                        col_str = f"  - {col_name} ({col_type})"
                        if col.get('description') or col.get('comment'):
                            col_str += f" -- {col.get('description') or col.get('comment')}"
                        if col.get('is_primary_key'):
                            col_str += " [PK]"
                        if col.get('is_foreign_key'):
                            col_str += " [FK]"
                        col_desc.append(col_str)
                    
                    full_table_name = self._format_schema_table_name(
                        db_service, schema_name, table_name
                    )
                    table_schema = f"表: {full_table_name}\n" + "\n".join(col_desc)
                    schema_parts.append(table_schema)
                    
                    # 获取外键关系（从数据库）
                    if include_relations:
                        try:
                            constraints = await db_service.get_table_constraints(table_name, schema_name)
                            for constraint in constraints:
                                constraint_type = (
                                    constraint.get('constraint_type')
                                    or constraint.get('type', '')
                                )
                                if constraint_type != 'FOREIGN KEY':
                                    continue
                                source_full_name = self._format_schema_table_name(
                                    db_service, schema_name, table_name
                                )
                                foreign_table = (
                                    constraint.get('referenced_table')
                                    or constraint.get('foreign_table')
                                    or ''
                                )
                                target_full_name = self._format_schema_table_name(
                                    db_service, schema_name, foreign_table
                                )
                                source_columns = (
                                    constraint.get('columns')
                                    or constraint.get('column')
                                    or ''
                                )
                                foreign_columns = (
                                    constraint.get('referenced_columns')
                                    or constraint.get('foreign_column')
                                    or ''
                                )
                                relation = (
                                    f"  {source_full_name}.{source_columns} → "
                                    f"{target_full_name}.{foreign_columns}"
                                )
                                table_relations.append(relation)
                        except Exception as e:
                            logger.debug(f'获取表 {table_name} 外键关系失败: {e}')
                    
                except Exception as e:
                    logger.warning(f'获取表 {table_name} 结构失败: {e}')
            
            # 添加手动指定的表关系（逻辑外键）
            if manual_relations:
                for rel in manual_relations:
                    source_table = rel.get('sourceTable', '')
                    source_field = rel.get('sourceField', '')
                    target_table = rel.get('targetTable', '')
                    target_field = rel.get('targetField', '')
                    if source_table and source_field and target_table and target_field:
                        relation = f"  {source_table}.{source_field} → {target_table}.{target_field} (逻辑外键)"
                        table_relations.append(relation)
            
            # 组装最终的 Schema 上下文
            result = "\n\n".join(schema_parts)
            
            if table_relations:
                result += "\n\n## 表关系（外键）\n" + "\n".join(table_relations)
            
            return result
            
        except Exception as e:
            logger.exception(f'获取 Schema 上下文失败: {e}')
            return ''
    
    @staticmethod
    def _resolve_schema_name(db_service, schema_name: str) -> str:
        """按连接类型解析有效 schema，避免非 PG 库误用 public。"""
        if db_service._uses_schema_layer():
            if not schema_name or (
                schema_name == 'public' and db_service.db_type != 'postgresql'
            ):
                return db_service._default_schema()
            return schema_name
        return schema_name or ''

    @staticmethod
    def _format_schema_table_name(db_service, schema_name: str, table_name: str) -> str:
        if not table_name:
            return table_name
        if db_service._uses_schema_layer() and schema_name:
            return f"{schema_name}.{table_name}"
        return table_name
    
    async def _get_db_type(self, db_connection: str) -> str:
        """获取数据库类型"""
        try:
            from core.database_manager.service import AsyncDatabaseManagerService
            db_service = await AsyncDatabaseManagerService.create(db_connection)
            return db_service.db_type.upper()
        except Exception:
            return 'PostgreSQL'
    
    def _call_llm_with_function_calling(
        self,
        llm_service,
        model_id: str,
        db_type: str,
        schema_context: str,
        user_question: str
    ) -> Generator[TextToSqlStreamEvent, None, Tuple[Dict[str, Any], int]]:
        """
        使用 Function Calling 调用 LLM 生成 SQL
        
        Args:
            llm_service: LLM 服务实例
            model_id: 模型 ID
            db_type: 数据库类型
            schema_context: Schema 上下文
            user_question: 用户问题
            
        Yields:
            TextToSqlStreamEvent: 流式输出事件
            
        Returns:
            tuple: (llm_result, total_tokens)
        """
        from datetime import datetime
        
        system_prompt = TEXT_TO_SQL_SYSTEM_PROMPT_FC.format(
            db_type=db_type,
            current_date=datetime.now().strftime('%Y-%m-%d'),
            schema_context=schema_context,
        )
        
        messages = [
            {'role': 'system', 'content': system_prompt},
            {'role': 'user', 'content': f'用户问题: {user_question}'},
        ]
        
        # 定义工具
        tools = [TEXT_TO_SQL_TOOL]
        
        # 累积工具调用参数
        accumulated_arguments = ''
        total_tokens = 0
        tool_call_received = False
        
        for chunk in llm_service.chat_stream_sync(
            model_id=model_id,
            messages=messages,
            temperature=0.1,
            max_tokens=2048,
            tools=tools,
            tool_choice='required',  # 强制使用工具
        ):
            # 处理工具调用增量
            if chunk.tool_call_delta:
                delta_args = chunk.tool_call_delta.get('arguments', '')
                if delta_args:
                    accumulated_arguments += delta_args
                    # 流式输出工具调用参数
                    yield TextToSqlStreamEvent(
                        event_type='sql_chunk',
                        content=delta_args,
                        is_finished=False,
                    )
            
            # 处理完整的工具调用
            if chunk.tool_calls:
                tool_call_received = True
                for tool_call in chunk.tool_calls:
                    if tool_call.name == 'generate_sql':
                        # 解析工具调用参数
                        arguments = tool_call.arguments
                        if isinstance(arguments, str):
                            arguments = json.loads(arguments)
                        
                        yield TextToSqlStreamEvent(
                            event_type='sql_chunk',
                            content='',
                            is_finished=True,
                        )
                        
                        return (arguments, total_tokens)
            
            # 记录 token 使用
            if chunk.is_finished and chunk.total_tokens > 0:
                total_tokens = chunk.total_tokens
        
        # 如果没有收到工具调用，尝试从累积的参数中解析
        if accumulated_arguments and not tool_call_received:
            try:
                arguments = json.loads(accumulated_arguments)
                yield TextToSqlStreamEvent(
                    event_type='sql_chunk',
                    content='',
                    is_finished=True,
                )
                return (arguments, total_tokens)
            except json.JSONDecodeError:
                pass
        
        # Function Calling 失败
        raise ValueError('Function Calling 未返回有效的工具调用')
    
    def _parse_llm_response(self, content: str) -> Optional[Dict[str, Any]]:
        """解析 LLM 响应（Fallback 方式）"""
        try:
            # 尝试直接解析 JSON
            return json.loads(content)
        except json.JSONDecodeError:
            pass
        
        # 尝试从 markdown 代码块中提取 JSON
        import re
        json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
        if json_match:
            try:
                return json.loads(json_match.group(1))
            except json.JSONDecodeError:
                pass
        
        # 尝试提取 {...} 部分
        brace_match = re.search(r'\{[\s\S]*\}', content)
        if brace_match:
            try:
                return json.loads(brace_match.group(0))
            except json.JSONDecodeError:
                pass
        
        logger.warning(f'无法解析 LLM 响应: {content[:500]}')
        return None
    
    def _validate_sql(self, sql: str) -> bool:
        """验证 SQL 安全性"""
        if not sql:
            logger.warning("SQL 验证失败: SQL 为空")
            return False
        
        # 去除前后空白和可能的 Markdown 代码块标记
        sql_cleaned = sql.strip()
        if sql_cleaned.startswith('```'):
            # 移除 Markdown 代码块
            lines = sql_cleaned.split('\n')
            # 移除第一行（```sql 或 ```）
            if lines:
                lines = lines[1:]
            # 移除最后一行（```）
            if lines and lines[-1].strip() == '```':
                lines = lines[:-1]
            sql_cleaned = '\n'.join(lines).strip()
        
        sql_upper = sql_cleaned.upper().strip()
        
        # 只允许 SELECT 语句（也允许 WITH ... SELECT 即 CTE）
        if not sql_upper.startswith('SELECT') and not sql_upper.startswith('WITH'):
            logger.warning(f"SQL 验证失败: 不是 SELECT/WITH 语句, SQL 开头: {sql_upper[:50]}")
            return False
        
        # 禁止危险关键字（作为独立语句，不在子查询或 CTE 中）
        dangerous_keywords = [
            'INSERT', 'UPDATE', 'DELETE', 'DROP', 'TRUNCATE',
            'ALTER', 'CREATE', 'GRANT', 'REVOKE', 'EXEC', 'EXECUTE',
        ]
        
        for keyword in dangerous_keywords:
            # 检查是否作为独立关键字出现（前后有空格或在开头/结尾）
            if re.search(rf'\b{keyword}\b', sql_upper):
                logger.warning(f"SQL 验证失败: 包含危险关键字 {keyword}, SQL: {sql_cleaned[:200]}")
                return False
        
        # 禁止多语句（分号后还有内容）
        if ';' in sql_cleaned:
            parts = sql_cleaned.split(';')
            if any(p.strip() for p in parts[1:]):
                logger.warning(f"SQL 验证失败: 包含多条语句")
                return False
        
        return True
    
    def _format_sql(self, sql: str) -> str:
        """
        格式化 SQL 语句
        
        Args:
            sql: 原始 SQL 语句
            
        Returns:
            格式化后的 SQL 语句
        """
        try:
            # 去除可能的 Markdown 代码块标记
            sql_cleaned = sql.strip()
            if sql_cleaned.startswith('```'):
                lines = sql_cleaned.split('\n')
                if lines:
                    lines = lines[1:]
                if lines and lines[-1].strip() == '```':
                    lines = lines[:-1]
                sql_cleaned = '\n'.join(lines).strip()
            
            # 使用 sqlparse 格式化
            formatted_sql = sqlparse.format(
                sql_cleaned,
                reindent=True,              # 重新缩进
                keyword_case='upper',       # 关键字大写
                identifier_case='lower',    # 标识符小写
                strip_comments=False,       # 保留注释
                use_space_around_operators=True,  # 操作符周围加空格
            )
            
            return formatted_sql.strip()
        except Exception as e:
            logger.warning(f"SQL 格式化失败: {e}, 返回原始 SQL")
            return sql
