"""
Streamlined KQL Query Execution Module

This module provides simplified KQL query execution with Azure authentication
and integrated schema management using the centralized SchemaManager.

Author: Arjun Trivedi
Email: arjuntrivedi42@yahoo.com
"""

import asyncio
import logging
import re
import time
import json
from typing import Any, Dict, List, Optional, Tuple

import pandas as pd
import numpy as np
from azure.kusto.data import KustoClient, KustoConnectionStringBuilder
from azure.kusto.data.exceptions import KustoServiceError

# Constants import
from .constants import DEFAULT_QUERY_TIMEOUT
from .memory import get_memory_manager
from .utils import (
    extract_cluster_and_database_from_query,
    extract_tables_from_query,
    generate_query_description,
    retry_on_exception
)

logger = logging.getLogger(__name__)



# Import schema validator at module level - now from memory.py
_schema_validator = None

def get_schema_validator():
    """Lazy load schema validator to avoid circular imports."""
    global _schema_validator  # pylint: disable=global-statement
    if _schema_validator is None:
        try:
            memory = get_memory_manager()
            # Schema validator is now part of MemoryManager
            _schema_validator = memory
        except (ImportError, AttributeError, ValueError) as e:
            logger.warning("Schema validator not available: %s", e)
            _schema_validator = None
    return _schema_validator













def validate_query(query: str) -> Tuple[str, str]:
    """
    Validate a KQL query and extract cluster and database information.

    Args:
        query: The KQL query to validate

    Returns:
        Tuple of (cluster_uri, database)

    Raises:
        ValueError: If query is invalid or missing required components
    """
    if not query or not query.strip():
        raise ValueError("Query cannot be empty")

    try:
        cluster, database = extract_cluster_and_database_from_query(query)

        if not cluster:
            raise ValueError("Query must include cluster specification")

        if not database:
            # According to test expectations, missing database should be treated as invalid cluster format
            raise ValueError("Query must include cluster specification - invalid cluster format without database")

        return cluster, database

    except Exception as e:
        if "cluster" in str(e).lower() or "database" in str(e).lower():
            raise
        raise ValueError(f"Invalid query format: {e}") from e


def _normalize_cluster_uri(cluster_uri: str) -> str:
    """Normalize cluster URI for connection."""
    if not cluster_uri:
        raise ValueError("Cluster URI cannot be None or empty")

    if not cluster_uri.startswith("https://"):
        cluster_uri = f"https://{cluster_uri}"
    return cluster_uri.rstrip("/")


# Client cache for pooling
_client_cache = {}

def _get_kusto_client(cluster_url: str) -> KustoClient:
    """Create and authenticate a Kusto client with pooling."""
    # Normalize URL for consistent caching
    normalized_url = _normalize_cluster_uri(cluster_url)
    
    if normalized_url not in _client_cache:
        logger.info("Creating new Kusto client for %s", normalized_url)
        kcsb = KustoConnectionStringBuilder.with_az_cli_authentication(normalized_url)
        _client_cache[normalized_url] = KustoClient(kcsb)
    
    return _client_cache[normalized_url]

def _parse_kusto_response(response) -> pd.DataFrame:
    """Parse a Kusto response into a pandas DataFrame."""
    if not response or not getattr(response, "primary_results", None):
        return pd.DataFrame()

    first_result = response.primary_results[0]
    df = None

    try:
        td = first_result.to_dict()
        if isinstance(td, dict) and "data" in td and td["data"] is not None:
            df = pd.DataFrame(td["data"])
    except (ValueError, TypeError, KeyError, AttributeError):
        df = None

    if df is None:
        try:
            rows = list(first_result)
            cols = [c.column_name for c in getattr(first_result, "columns", []) if hasattr(c, "column_name")]
            if rows and isinstance(rows[0], (list, tuple)) and cols:
                df = pd.DataFrame(rows, columns=cols)
            else:
                df = pd.DataFrame(rows)
        except (ValueError, TypeError, KeyError, AttributeError):
            df = pd.DataFrame()

    return df

