"""
SQL Execution Service.

Handles:
- Connection and dialect resolution from semantic model specs
- Table name resolution (dataset -> physical SQL table)
- SQL query execution against DSS connections

Used by:
- Playground routes (webapp)
- Agent Tools
"""

import logging
from typing import Any, Optional, Set

from semantic_models.schemas import EntityClass, SemanticModelSpec

from semantic_models.exceptions import ExecutionError, ValidationError
from semantic_models.logging_utils import summarize_for_logging, truncate_text
from semantic_models.services.dss_utils import get_connection_dialect, get_dataset_location_info
from semantic_models.utils import add_limit_to_query

logger = logging.getLogger(__name__)
SQL_LOG_CHARS = 3000

# =============================================================================
# Connection Resolution
# =============================================================================


def _split_dataset_ref(dataset_ref: Optional[str]) -> tuple[Optional[str], Optional[str]]:
    if not dataset_ref or "." not in dataset_ref:
        return None, None
    project_key, dataset_name = dataset_ref.split(".", 1)
    return project_key, dataset_name


def get_entity_connection(entity: EntityClass) -> Optional[str]:
    """Resolve the DSS connection name for an entity."""
    if entity.type != "DATASET":
        raise ValidationError(f"Unsupported entity type for SQL execution: {entity.type}")

    project_key, dataset_name = _split_dataset_ref(entity.datasetRef)
    if not project_key or not dataset_name:
        logger.warning("Dataset entity missing datasetRef: %s", entity.name)
        return None

    try:
        full_dataset_name = f"{project_key}.{dataset_name}"
        logger.debug("Getting connection for dataset: %s", full_dataset_name)
        location = get_dataset_location_info(project_key, dataset_name)
        return location.connection
    except ValidationError:
        raise
    except Exception as e:
        logger.warning("Failed to get connection for dataset %s: %s", dataset_name, e)
        return None


def ensure_single_connection(spec: SemanticModelSpec) -> str:
    """
    Ensure all entities in the semantic model use the same DSS connection.

    Returns the connection name.
    Raises ValidationError if entities use different connections.
    """
    if not spec.entities:
        raise ValidationError("Semantic model has no entities")

    connections: Set[str] = set()
    entity_connections: dict[str, Optional[str]] = {}

    for entity in spec.entities:
        conn = get_entity_connection(entity)
        entity_connections[entity.name] = conn
        if conn:
            connections.add(conn)

    # Check for entities without connections
    missing = [eid for eid, conn in entity_connections.items() if not conn]
    if missing:
        raise ValidationError(f"Entities without connection: {missing}")

    if len(connections) != 1:
        raise ValidationError(f"All entities must use the same connection. Found: {sorted(connections)}")

    connection_name = next(iter(connections))
    logger.debug("Using single connection '%s' for semantic model", connection_name)
    return connection_name


def get_dialect_from_spec(spec: SemanticModelSpec) -> str:
    """
    Infer the SQL dialect from the semantic model's underlying connection.

    Returns a SQLGlot-compatible dialect name, or "generic" if unknown.
    """
    try:
        connection_name = ensure_single_connection(spec)
    except ValidationError as e:
        logger.warning("Cannot determine dialect: %s", e)
        return "generic"

    dialect = get_connection_dialect(connection_name)
    logger.debug("Connection '%s' -> dialect='%s'", connection_name, dialect)
    return dialect


# =============================================================================
# Table Name Resolution
# =============================================================================


def resolve_dataset_sql_table(project_key: str, dataset_name: str, dialect: str = "generic") -> str:
    try:
        # Get dataset location info
        full_dataset_name = f"{project_key}.{dataset_name}"
        logger.info("Resolving SQL table for dataset: %s", full_dataset_name)
        location = get_dataset_location_info(project_key, dataset_name)
        logger.info("Resolved %s -> %s", full_dataset_name, location.full_table)
        return location.full_table

    except ValidationError:
        raise
    except Exception as e:
        logger.warning("Failed to resolve SQL table for %s.%s: %s", project_key, dataset_name, e)
        raise ValidationError(f"Failed to resolve SQL table for {project_key}.{dataset_name}: {e}")


def resolve_table_name(entity: EntityClass, dialect: str = "generic") -> str:
    """Resolve the actual SQL table name for an entity.
    """
    if entity.type != "DATASET":
        raise ValueError(f"Unsupported entity type for SQL execution: {entity.type}")

    project_key, dataset_name = _split_dataset_ref(entity.datasetRef)
    if not project_key or not dataset_name:
        raise ValueError(
            f"Dataset entity '{entity.name}' missing required datasetRef. "
            f"Got datasetRef={entity.datasetRef!r}"
        )

    return resolve_dataset_sql_table(project_key, dataset_name, dialect)


# =============================================================================
# SQL Execution
# =============================================================================


def execute_sql(
    spec: SemanticModelSpec, sql: str, max_rows: int = 1000
) -> tuple[list[str], list[list[Any]], int, bool]:
    """
    Execute SQL against the semantic model's single connection.

    Args:
        spec: The semantic model specification (used to determine connection)
        sql: The SQL query to execute
        max_rows: Maximum number of rows to return (default 1000)

    Returns:
        Tuple of (columns, rows, row_count, truncated)
    """
    # Get the validated single connection and dialect
    connection_name = ensure_single_connection(spec)
    dialect = get_dialect_from_spec(spec)

    # Enforce read-only/safety checks at execution boundary.
    # This remains mandatory even if the agent already called validate_sql.
    try:
        from semantic_models.services.sql_validation import validate_sql as validate_sql_service

        validate_sql_service(sql=sql, dialect=dialect, spec=spec)
    except ValidationError as e:
        logger.info("=== SQL EXECUTION BLOCKED ===")
        logger.info("  Validation error: %s", e)
        raise ExecutionError(f"SQL validation failed: {e}")

    limited_sql = add_limit_to_query(sql, max_rows + 1, dialect)

    logger.info("=== SQL EXECUTION START ===")
    logger.info("  Connection: %s", connection_name)
    logger.info("  Dialect: %s", dialect)
    logger.info("  Max rows: %d", max_rows)
    logger.info("  SQL: %s", truncate_text(limited_sql, max_chars=SQL_LOG_CHARS))

    try:
        import dataiku

        executor = dataiku.core.sql.SQLExecutor2(connection=connection_name)
        result = executor.query_to_df(limited_sql)

        # Convert to response format
        columns = list(result.columns)
        truncated = len(result) > max_rows
        if truncated:
            result = result.head(max_rows)

        # Convert DataFrame to list of lists, handling NaN/None
        rows = result.where(result.notnull(), None).values.tolist()

        logger.info("=== SQL EXECUTION SUCCESS ===")
        logger.info("  Columns: %s", summarize_for_logging(columns, max_chars=100, max_items=100))
        logger.info("  Rows returned: %d", len(rows))
        logger.info("  Truncated: %s", truncated)

        return columns, rows, len(rows), truncated

    except Exception as e:
        logger.info("=== SQL EXECUTION FAILED ===")
        logger.info("  Error: %s", truncate_text(e, max_chars=1000))
        logger.exception("SQL execution failed")
        raise ExecutionError(f"SQL execution failed: {e}")
