"""
Enhanced SQL Database Agent Implementation for AI-Parrot.

Concrete implementation of AbstractDbAgent for SQL databases with support for:
- PostgreSQL, MySQL, and SQL Server
- Dictionary and string credentials
- Dual DSN generation for SQLAlchemy and asyncdb
- DatabaseQueryTool integration for query validation and execution
"""

from typing import Dict, Any, List, Optional, Union
import re
from urllib.parse import urlparse, quote_plus
from datetime import datetime
import pandas as pd
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from sqlalchemy import text
from .abstract import AbstractDBAgent
from .tools import DatabaseSchema, TableMetadata
from ...models import AIMessage
from ...tools.databasequery import DatabaseQueryTool
from ...tools import ToolResult


class SQLAgent(AbstractDBAgent):
    """
    SQL Database Agent with dual DSN support and DatabaseQueryTool integration.

    Supports PostgreSQL, MySQL, and SQL Server with both dictionary and string credentials.
    """

    # Database flavor mappings for SQLAlchemy
    SQLALCHEMY_DIALECT_MAPPING = {
        'postgresql': 'postgresql+asyncpg',
        'pg': 'postgresql+asyncpg',
        'postgres': 'postgresql+asyncpg',
        'mysql': 'mysql+aiomysql',
        'sqlserver': 'mssql+aioodbc',
        'mssql': 'mssql+aioodbc'
    }

    # Default ports for databases
    DEFAULT_PORTS = {
        'postgresql': 5432,
        'postgres': 5432,
        'mysql': 3306,
        'sqlserver': 1433,
        'mssql': 1433
    }

    def __init__(
        self,
        name: str = "SQLAgent",
        credentials: Union[str, Dict[str, Any]] = None,
        database_flavor: str = "postgresql",
        schema_name: str = "public",
        max_sample_rows: int = 2,
        **kwargs
    ):
        """
        Initialize SQL Database Agent.

        Args:
            name: Agent name
            credentials: Connection credentials (dict or connection string)
            database_flavor: Database type (postgresql, mysql, sqlserver)
            schema_name: Target schema name
            max_sample_rows: Maximum rows to sample from each table
        """
        self.database_flavor = database_flavor.lower()
        self.max_sample_rows = max_sample_rows
        self.async_session_maker = None

        # DSN strings for different purposes
        self.discovery_dsn = None  # SQLAlchemy format for schema discovery
        self.dsn = None            # asyncdb format for DatabaseQueryTool
        self.credentials = None
        self.connection_dict = None
        if isinstance(credentials, dict):
            self.connection_dict = credentials

        # Validate database flavor
        if self.database_flavor not in self.SQLALCHEMY_DIALECT_MAPPING:
            raise ValueError(
                f"Unsupported database flavor: {database_flavor}"
            )
        
        # Force low temperature to minimize hallucinations
        kwargs['temperature'] = kwargs.get('temperature', 0.0)

        super().__init__(
            name=name,
            credentials=credentials,
            schema_name=schema_name,
            **kwargs
        )

        # Process credentials and generate DSNs
        self._process_credentials(credentials)

        # Add SQL-specific tools
        self._setup_sql_tools()

    def _dsn_for_sqlalchemy(self, connection_string: str) -> str:
        """Adapt connection string for SQLAlchemy async drivers."""
        parsed = urlparse(connection_string)

        if parsed.scheme.startswith('postgresql') and '+asyncpg' not in parsed.scheme:
            return connection_string.replace('postgresql://', 'postgresql+asyncpg://')
        elif parsed.scheme.startswith('postgres') and '+asyncpg' not in parsed.scheme:
            return connection_string.replace('postgres://', 'postgresql+asyncpg://')
        elif parsed.scheme.startswith('mysql') and '+aiomysql' not in parsed.scheme:
            return connection_string.replace('mysql://', 'mysql+aiomysql://')
        elif parsed.scheme.startswith('mssql') and '+aioodbc' not in parsed.scheme:
            return connection_string.replace('mssql://', 'mssql+aioodbc://')

        return connection_string

    def _dsn_for_asyncdb(self, connection_string: str) -> str:
        """Adapt connection string for asyncdb format."""
        parsed = urlparse(connection_string)

        # Check if already in asyncdb format:
        if parsed.scheme in ['postgres', 'mysql', 'mssql']:
            return connection_string

        # Convert SQLAlchemy formats to asyncdb formats
        if parsed.scheme.startswith('postgresql'):
            return connection_string.replace(
                'postgresql+asyncpg://', 'postgres://'
            ).replace('postgresql://', 'postgres://')
        elif parsed.scheme.startswith('mysql'):
            return connection_string.replace(
                'mysql+aiomysql://', 'mysql://'
            ).replace('mysql://', 'mysql://')
        elif parsed.scheme.startswith('mssql'):
            return connection_string.replace(
                'mssql+aioodbc://', 'mssql://'
            ).replace('mssql://', 'mssql://')

        return connection_string

    def _process_credentials(self, credentials: Union[str, Dict[str, Any]]) -> None:
        """
        Process credentials and generate both discovery_dsn and dsn.

        Args:
            credentials: Either connection string or dictionary with connection params
        """
        if isinstance(credentials, str):
            # Connection string provided
            self.connection_string = credentials
            self.discovery_dsn = self._dsn_for_sqlalchemy(credentials)
            self.dsn = self._dsn_for_asyncdb(credentials)
            self.credentials = {}
        elif isinstance(credentials, dict):
            # Dictionary credentials provided
            self.connection_dict = credentials
            self.discovery_dsn = self._build_sqlalchemy_dsn_from_dict(credentials)
            self.dsn = self._build_asyncdb_dsn_from_dict(credentials)
            self.connection_string = self.discovery_dsn
            self.credentials = credentials
        else:
            raise ValueError(
                "Credentials must be either a connection string or dictionary"
            )

    def _build_sqlalchemy_dsn_from_dict(self, creds: Dict[str, Any]) -> str:
        """
        Build SQLAlchemy DSN from credentials dictionary.

        Args:
            creds: Dictionary with keys like host, port, database, username, password

        Returns:
            SQLAlchemy-compatible connection string
        """
        # Extract credentials with defaults
        host = creds.get('host', 'localhost')
        port = creds.get('port', self.DEFAULT_PORTS.get(self.database_flavor, 5432))
        database = creds.get('database', creds.get('dbname', 'postgres'))
        username = creds.get('username', creds.get('user', 'postgres'))
        password = creds.get('password', creds.get('pwd', ''))

        # URL encode password to handle special characters
        encoded_password = quote_plus(str(password)) if password else ''

        # Get SQLAlchemy dialect
        dialect = self.SQLALCHEMY_DIALECT_MAPPING[self.database_flavor]

        # Build connection string
        if encoded_password:
            dsn = f"{dialect}://{username}:{encoded_password}@{host}:{port}/{database}"
        else:
            dsn = f"{dialect}://{username}@{host}:{port}/{database}"

        # Add any additional parameters
        params = []
        for key, value in creds.items():
            if key not in ['host', 'port', 'database', 'dbname', 'username', 'user', 'password', 'pwd']:
                params.append(f"{key}={value}")

        if params:
            dsn += "?" + "&".join(params)

        return dsn

    def _build_asyncdb_dsn_from_dict(self, creds: Dict[str, Any]) -> str:
        """
        Build asyncdb DSN from credentials dictionary.

        Args:
            creds: Dictionary with connection parameters

        Returns:
            asyncdb-compatible connection string (postgres://...)
        """
        # Extract credentials
        host = creds.get('host', 'localhost')
        port = creds.get('port', self.DEFAULT_PORTS.get(self.database_flavor, 5432))
        database = creds.get('database', creds.get('dbname', 'postgres'))
        username = creds.get('username', creds.get('user', 'postgres'))
        password = creds.get('password', creds.get('pwd', ''))

        # URL encode password
        encoded_password = quote_plus(str(password)) if password else ''

        # Get asyncdb scheme (postgres for PostgreSQL regardless of flavor name)
        if self.database_flavor in ['postgresql', 'postgres']:
            scheme = 'postgres'
        elif self.database_flavor == 'mysql':
            scheme = 'mysql'
        elif self.database_flavor in ['sqlserver', 'mssql']:
            scheme = 'mssql'
        else:
            scheme = 'postgres'  # Default fallback

        # Build DSN
        if encoded_password:
            dsn = f"{scheme}://{username}:{encoded_password}@{host}:{port}/{database}"
        else:
            dsn = f"{scheme}://{username}@{host}:{port}/{database}"

        return dsn

    def _setup_sql_tools(self):
        """Setup SQL-specific tools including DatabaseQueryTool."""
        # The DatabaseQueryTool should already be registered in the parent class
        # We just need to ensure it's configured properly
        pass

    async def connect_database(self) -> None:
        """Connect to the SQL database using SQLAlchemy async engine."""
        if not self.discovery_dsn:
            raise ValueError("Discovery DSN is required")

        try:
            # Create async engine for schema discovery
            self.engine = create_async_engine(
                self.discovery_dsn,
                echo=False,
                pool_pre_ping=True,
                pool_recycle=3600
            )

            # Create session maker
            self.async_session_maker = sessionmaker(
                self.engine,
                class_=AsyncSession,
                expire_on_commit=False
            )

            # Test connection
            async with self.engine.begin() as conn:
                await conn.execute(text("SELECT 1"))

            self.logger.info(
                f"Successfully connected to {self.database_flavor} database using SQLAlchemy"
            )

            # Test DatabaseQueryTool connection
            await self._test_database_query_tool()

        except Exception as e:
            self.logger.error(f"Failed to connect to database: {e}")
            raise

    async def _test_database_query_tool(self) -> None:
        """Test DatabaseQueryTool connection."""
        try:
            # Get database query tool from registered tools
            db_tool = self.tool_manager.get_tool('database_query')
            if db_tool:
                # Test with a simple query
                test_result = await db_tool.execute(
                    driver='pg' if self.database_flavor in ['postgresql', 'postgres', 'pg'] else self.database_flavor,
                    query="SELECT 1 as test_column LIMIT 1",
                    dsn=self.dsn,
                    credentials=self.credentials or None,
                    output_format='native'
                )

                if test_result.status == "success":
                    self.logger.debug(
                        "DatabaseQueryTool connection test successful"
                    )
                else:
                    self.logger.warning(
                        f"DatabaseQueryTool test failed: {test_result.error}"
                    )
            else:
                self.logger.warning(
                    "DatabaseQueryTool not found in registered tools"
                )

        except Exception as e:
            self.logger.warning(
                f"DatabaseQueryTool test failed: {e}"
            )

    async def extract_schema_metadata(self) -> DatabaseSchema:
        """Extract complete schema metadata from SQL database."""
        if not self.engine:
            await self.connect_database()

        try:
            async with self.engine.begin() as conn:
                # Get database name
                db_name_query = await self._get_database_name_query()
                result = await conn.execute(text(db_name_query))
                database_name = result.scalar()

                # Extract tables metadata
                tables = await self._extract_tables_metadata(conn)

                # Extract views metadata (simplified for now)
                views = []

                schema_metadata = DatabaseSchema(
                    database_name=database_name or "unknown",
                    database_type=self.database_flavor,
                    tables=tables,
                    views=views,
                    functions=[],
                    procedures=[],
                    metadata={
                        "schema_name": self.schema_name,
                        "extraction_timestamp": datetime.now().isoformat(),
                        "total_tables": len(tables),
                        "total_views": len(views),
                        "discovery_dsn": self.discovery_dsn,
                        "asyncdb_dsn": self.dsn
                    }
                )

                self.logger.info(
                    f"Extracted metadata for {len(tables)} tables"
                )

                return schema_metadata

        except Exception as e:
            self.logger.error(f"Failed to extract schema metadata: {e}")
            raise

    async def _get_database_name_query(self) -> str:
        """Get database name query based on database flavor."""
        if self.database_flavor in ['postgresql', 'postgres']:
            return "SELECT current_database()"
        elif self.database_flavor == 'mysql':
            return "SELECT database()"
        elif self.database_flavor in ['sqlserver', 'mssql']:
            return "SELECT DB_NAME()"
        else:
            return "SELECT 'unknown' as database_name"

    async def _extract_tables_metadata(self, conn) -> List[TableMetadata]:
        """Extract metadata for all tables in the schema."""
        tables = []

        # Get table names
        if self.database_flavor in ['postgresql', 'postgres']:
            table_query = """
                SELECT table_name
                FROM information_schema.tables
                WHERE table_schema = :schema_name
                AND table_type = 'BASE TABLE'
                ORDER BY table_name
            """
        elif self.database_flavor == 'mysql':
            table_query = """
                SELECT table_name
                FROM information_schema.tables
                WHERE table_schema = :schema_name
                AND table_type = 'BASE TABLE'
                ORDER BY table_name
            """
        else:  # SQL Server
            table_query = """
                SELECT table_name
                FROM information_schema.tables
                WHERE table_schema = :schema_name
                AND table_type = 'BASE TABLE'
                ORDER BY table_name
            """

        result = await conn.execute(
            text(table_query), {"schema_name": self.schema_name}
        )
        table_rows = result.fetchall()

        for row in table_rows:
            table_name = row[0]
            # Extract detailed table metadata
            table_metadata = await self._extract_single_table_metadata(conn, table_name)
            tables.append(table_metadata)

        return tables

    async def _extract_single_table_metadata(self, conn, table_name: str) -> TableMetadata:
        """Extract detailed metadata for a single table."""
        # Get column information
        columns = await self._get_table_columns(conn, table_name)

        # Get primary keys
        primary_keys = await self._get_primary_keys(conn, table_name)

        # Get foreign keys
        foreign_keys = await self._get_foreign_keys(conn, table_name)

        # Get sample data using DatabaseQueryTool
        sample_data = await self._get_sample_data_via_tool(table_name)

        return TableMetadata(
            name=table_name,
            schema=self.schema_name,
            columns=columns,
            primary_keys=primary_keys,
            foreign_keys=foreign_keys,
            indexes=[],  # Simplified for now
            description=None,  # Simplified for now
            sample_data=sample_data
        )

    async def _get_table_columns(self, conn, table_name: str) -> List[Dict[str, Any]]:
        """Get column information for a table."""
        if self.database_flavor in ['postgresql', 'postgres']:
            query = """
                SELECT
                    column_name,
                    data_type,
                    is_nullable,
                    column_default,
                    character_maximum_length,
                    numeric_precision,
                    numeric_scale
                FROM information_schema.columns
                WHERE table_schema = :schema_name
                AND table_name = :table_name
                ORDER BY ordinal_position
            """
        elif self.database_flavor == 'mysql':
            query = """
                SELECT
                    column_name,
                    data_type,
                    is_nullable,
                    column_default,
                    character_maximum_length,
                    numeric_precision,
                    numeric_scale
                FROM information_schema.columns
                WHERE table_schema = :schema_name
                AND table_name = :table_name
                ORDER BY ordinal_position
            """
        else:  # SQL Server
            query = """
                SELECT
                    column_name,
                    data_type,
                    is_nullable,
                    column_default,
                    character_maximum_length,
                    numeric_precision,
                    numeric_scale
                FROM information_schema.columns
                WHERE table_schema = :schema_name
                AND table_name = :table_name
                ORDER BY ordinal_position
            """

        result = await conn.execute(text(query), {
            "schema_name": self.schema_name,
            "table_name": table_name
        })

        columns = []
        for row in result.fetchall():
            columns.append({
                "name": row[0],
                "type": row[1],
                "nullable": row[2] == "YES",
                "default": row[3],
                "max_length": row[4],
                "precision": row[5],
                "scale": row[6]
            })

        return columns

    async def _get_primary_keys(self, conn, table_name: str) -> List[str]:
        """Get primary key columns for a table."""
        if self.database_flavor in ['postgresql', 'postgres']:
            query = """
                SELECT column_name
                FROM information_schema.key_column_usage
                WHERE table_schema = :schema_name
                AND table_name = :table_name
                AND constraint_name IN (
                    SELECT constraint_name
                    FROM information_schema.table_constraints
                    WHERE table_schema = :schema_name
                    AND table_name = :table_name
                    AND constraint_type = 'PRIMARY KEY'
                )
                ORDER BY ordinal_position
            """
        else:  # MySQL and SQL Server
            query = """
                SELECT column_name
                FROM information_schema.key_column_usage
                WHERE table_schema = :schema_name
                AND table_name = :table_name
                AND constraint_name = 'PRIMARY'
                ORDER BY ordinal_position
            """

        result = await conn.execute(text(query), {
            "schema_name": self.schema_name,
            "table_name": table_name
        })

        return [row[0] for row in result.fetchall()]

    async def _get_foreign_keys(self, conn, table_name: str) -> List[Dict[str, Any]]:
        """Get foreign key information for a table."""
        query = """
            SELECT
                kcu.column_name,
                ccu.table_schema AS referenced_table_schema,
                ccu.table_name AS referenced_table_name,
                ccu.column_name AS referenced_column_name
            FROM information_schema.key_column_usage kcu
            JOIN information_schema.constraint_column_usage ccu
                ON kcu.constraint_name = ccu.constraint_name
            WHERE kcu.table_schema = :schema_name
            AND kcu.table_name = :table_name
            AND kcu.constraint_name IN (
                SELECT constraint_name
                FROM information_schema.table_constraints
                WHERE table_schema = :schema_name
                AND table_name = :table_name
                AND constraint_type = 'FOREIGN KEY'
            )
        """

        result = await conn.execute(text(query), {
            "schema_name": self.schema_name,
            "table_name": table_name
        })

        foreign_keys = []
        for row in result.fetchall():
            foreign_keys.append({
                "column": row[0],
                "referenced_table_schema": row[1],
                "referenced_table": row[2],
                "referenced_column": row[3]
            })

        return foreign_keys

    async def _get_sample_data_via_tool(self, table_name: str) -> List[Dict[str, Any]]:
        """Get sample data using DatabaseQueryTool."""
        try:
            # Get database query tool
            db_tool = self.tool_manager.get_tool('database_query')
            if not db_tool:
                self.logger.warning("DatabaseQueryTool not found")
                return []

            # Build sample query
            full_table_name = f'"{self.schema_name}"."{table_name}"' if self.schema_name != 'public' else f'"{table_name}"'
            sample_query = f"SELECT * FROM {full_table_name} LIMIT {self.max_sample_rows}"

            # Execute query
            result = await db_tool.execute(
                driver='pg' if self.database_flavor in ['postgresql', 'postgres'] else self.database_flavor,
                query=sample_query,
                dsn=self.dsn,
                credentials=self.connection_dict,
                output_format='json'
            )
            if result.status == "success":
                return result.result
            else:
                self.logger.warning(f"Could not get sample data for {table_name}: {result.error}")
                return []

        except Exception as e:
            self.logger.warning(f"Error getting sample data for {table_name}: {e}")
            return []

    async def generate_query(
        self,
        natural_language_query: str,
        target_tables: Optional[List[str]] = None,
        query_type: str = "SELECT"
    ) -> Dict[str, Any]:
        """Generate SQL query from natural language and validate it."""
        try:
            # Get schema context
            schema_context = await self._get_schema_context_for_query(
                natural_language_query, target_tables
            )

            # Build prompt for LLM
            prompt = self._build_query_generation_prompt(
                natural_language_query=natural_language_query,
                schema_context=schema_context,
                query_type=query_type,
                database_flavor=self.database_flavor
            )

            # Generate query using LLM
            response = await self._llm.ask(
                prompt=prompt,
                model=self._llm_model,
                temperature=0.0,  # Zero temperature for deterministic results
                use_tools=False, # Explicitly disable tools to prevent recursion
                tools=[]
            )

            # Extract SQL query from response
            generated_query = self._extract_sql_from_response(str(response.output))

            # Validate query using DatabaseQueryTool with LIMIT 0
            validation_result = await self._validate_query_with_tool(generated_query)

            result = {
                "query": generated_query,
                "query_type": query_type,
                "tables_used": self._extract_tables_from_query(generated_query),
                "schema_context_used": len(schema_context),
                "validation": validation_result,
                "natural_language_input": natural_language_query
            }

            return result

        except Exception as e:
            self.logger.error(f"Failed to generate query: {e}")
            raise

    async def _validate_query_with_tool(self, query: str) -> Dict[str, Any]:
        """Validate query using DatabaseQueryTool with LIMIT 0."""
        try:
            # Get database query tool
            db_tool = None
            for tool in self.tools:
                if isinstance(tool, DatabaseQueryTool):
                    db_tool = tool
                    break

            if not db_tool:
                return {
                    "valid": False,
                    "error": "DatabaseQueryTool not available",
                    "method": "tool_validation"
                }

            # Modify query to add LIMIT 0 for validation (no data returned)
            if query.strip().upper().startswith('SELECT'):
                validation_query = f"SELECT * FROM ({query.rstrip(';')}) AS validation_subquery LIMIT 0"
            else:
                # For non-SELECT queries, we can't easily validate without risk
                validation_query = query

            # Execute validation query
            result = await db_tool.execute(
                driver='pg' if self.database_flavor in ['postgresql', 'postgres'] else self.database_flavor,
                query=validation_query,
                dsn=self.dsn,
                credentials=self.connection_dict,
                output_format='native'
            )

            return {
                "valid": result.status == "success",
                "error": result.error if result.status == "error" else None,
                "method": "database_query_tool",
                "validation_query": validation_query
            }

        except Exception as e:
            return {
                "valid": False,
                "error": str(e),
                "method": "tool_validation"
            }

    async def explain_query(self, query: str) -> str:
        """
        Explain a database query (e.g. EXPLAIN ANALYZE).
        
        Args:
            query: The SQL query to explain
            
        Returns:
            The execution plan as a string
        """
        try:
            # Construct EXPLAIN query based on flavor
            if self.database_flavor in ['postgresql', 'postgres', 'pg']:
                # Use JSON format for better parsing if needed, and ANALYZE for actual execution stats
                explain_query = f"EXPLAIN (FORMAT JSON, ANALYZE) {query}"
            elif self.database_flavor == 'mysql':
                explain_query = f"EXPLAIN ANALYZE {query}"
            else:
                explain_query = f"EXPLAIN {query}"

            # Execute the explain query
            # We use execute_query but need to handle the result format
            result = await self.execute_query(explain_query, limit=0) # limit=0 is ignored for EXPLAIN usually

            if result["success"]:
                 # Format the result
                 data = result["data"]
                 if self.database_flavor in ['postgresql', 'postgres', 'pg']:
                     # Postgres JSON output usually comes as a single cell with lists
                     try:
                         # It might be a list of dicts in the first column
                         plan = data.iloc[0, 0]
                         if isinstance(plan, list) or isinstance(plan, dict):
                             return json.dumps(plan, indent=2)
                         return str(plan)
                     except Exception:
                         return data.to_string()
                 else:
                     return data.to_string()
            else:
                return f"Failed to explain query: {result['error']}"

        except Exception as e:
            self.logger.error(f"Error explaining query: {e}")
            return f"Error explaining query: {str(e)}"

    async def execute_query(self, query: str, limit: int = 200) -> Dict[str, Any]:
        """Execute SQL query and return results using DatabaseQueryTool."""
        try:
            # Get database query tool
            db_tool = self.tool_manager.get_tool('database_query')
            if not db_tool:
                db_tool = None

            if not db_tool:
                return {
                    "success": False,
                    "error": "DatabaseQueryTool not available",
                    "query": query
                }

            # Add limit for SELECT queries if not present
            execution_query = query
            result = None
            if query.strip().upper().startswith('SELECT') and 'LIMIT' not in query.upper():
                execution_query = f"{query.rstrip(';')} LIMIT {limit}"

            # Execute query (return a ToolResult)
            result = await db_tool.execute(
                driver='pg' if self.database_flavor in ['postgresql', 'postgres'] else self.database_flavor,
                query=execution_query,
                dsn=self.dsn,
                credentials=self.connection_dict,
                output_format='pandas'
            )

            if result.status == "success":
                data = result.result
                columns = data.columns.tolist() if not data.empty else []
                row_count = len(data) if not data.empty else 0
                return {
                    "success": True,
                    "data": data,
                    "columns": columns,
                    "row_count": row_count,
                    "query": execution_query,
                    "tool_used": "DatabaseQueryTool",
                    "raw_result": result
                }
            else:
                return {
                    "success": False,
                    "error": result.error,
                    "query": execution_query,
                    "tool_used": "DatabaseQueryTool",
                    "raw_result": result
                }

        except Exception as e:
            self.logger.error(f"Query execution failed: {e}")
            return {
                "success": False,
                "error": str(e),
                "query": query,
                "tool_used": "DatabaseQueryTool",
                "raw_result": None
            }

    async def _get_schema_context_for_query(
        self,
        natural_language_query: str,
        target_tables: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """Get relevant schema context for query generation."""
        if target_tables:
            context = []
            for table_name in target_tables:
                table_info = await self.search_schema(
                    search_term=table_name,
                    search_type="tables",
                    limit=1
                )
                if table_info:
                    context.extend(table_info)
            return context
        else:
            return await self.search_schema(
                search_term=natural_language_query,
                search_type="all",
                limit=5
            )

    def _build_query_generation_prompt(
        self,
        natural_language_query: str,
        schema_context: List[Dict[str, Any]],
        query_type: str,
        database_flavor: str
    ) -> str:
        """Build prompt for LLM query generation."""
        prompt = f"""
You are an expert SQL developer working with a {database_flavor} database.
Generate a clean, efficient {query_type} SQL query based on the natural language request and schema information.

Natural Language Request: {natural_language_query}

Available Schema Information:
"""

        for i, context in enumerate(schema_context[:3], 1):
            prompt += f"\n{i}. {context.get('content', '')}\n"

        prompt += f"""
Requirements:
1. Generate valid {database_flavor} SQL with clean formatting
2. Use appropriate {database_flavor} syntax and functions
3. Use simple column names unless JOINs require qualification
4. Use table aliases for readability in JOINs
5. Only use double quotes for identifiers with special characters
6. Include appropriate WHERE clauses and filters
7. Optimize for performance and readability
8. Return ONLY the SQL query without explanations or formatting

Query Type: {query_type}
Database: {database_flavor}

SQL Query:"""

        return prompt

    def _extract_sql_from_response(self, response_text: str) -> str:
        """Extract SQL query from LLM response."""
        # Remove markdown code blocks if present
        if "```sql" in response_text:
            lines = response_text.split('\n')
            sql_lines = []
            in_sql_block = False

            for line in lines:
                if line.strip().startswith("```sql"):
                    in_sql_block = True
                    continue
                elif line.strip() == "```" and in_sql_block:
                    break
                elif in_sql_block:
                    sql_lines.append(line)

            return '\n'.join(sql_lines).strip()
        else:
            return response_text.strip()

    async def ask(
        self,
        question: str = None,
        user_context: str = "",
        context: str = "",
        return_results: bool = True,  # New parameter to control query execution
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        use_conversation_history: bool = True,
        **kwargs
    ) -> AIMessage:
        """
        Enhanced ask method that can automatically execute generated SQL queries.

        Args:
            question: The user's question about the database
            user_context: User-specific context for database interaction
            context: Additional context about data location, schema guidance
            return_results: If True, automatically execute generated SQL queries and return data
            session_id: Session identifier for conversation history
            user_id: User identifier
            use_conversation_history: Whether to use conversation history
            **kwargs: Additional arguments for LLM

        Returns:
            AIMessage: The response from the LLM, potentially enhanced with query results
        """
        # Backwards compatibility
        if question is None:
            question = kwargs.get('prompt')
            
        # First, get the standard response from the parent method
        response = await super().ask(
            question=question,
            user_context=user_context,
            context=context,
            session_id=session_id,
            user_id=user_id,
            use_conversation_history=use_conversation_history,
            **kwargs
        )

        # If return_results is False, return the response as-is
        if not return_results:
            return response

        # Try to extract and execute SQL queries from the response
        try:
            response_text = str(response.output) if response.output else ""

            # Extract SQL queries from the response
            sql_queries = self._extract_queries(response_text)

            if sql_queries:
                # Execute the first/main SQL query
                main_query = sql_queries[0]
                self.logger.debug(
                    f"Auto-executing extracted query: {main_query[:100]}..."
                )

                # Execute the query
                result = await self.execute_query(
                    query=main_query
                )
                # Preserve original response
                response.response = response_text
                # is the dataframe:
                response.output = result.get('data', None)
                response.raw_response = result # Preserve raw ToolResult

                # Add execution metadata if response has metadata attribute
                if hasattr(response, 'metadata') and response.metadata:
                    response.metadata.update({
                        'auto_executed_query': True,
                        'executed_query': main_query,
                        'execution_success': result.get('status') == 'success',
                        'row_count': result.get('row_count', 0),
                        'columns': result.get('columns', []),
                        'error': result.get('error', None)
                    })

        except Exception as e:
            self.logger.warning(
                f"Failed to auto-execute query: {e}"
            )
            # Don't fail the entire request, just log the warning
            # The user still gets the explanation even if execution fails

        return response

    async def search_schema(
        self,
        search_term: str,
        search_type: str = "all",
        limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Search the database schema using SQL queries against information_schema.
        
        Args:
            search_term: Term to search for (supports LIKE patterns implicitly)
            search_type: Type of search ('tables', 'columns', 'all')
            limit: Maximum number of results
            
        Returns:
            List of matching schema objects
        """
        results = []
        
        # Check cache first
        if self.cache:
            cached_results = await self.cache.get(search_term, search_type, limit)
            if cached_results is not None:
                self.logger.info(f"Schema search cache hit for term: {search_term}")
                return cached_results

        term_pattern = f"%{search_term}%"
        
        try:
            # Determine logic based on search_type
            search_tables = search_type in ["all", "tables"]
            search_columns = search_type in ["all", "columns"]
            
            # --- Search Tables ---
            if search_tables:
                if self.database_flavor in ['postgresql', 'postgres', 'pg']:
                    # Support schema.table search
                    query = """
                        SELECT table_schema, table_name, 'TABLE' as type
                        FROM information_schema.tables
                        WHERE (table_name ILIKE :term 
                           OR table_schema || '.' || table_name ILIKE :term)
                          AND table_schema NOT IN ('information_schema', 'pg_catalog')
                          AND table_type = 'BASE TABLE'
                        LIMIT :limit
                    """
                elif self.database_flavor == 'mysql':
                    query = """
                        SELECT table_schema, table_name, 'TABLE' as type
                        FROM information_schema.tables
                        WHERE (table_name LIKE :term 
                           OR CONCAT(table_schema, '.', table_name) LIKE :term)
                          AND table_schema = DATABASE()
                          AND table_type = 'BASE TABLE'
                        LIMIT :limit
                    """
                else: # Generic/SQL Server
                    query = """
                        SELECT table_schema, table_name, 'TABLE' as type
                        FROM information_schema.tables
                        WHERE table_name LIKE :term
                        LIMIT :limit
                    """
                
                if self.engine:
                    async with self.engine.connect() as conn:
                        result_proxy = await conn.execute(text(query), {"term": term_pattern, "limit": limit})
                        rows = result_proxy.fetchall()
                        for row in rows:
                            results.append({
                                "type": "table",
                                "name": row[1],
                                "schema": row[0],
                                "description": f"Table: {row[0]}.{row[1]}"
                            })
            
            # --- Search Columns ---
            if search_columns and len(results) < limit:
                current_limit = limit - len(results)
                if self.database_flavor in ['postgresql', 'postgres', 'pg']:
                    query = """
                        SELECT table_schema, table_name, column_name, data_type
                        FROM information_schema.columns
                        WHERE column_name ILIKE :term
                          AND table_schema NOT IN ('information_schema', 'pg_catalog')
                        LIMIT :limit
                    """
                elif self.database_flavor == 'mysql':
                    query = """
                        SELECT table_schema, table_name, column_name, data_type
                        FROM information_schema.columns
                        WHERE column_name LIKE :term
                          AND table_schema = DATABASE()
                        LIMIT :limit
                    """
                else: # Generic/SQL Server
                   query = """
                        SELECT table_schema, table_name, column_name, data_type
                        FROM information_schema.columns
                        WHERE column_name LIKE :term
                        LIMIT :limit
                    """

                if self.engine:
                    async with self.engine.connect() as conn:
                         result_proxy = await conn.execute(text(query), {"term": term_pattern, "limit": current_limit})
                         rows = result_proxy.fetchall()
                         for row in rows:
                             results.append({
                                 "type": "column",
                                 "table": row[1],
                                 "schema": row[0],
                                 "name": row[2], 
                                 "description": f"Column: {row[2]} (Type: {row[3]}) in {row[0]}.{row[1]}",
                                 "metadata": f"Type: {row[3]}"
                             })

            # Cache the results ONLY if we found something
            # This prevents caching False Negatives (empty results) which might be due to transient issues or bad queries
            if self.cache and results:
                await self.cache.set(search_term, search_type, limit, results)
                
            return results

        except Exception as e:
            self.logger.error(f"Error in SQL-based search_schema: {e}")
            return []

    def _extract_queries(self, response_text: str) -> List[str]:
        """
        Extract SQL queries from LLM response text.

        Args:
            response_text: The full response text from the LLM

        Returns:
            List of extracted SQL queries
        """
        queries = []

        # Method 1: Extract from markdown code blocks
        sql_pattern = r'```sql\n(.*?)\n```'
        matches = re.findall(sql_pattern, response_text, re.DOTALL | re.IGNORECASE)

        for match in matches:
            cleaned_query = match.strip()
            if cleaned_query and not cleaned_query.lower().startswith('--'):
                queries.append(cleaned_query)

        # Method 2: If no markdown blocks, look for SQL-like patterns
        # CAUTION: This fallback generates false positives for explanations.
        # We will disable aggressive line scanning and only support markdown blocks or single-line exact queries.
        if not queries:
            cleaned_text = response_text.strip()
            # If the whole text looks like a query (starts with keyword, ends with ;)
            if re.match(r'^(SELECT|WITH|SHOW|DESCRIBE|EXPLAIN)\b.*?;$', cleaned_text, re.IGNORECASE | re.DOTALL):
                queries.append(cleaned_text)

        # Clean up queries
        cleaned_queries = []
        for query in queries:
            # Remove common prefixes/suffixes
            query = re.sub(r'^```sql\s*', '', query, flags=re.IGNORECASE)
            query = re.sub(r'\s*```$', '', query)
            query = query.strip()

            # Basic validation - should contain SELECT, WITH, etc.
            if re.search(r'\b(SELECT|WITH|SHOW|DESCRIBE|EXPLAIN)\b', query, re.IGNORECASE):
                cleaned_queries.append(query)

        return cleaned_queries

    def _extract_tables_from_query(self, query: str) -> List[str]:
        """Extract table names from SQL query."""
        pattern = r'(?:FROM|JOIN)\s+(?:[\w\.]*\.)?(\w+)'
        matches = re.findall(pattern, query.upper())
        return list(set(matches))

    async def cleanup(self) -> None:
        """Cleanup resources."""
        if self.engine:
            await self.engine.dispose()
        await super().cleanup()




# Factory function for creating enhanced SQL agents
def create_sql_agent(
    database_flavor: str,
    credentials: Union[str, Dict[str, Any]],
    schema_name: str = None,
    **kwargs
) -> SQLAgent:
    """
    Factory function to create SQL database agents.

    Args:
        database_flavor: Database type ('postgresql', 'mysql', 'sqlserver')
        credentials: Connection credentials (string or dict)
        schema_name: Target schema name
        **kwargs: Additional arguments

    Returns:
        Configured SQLAgent instance
    """
    # Set default schema names
    if schema_name is None:
        if database_flavor.lower() in ['postgresql', 'postgres']:
            schema_name = 'public'
        elif database_flavor.lower() == 'mysql':
            schema_name = 'mysql'
        elif database_flavor.lower() in ['sqlserver', 'mssql']:
            schema_name = 'dbo'
        else:
            schema_name = 'public'

    return SQLAgent(
        database_flavor=database_flavor,
        credentials=credentials,
        schema_name=schema_name,
        **kwargs
    )


# Example usage
"""
# Dictionary credentials example
pg_creds = {
    'host': 'localhost',
    'port': 5432,
    'database': 'sales_db',
    'username': 'user',
    'password': 'password'
}

pg_agent = create_sql_agent(
    database_flavor='postgresql',
    credentials=pg_creds,
    schema_name='public'
)

# Connection string example
mysql_agent = create_sql_agent(
    database_flavor='mysql',
    credentials='mysql://user:pass@localhost/dbname'
)

# Usage
await pg_agent.initialize_schema()

# Generate and execute query
query_result = await pg_agent.generate_query(
    "Show me all customers from the East region with their order totals"
)

execution_result = await pg_agent.execute_query(query_result['query'])
print(f"Query: {execution_result['query']}")
print(f"Data: {execution_result['data']}")
"""