@retry_on_exception()
def _execute_kusto_query_sync(kql_query: str, cluster: str, database: str, _timeout: int = DEFAULT_QUERY_TIMEOUT) -> pd.DataFrame:
    """
    Core synchronous function to execute a KQL query against a Kusto cluster.
    Adds configurable request timeout and uses retry decorator for transient failures.
    """
    cluster_url = _normalize_cluster_uri(cluster)
    logger.info("Executing KQL on %s/%s: %s...", cluster_url, database, kql_query[:150])

    client = _get_kusto_client(cluster_url)
    start_time = time.time()
    try:
        is_mgmt_query = kql_query.strip().startswith('.')

        # First execution attempt
        try:
            if is_mgmt_query:
                response = client.execute_mgmt(database, kql_query)
            else:
                response = client.execute(database, kql_query)

            execution_time_ms = (time.time() - start_time) * 1000
            df = _parse_kusto_response(response)
            logger.debug("Query returned %d rows in %.2fms.", len(df), execution_time_ms)

            try:
                loop = asyncio.get_running_loop()
                loop.create_task(_post_execution_learning_bg(kql_query, cluster, database, df, execution_time_ms))
            except RuntimeError:
                logger.debug("No event loop running - skipping background learning task")

            return df

        except KustoServiceError as e:
            error_str = str(e)
            # Check for SEM0100 (Missing Column) or "Failed to resolve" errors
            if 'sem0100' in error_str.lower() or "failed to resolve scalar expression" in error_str.lower():
                logger.info("SEM0100 error detected: %s", error_str[:100])

                # Trigger schema refresh for involved tables to prevent future errors
                try:
                    tables = extract_tables_from_query(kql_query)
                    if tables:
                        loop = asyncio.get_running_loop()
                        # Force refresh schema for these tables
                        loop.create_task(_ensure_schema_discovered(cluster, database, tables))
                except (RuntimeError, Exception) as refresh_error:
                    logger.debug("Failed to trigger schema refresh on error: %s", refresh_error)

            # Re-raise the original error
            raise

    finally:
        # Do not close the client as it is now pooled/cached
        pass


def execute_large_query(query: str, cluster: str, database: str, _chunk_size: int = 1000, timeout: int = DEFAULT_QUERY_TIMEOUT) -> pd.DataFrame:
    """
    Minimal query chunking helper.
    - If the query already contains explicit 'take' or 'limit', execute as-is.
    - Otherwise run a single timed execution (safe fallback).
    This conservative approach avoids aggressive query rewriting while enabling
    an explicit place to improve chunking later.
    """
    if ' take ' in (query or "").lower() or ' limit ' in (query or "").lower():
        return _execute_kusto_query_sync(query, cluster, database, timeout)
    # Fallback: single execution with configured timeout & retries
    return _execute_kusto_query_sync(query, cluster, database, timeout)




# Essential functions for compatibility
def validate_kql_query_advanced(query: str, cluster: Optional[str] = None, database: Optional[str] = None) -> Dict[str, Any]:
    """
    Simplified KQL query validation.
    """
    try:
        if not query or not query.strip():
            return {
                "valid": False,
                "error": "Query cannot be empty",
                "suggestions": []
            }

        # Basic KQL syntax validation
        query_lower = query.lower().strip()

        # Check for management commands
        if query_lower.startswith('.') and not any(cmd in query_lower for cmd in ['.show', '.list', '.help']):
            return {
                "valid": False,
                "error": "Invalid management command",
                "suggestions": ["Use .show tables or .show databases for management commands"]
            }

        return {
            "valid": True,
            "cluster": cluster,
            "database": database,
            "suggestions": []
        }

    except (ValueError, TypeError, AttributeError) as e:
        return {
            "valid": False,
            "error": str(e),
            "suggestions": []
        }


def kql_execute_tool(kql_query: str, cluster_uri: Optional[str] = None, database: Optional[str] = None) -> pd.DataFrame:
    """
    Simplified KQL execution function.
    Executes the query directly without internal generation or complex pre-processing.
    """
    try:
        if not kql_query or not kql_query.strip():
            logger.error("Empty query provided to kql_execute_tool")
            raise ValueError("KQL query cannot be None or empty")

        clean_query = kql_query.strip()

        # Check if query already contains cluster/database specification
        has_cluster_spec = "cluster(" in clean_query and "database(" in clean_query

        if has_cluster_spec:
            # Query already has cluster/database - extract them and use the base query
            try:
                extracted_cluster, extracted_database = extract_cluster_and_database_from_query(clean_query)
                cluster = cluster_uri or extracted_cluster
                db = database or extracted_database
            except (ValueError, TypeError, AttributeError) as extract_error:
                logger.warning("Failed to extract cluster/database: %s", extract_error)
                cluster = cluster_uri
                db = database
        else:
            # No cluster specification in query - use parameters
            cluster = cluster_uri
            db = database

        if not cluster:
            raise ValueError("Cluster URI must be specified in query or parameters. Example: 'https://help.kusto.windows.net'")

        # Check if this is a management command that doesn't require a database
        is_mgmt_command = clean_query.strip().startswith('.')
        mgmt_commands_no_db = ['.show databases', '.show clusters', '.help']
        mgmt_needs_no_db = any(cmd in clean_query.lower() for cmd in mgmt_commands_no_db)

        if not db and not (is_mgmt_command and mgmt_needs_no_db):
            raise ValueError("Database must be specified in query or parameters. Example: 'Samples'")

        # Use "master" database for management commands that don't require specific database
        db_for_execution = db if db else "master"

        # Execute with enhanced error handling that propagates all exceptions
        try:
            return _execute_kusto_query_sync(clean_query, cluster, db_for_execution)
        except KustoServiceError as e:
            logger.error("Kusto service error during execution: %s", e)
            raise  # Re-raise to be handled by the MCP tool
        except (ValueError, TypeError, RuntimeError) as exec_error:
            logger.error("Generic query execution failed: %s", exec_error)
            raise  # Re-raise instead of returning empty DataFrame

    except (ValueError, TypeError, RuntimeError, AttributeError) as e:
        logger.error("kql_execute_tool failed pre-execution: %s", e)
        logger.error("Original query was: %s", kql_query if 'kql_query' in locals() else 'Unknown')
        raise  # Re-raise instead of returning empty DataFrame



async def _post_execution_learning_bg(query: str, cluster: str, database: str, _df: pd.DataFrame, execution_time_ms: float = 0.0):
    """
    Enhanced background learning task with automatic schema discovery triggering.
    This runs asynchronously to avoid blocking query response.
    """
    try:
        # Extract table names from the executed query using the enhanced parse_query_entities
        from .utils import parse_query_entities
        entities = parse_query_entities(query)
        tables = entities.get("tables", [])

        # If no tables extracted, try fallback extraction
        if not tables:
            try:
                # Fallback: extract table names using simpler pattern matching
                table_pattern = r'\b([A-Za-z_][A-Za-z0-9_]*)\s*\|'
                fallback_tables = re.findall(table_pattern, query)
                if fallback_tables:
                    tables = fallback_tables[:1]  # Take first table found
                    logger.debug("Fallback table extraction found: %s", tables)
                else:
                    # Even without table extraction, store successful query globally
                    logger.debug("No tables extracted but storing successful query globally")
                
                    memory_manager = get_memory_manager()
                    description = generate_query_description(query)
                    try:
                        # Store in global successful queries without table association
                        memory_manager.add_global_successful_query(cluster, database, query, description, execution_time_ms)
                        logger.debug("Stored global successful query: %s", description)
                    except (ValueError, TypeError, AttributeError) as e:
                        logger.debug("Failed to store global successful query: %s", e)
                    return
            except (ValueError, TypeError, AttributeError) as fallback_error:
                logger.debug("Fallback table extraction failed: %s", fallback_error)
                return

        # Store successful query for each table involved
    
        memory_manager = get_memory_manager()
        description = generate_query_description(query)

        for table in tables:
            try:
                # Add successful query to table-specific memory
                memory_manager.add_successful_query(cluster, database, query, description, execution_time_ms)
                logger.debug("Stored successful query for %s: %s", table, description)
            except (ValueError, TypeError, AttributeError) as e:
                logger.debug("Failed to store successful query for %s: %s", table, e)

        # ENHANCED: Force schema discovery for all tables involved in the query
        await _ensure_schema_discovered(cluster, database, tables)

    except (ValueError, TypeError, RuntimeError) as e:
        logger.debug("Background learning task failed: %s", e)


async def _ensure_schema_discovered(cluster_uri: str, database: str, tables: List[str]):
    """
    Force schema discovery if not in memory.
    This is the implementation recommended in the analysis.
    """

    from .utils import SchemaManager

    memory = get_memory_manager()
    schema_manager = SchemaManager(memory)

    for table in tables:
        try:
            # Get schema from database using internal method that returns list of schemas
            schemas = memory._get_database_schema(cluster_uri, database)

            # Find the schema for this specific table
            schema = next((s for s in schemas if s.get("table") == table), None)

            if not schema or not schema.get("columns"):
                # Trigger live discovery with force refresh
                logger.info("Auto-triggering schema discovery for %s.%s", database, table)
                discovered_schema = await schema_manager.get_table_schema(  # pylint: disable=too-many-function-args
                    cluster=cluster_uri, database=database, table=table, _force_refresh=True
                )

                if discovered_schema and discovered_schema.get("columns"):
                    logger.info("Successfully auto-discovered schema for %s with %d columns", table, len(discovered_schema['columns']))
                elif discovered_schema and discovered_schema.get("is_not_found"):
                    logger.info("Schema discovery skipped for '%s' (likely a local variable or function)", table)
                else:
                    logger.warning("Auto-discovery failed for %s - no columns found", table)
            else:
                logger.debug("Schema already exists for %s, skipping auto-discovery", table)

        except (ValueError, TypeError, KeyError, AttributeError) as e:
            logger.warning("Auto schema discovery failed for %s: %s", table, e)
            # Continue with other tables even if one fails

def get_knowledge_corpus():
    """Backward-compatible wrapper to memory.get_knowledge_corpus"""
    try:
        return get_memory_manager()
    except (ImportError, AttributeError, ValueError):
        # Fallback mock for tests if import fails
        class MockCorpus:
            def get_ai_context_for_query(self, _query: str) -> Dict[str, Any]:
                return {}

            @property
            def memory_manager(self):
                return None
        return MockCorpus()




async def execute_kql_query(
    query: str,
    cluster: Optional[str] = None,
    database: Optional[str] = None,
    visualize: bool = False,
    use_schema_context: bool = True,
    timeout: int = 300,
    use_cache: bool = True
) -> Any: # pylint: disable=too-many-arguments
    """
    Legacy compatibility function for __init__.py import.

    Returns a list of dictionaries (test compatibility) or dictionary with success/error status.
    Enhanced with background learning integration and caching.

    Args:
        query: KQL query to execute
        cluster: Cluster URI (optional)
        database: Database name (optional)
        visualize: Whether to include visualization (ignored for now)
        use_schema_context: Whether to use schema context (ignored for now)
        use_cache: Whether to use query result caching (default: True)
    """
    try:
        # Check cache first
        if use_cache:
            try:
                memory = get_memory_manager()
                cached_json = memory.get_cached_result(query)
                if cached_json:
                    logger.info("Returning cached result for query")
                    return json.loads(cached_json)
            except (ImportError, AttributeError, ValueError, json.JSONDecodeError) as e:
                logger.debug("Cache lookup failed: %s", e)

        # Optionally load schema context prior to execution (tests may patch get_knowledge_corpus)
        if use_schema_context:
            try:
                corpus = get_knowledge_corpus()
                # Call the method so tests can patch and assert it was invoked
                if hasattr(corpus, 'get_ai_context_for_query'):
                    _ = corpus.get_ai_context_for_query(query)  # type: ignore[attr-defined]
            except (ValueError, TypeError, AttributeError, KeyError):
                # Ignore failures to keep function resilient in test and runtime environments
                pass

        # Extract cluster and database if not provided
        if not cluster or not database:
            extracted_cluster, extracted_database = extract_cluster_and_database_from_query(query)
            cluster = cluster or extracted_cluster
            database = database or extracted_database

        if not cluster or not database:
            raise ValueError("Query must include cluster and database specification")

        # Execute using the core sync function wrapped in asyncio.to_thread and enforce overall timeout
        df = await asyncio.wait_for(
            asyncio.to_thread(_execute_kusto_query_sync, query, cluster, database, timeout),
            timeout=timeout + 5,
        )

        # Return list format for test compatibility with proper serialization
        if hasattr(df, 'to_dict'):
            # Convert DataFrame to serializable records
            records = []
            try:
                for _, row in df.iterrows():
                    record = {}
                    for col, value in row.items():
                        # Handle complex types (lists, dicts, arrays) first to avoid ambiguity
                        if isinstance(value, (list, dict, tuple, np.ndarray)):
                            if isinstance(value, np.ndarray):
                                record[col] = value.tolist()
                            else:
                                record[col] = value
                            continue

                        # Safe check for null/NaN for scalar types
                        is_null = False
                        try:
                            if pd.isna(value):
                                is_null = True
                        except (ValueError, TypeError):
                            # Fallback for ambiguous cases
                            is_null = False

                        if is_null:
                            record[col] = None
                        elif hasattr(value, 'isoformat'):  # Timestamp objects
                            record[col] = value.isoformat()
                        elif hasattr(value, 'strftime'):  # datetime objects
                            record[col] = value.strftime('%Y-%m-%d %H:%M:%S')
                        elif isinstance(value, type):  # type objects
                            record[col] = value.__name__
                        elif hasattr(value, 'item'):  # numpy types
                            record[col] = value.item()
                        else:
                            record[col] = value
                    records.append(record)
            except (ValueError, TypeError, KeyError, AttributeError) as e:
                logger.warning("DataFrame serialization failed: %s", e)
                # Fallback to string conversion
                records = df.astype(str).to_dict("records")

            if visualize and records:
                # Add simple visualization marker for tests
                records.append({"visualization": "chart_data"})

            # Cache the result
            if use_cache and records:
                try:

                    memory = get_memory_manager()
                    memory.cache_query_result(query, json.dumps(records), len(records))
                except Exception as e:
                    logger.debug("Failed to cache result: %s", e)

            return records
        else:
            return []

    except Exception as e:
        logger.error("Query execution failed: %s", e)
        raise  # Re-raise for test compatibility


async def execute_with_full_flow(query: str, user_context: Optional[str] = None) -> Dict:
    """
    Implement complete execution flow with learning as recommended in the analysis.
    This implements the expected flow: execute â†’ learn â†’ discover â†’ refine.
    """
    try:
        # Step 1: Execute initial query
        result = await execute_kql_query(query)

        # Step 2: Extract and learn from context
        if user_context:
            context = await extract_context_from_prompt(user_context)
            await learn_from_data(result, context)

        # Step 3: Trigger background schema discovery
        from .utils import parse_query_entities
        entities = parse_query_entities(query)
        cluster, database, tables = entities["cluster"], entities["database"], entities["tables"]

        if cluster and database and tables:
            asyncio.create_task(_ensure_schema_discovered(cluster, database, tables))

        # Step 4: Generate enhanced query if needed (simplified for now)
        enhanced_result = {
            "initial_result": result,
            "learning_complete": True,
            "schema_discovery_triggered": bool(tables),
            "entities_extracted": entities
        }

        return enhanced_result

    except (ValueError, TypeError, RuntimeError) as e:
        logger.error("Full flow execution failed: %s", e)
        return {
            "initial_result": None,
            "error": str(e),
            "learning_complete": False
        }


async def extract_context_from_prompt(user_context: str) -> Dict:
    """Extract meaningful context from user input for learning."""
    return {
        "user_intent": user_context,
        "needs_refinement": len(user_context.split()) > 10,  # Simple heuristic
        "context_type": "natural_language"
    }


async def learn_from_data(result_data: Any, context: Dict):
    """Store learning results in memory for future use."""
    try:

        memory = get_memory_manager()

        # Convert result to learnable format
        if isinstance(result_data, list) and result_data:
            learning_data = {
                "row_count": len(result_data),
                "columns": list(result_data[0].keys()) if result_data else [],
                "success": True,
                "context": context
            }

            # Store learning result using the context info
            memory.store_learning_result(
                query=context.get("user_intent", ""),
                result_data=learning_data,
                execution_type="enhanced_flow_execution"
            )

    except (ValueError, TypeError, AttributeError) as e:
        logger.warning("Learning from data failed: %s", e)
