from __future__ import annotations

import json
import logging
import re
from typing import Any, Dict, List, Optional

from actor.connectors.db import DBClient
from actor.runtime import RuntimeConfig
from actor.authentication.secrets import build_secret_provider, CompositeSecretProvider
from actor.authentication.database import build_database_connection
from actor.authentication.config import SecretStoreConfig
from actor.authentication.exceptions import ConfigurationError, SecretProviderError, AuthMaterializationError

logger = logging.getLogger(__name__)


class DBTaskExecutor:
    """Handles database-related tasks from queue, with optional LLM desensitization."""

    def __init__(self, runtime_config: RuntimeConfig, observability, llm_client=None) -> None:
        self.runtime_config = runtime_config
        self.observability = observability
        self.llm_client = llm_client
        self._db_client: DBClient | None = None

    def _ensure_client(self) -> DBClient:
        """Create or return the cached DBClient instance."""
        if self._db_client:
            return self._db_client

        db_cfg = self.runtime_config.get_database()
        if not db_cfg:
            raise RuntimeError("No database config available from SaaS")

        # Extract database_type from config if available
        database_type: Optional[str] = getattr(db_cfg, "database_type", None)

        # If URL already present, use it directly
        if db_cfg.url:
            client = DBClient(
                url=db_cfg.url,
                connect_args=db_cfg.connect_args,
                observability=self.observability,
                database_type=database_type,
            )
            client.validate()  # SELECT 1 — fail fast if unreachable
            self._db_client = client
            return self._db_client

        # Otherwise, attempt credential materialization
        if db_cfg.auth_type and db_cfg.secret_store:
            try:
                secret_store = SecretStoreConfig.parse_obj(db_cfg.secret_store)
                base_provider = build_secret_provider(secret_store)
                literal_values: Dict[str, Any] = db_cfg.literals or {}
                mapping: Dict[str, str] = {}

                if db_cfg.secret_mapping:
                    mapping.update(db_cfg.secret_mapping)
                for key, value in literal_values.items():
                    if key not in mapping and value is not None:
                        mapping[key] = key
                provider = CompositeSecretProvider(base_provider, literal_values)

                # Pass database_type for validation if available
                materialized = build_database_connection(
                    auth_type=db_cfg.auth_type,
                    mapping=mapping,
                    provider=provider,
                    database_type=database_type,
                )

                if isinstance(materialized, dict):
                    # Check for special Databricks connection marker
                    if "_databricks_connection" in materialized:
                        logger.warning("Databricks oauth_m2m requires special handling")
                        raise RuntimeError("Databricks oauth_m2m connection not yet supported in agent")

                    url = materialized.get("url")
                    connect_args = materialized.get("connect_args") or {}
                else:
                    url = materialized
                    connect_args = {}

                client = DBClient(
                    url=url,
                    connect_args=connect_args,
                    observability=self.observability,
                    database_type=database_type,
                )
                client.validate()  # SELECT 1 — fail fast if unreachable
                self._db_client = client
                return self._db_client

            except (ConfigurationError, SecretProviderError, AuthMaterializationError) as exc:
                raise RuntimeError(f"Failed to materialize database connection: {exc}") from exc

        raise RuntimeError("Database URL missing; cannot build connection")

    def handle(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """Handle a database task and return the result."""
        task_type = task.get("task_type")
        payload = task.get("payload") or {}
        client = self._ensure_client()

        if task_type == "db_validate":
            success, _, error = client.safe_connect()
            if success:
                schemas = client.list_schemas()
                tables: Dict[str, Any] = {}
                columns: Dict[str, Any] = {}
                for schema in schemas:
                    try:
                        schema_tables = client.list_tables(schema=schema)
                    except Exception:
                        schema_tables = []
                    tables[schema] = schema_tables
                    for table in schema_tables:
                        key = f"{schema}.{table}"
                        try:
                            cols = client.describe_table(table=table, schema=schema)
                            columns[key] = client.normalize_columns(cols)
                        except Exception:
                            columns[key] = []
                result = {
                    "status": "ok",
                    "database_type": client.database_type,
                    "schemas": schemas,
                    "tables": tables,
                    "columns": columns,
                }
                try:
                    total_tables = sum(len(v) for v in tables.values())
                    total_columns = sum(len(v) for v in columns.values())
                    logger.info(
                        "db_validate result: schemas=%s tables=%s columns=%s",
                        len(schemas),
                        total_tables,
                        total_columns,
                    )
                except Exception:
                    logger.info("db_validate result prepared")
                return result
            message = None
            if isinstance(error, dict):
                message = (error.get("error") or {}).get("message")
            raise RuntimeError(message or "database validation failed")

        if task_type == "db_schemas":
            return {"schemas": client.list_schemas()}

        if task_type == "db_tables":
            schema = payload.get("schema")
            return {"tables": client.list_tables(schema=schema)}

        if task_type == "db_columns":
            schema = payload.get("schema")
            table = payload.get("table")
            if not table:
                raise ValueError("table is required")
            return {"columns": client.describe_table(table=table, schema=schema)}

        if task_type == "db_sample":
            table = payload.get("table")
            schema = payload.get("schema")
            limit = payload.get("limit", 5)
            if not table:
                raise ValueError("table is required")
            rows = client.sample_table(table=table, schema=schema, limit=limit)
            return {"rows": rows}

        if task_type == "db_query":
            return self._handle_db_query(client, payload)

        if task_type == "classify_fields":
            return self._handle_classify_fields(client, payload)

        if task_type == "entity_resolve":
            return self._handle_entity_resolve(client, payload)

        raise ValueError(f"Unsupported db task_type '{task_type}'")

    def _handle_db_query(self, client: DBClient, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Handle db_query with desensitization pipeline."""
        query = payload.get("query")
        params = payload.get("params") or {}
        if not query:
            raise ValueError("query is required")

        success, rows, error = client.safe_query(query, params=params)
        if not success:
            return {"status": "error", "error": error}

        # If LLM client available, desensitize
        if self.llm_client:
            from actor.connectors.desensitize import desensitize

            table_context = payload.get("table_context") or {}
            entity_ids = payload.get("entity_ids") or []
            llm_cfg = self.runtime_config.get_llm()
            token_limit = llm_cfg.token_limit if llm_cfg else 4096
            model = self.llm_client.model

            result_str = desensitize(
                data=rows,
                table_context=table_context,
                entity_ids=entity_ids,
                llm_client=self.llm_client,
                model=model,
                chunk_max_tokens=token_limit,
            )
            return {"desensitized_text": result_str, "rows_processed": len(rows)}

        # No LLM — check passthrough
        if self.runtime_config.llm_passthrough:
            return {"rows": rows, "desensitization_status": "passthrough"}

        raise RuntimeError("LLM not configured and llm_passthrough is disabled; cannot return raw data")

    def _handle_classify_fields(self, client: DBClient, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Classify field sensitivity using LLM."""
        if not self.llm_client:
            raise RuntimeError("LLM client required for classify_fields task")

        from actor.connectors.desensitize import _classify_fields

        table = payload.get("table")
        schema = payload.get("schema")
        table_context = payload.get("table_context") or {}
        entity_ids = payload.get("entity_ids") or []

        # Get sample data for classification
        if table:
            sample_data = client.sample_table(table=table, schema=schema, limit=5)
        else:
            sample_data = payload.get("sample_data") or []

        classification = _classify_fields(
            data=sample_data,
            table_context=table_context,
            entity_ids=entity_ids,
            llm_client=self.llm_client,
            model=self.llm_client.model,
        )
        return {"classification": classification}

    def _handle_entity_resolve(self, client: DBClient, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Resolve entity references locally using direct DB access.

        Runs all column probing against the local database connection,
        avoiding hundreds of round trips through Mars.

        Payload:
            entity_type: str - entity type hint (e.g. "Patient")
            entity_id: str - the identifier to find
            target_tables: list[str] - fully qualified table names to search
            dialect: str - SQL dialect name (postgresql, mysql, etc.)

        Returns:
            dict with matches list and best_match
        """
        entity_type = payload.get("entity_type", "")
        entity_id = payload.get("entity_id")
        target_tables = payload.get("target_tables", [])
        dialect = payload.get("dialect") or client.database_type or "postgresql"

        if not entity_id:
            raise ValueError("entity_id is required")
        if not target_tables:
            raise ValueError("target_tables is required")

        resolver = _EntityProber(client, dialect)
        result = resolver.resolve(entity_type, entity_id, target_tables)
        logger.info(
            "entity_resolve: type=%s id=%s tables=%d matches=%d best=%s",
            entity_type, entity_id, len(target_tables),
            len(result.get("matches", [])),
            result.get("best_match"),
        )
        return result

    def validate(self) -> None:
        """Validate database connectivity (raises on failure)."""
        client = self._ensure_client()
        client.validate()

    def dispose(self) -> None:
        """Clean up database resources."""
        if self._db_client:
            self._db_client.dispose()
            self._db_client = None


# ---------------------------------------------------------------------------
# Entity resolution helpers (runs entirely on agent, no Mars round trips)
# ---------------------------------------------------------------------------

# Exact data_type values that are strictly numeric.
_NUMERIC_TYPE_EXACT = frozenset({
    'int', 'int2', 'int4', 'int8', 'int16', 'int32', 'int64',
    'integer', 'smallint', 'bigint', 'tinyint', 'mediumint',
    'serial', 'bigserial', 'smallserial',
    'numeric', 'decimal',
    'float', 'float4', 'float8', 'float16', 'float32', 'float64',
    'double', 'double precision', 'real',
    'number', 'money', 'smallmoney',
})

# Type tokens that indicate a column cannot hold an entity reference.
_SKIP_TYPE_TOKENS = frozenset({
    'timestamp', 'datetime', 'date', 'time',
    'bool',
    'bytea', 'binary', 'blob', 'varbinary', 'image',
    'geometry', 'geography', 'polygon', 'linestring',
    'interval',
})


class _DialectSQL:
    """Minimal dialect-aware SQL generator for entity probing."""

    _DIALECTS = {
        'postgresql', 'redshift', 'mysql', 'mariadb', 'snowflake', 'bigquery',
        'clickhouse', 'sqlserver', 'oracle', 'trino', 'sparksql', 'databricks',
        'athena', 'vertica', 'teradata', 'db2', 'duckdb',
    }

    def __init__(self, dialect: str):
        self.dialect = dialect.lower()
        if self.dialect not in self._DIALECTS:
            self.dialect = 'postgresql'

    def _quote(self, name: str) -> str:
        if self.dialect in ('mysql', 'mariadb', 'sparksql', 'databricks'):
            return f'`{name}`'
        if self.dialect == 'sqlserver':
            return f'[{name}]'
        return f'"{name}"'

    def _cast_text(self, column: str) -> str:
        col = self._quote(column)
        d = self.dialect
        if d in ('postgresql', 'redshift', 'duckdb'):
            return f'{col}::text'
        if d in ('mysql', 'mariadb'):
            return f'CAST({col} AS CHAR)'
        if d == 'snowflake':
            return f'{col}::STRING'
        if d in ('bigquery', 'sparksql', 'databricks', 'athena'):
            return f'CAST({col} AS STRING)'
        if d == 'clickhouse':
            return f'toString({col})'
        if d == 'sqlserver':
            return f'CAST({col} AS NVARCHAR(MAX))'
        if d == 'oracle':
            return f'TO_CHAR({col})'
        if d in ('vertica', 'teradata', 'db2'):
            return f'CAST({col} AS VARCHAR(4000))'
        if d == 'trino':
            return f'CAST({col} AS VARCHAR)'
        return f'CAST({col} AS TEXT)'

    def contains(self, column: str, value: str) -> str:
        safe = value.replace("'", "''")
        return f"{self._cast_text(column)} LIKE '%{safe}%'"

    def exact(self, column: str, value: str) -> str:
        safe = value.replace("'", "''")
        return f"{self._quote(column)} = '{safe}'"

    def cast_match(self, column: str, value: str) -> str:
        safe = value.replace("'", "''")
        return f"{self._cast_text(column)} = '{safe}'"

    def json_extract_text(self, column: str, path: str) -> str:
        col = self._quote(column)
        d = self.dialect
        if d in ('postgresql', 'duckdb'):
            return f"{col}->>'{path}'"
        if d == 'redshift':
            return f"JSON_EXTRACT_PATH_TEXT({col}, '{path}')"
        if d in ('mysql', 'mariadb'):
            return f"JSON_UNQUOTE(JSON_EXTRACT({col}, '$.{path}'))"
        if d == 'snowflake':
            return f"{col}:{path}::STRING"
        if d in ('bigquery', 'sqlserver', 'oracle', 'teradata', 'db2'):
            return f"JSON_VALUE({col}, '$.{path}')"
        if d == 'clickhouse':
            return f"JSONExtractString({col}, '{path}')"
        if d in ('trino', 'athena'):
            return f"json_extract_scalar({col}, '$.{path}')"
        if d in ('sparksql', 'databricks'):
            return f"get_json_object({col}, '$.{path}')"
        if d == 'vertica':
            return f"MAPLOOKUP({col}, '{path}')"
        return f"JSON_VALUE({col}, '$.{path}')"

    def json_match(self, column: str, path: str, value: str) -> str:
        safe = value.replace("'", "''")
        return f"{self.json_extract_text(column, path)} = '{safe}'"

    def json_keys_query(self, table: str, column: str, limit: int = 50) -> Optional[str]:
        col = self._quote(column)
        d = self.dialect
        if d in ('postgresql', 'duckdb'):
            return f"SELECT DISTINCT jsonb_object_keys({col}::jsonb) as key FROM {table} WHERE {col} IS NOT NULL LIMIT {limit}"
        if d == 'snowflake':
            return f"SELECT DISTINCT f.key FROM {table}, LATERAL FLATTEN(input => {col}) f WHERE {col} IS NOT NULL LIMIT {limit}"
        if d == 'bigquery':
            return f"SELECT DISTINCT key FROM {table}, UNNEST(JSON_KEYS({col})) AS key WHERE {col} IS NOT NULL LIMIT {limit}"
        if d in ('mysql', 'mariadb'):
            return f"SELECT DISTINCT jt.key_name as key FROM {table}, JSON_TABLE({col}, '$.*' COLUMNS(key_name VARCHAR(255) PATH '$')) jt LIMIT {limit}"
        return None

    def schema_columns_query(self, table: str) -> Optional[str]:
        parts = re.sub(r'["\[\]`]', '', table).split('.')
        if len(parts) == 3:
            catalog, schema, table_name = parts
        elif len(parts) == 2:
            schema, table_name = parts
            catalog = None
        else:
            table_name = parts[0]
            schema = 'public'
            catalog = None
        d = self.dialect
        if d in ('postgresql', 'redshift', 'duckdb'):
            return f"SELECT column_name, data_type FROM information_schema.columns WHERE table_schema = '{schema}' AND table_name = '{table_name}'"
        if d in ('mysql', 'mariadb'):
            return f"SELECT COLUMN_NAME as column_name, DATA_TYPE as data_type FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = '{schema}' AND TABLE_NAME = '{table_name}'"
        if d == 'snowflake':
            return f"SELECT COLUMN_NAME as column_name, DATA_TYPE as data_type FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = '{schema.upper()}' AND TABLE_NAME = '{table_name.upper()}'"
        if d == 'bigquery':
            proj = catalog or 'project'
            return f"SELECT column_name, data_type FROM `{proj}.{schema}.INFORMATION_SCHEMA.COLUMNS` WHERE table_name = '{table_name}'"
        if d == 'sqlserver':
            return f"SELECT COLUMN_NAME as column_name, DATA_TYPE as data_type FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = '{schema}' AND TABLE_NAME = '{table_name}'"
        return f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name}'"


class _EntityProber:
    """Discovers entity references by probing columns locally.

    Mirrors the probing logic from gnosis EntityResolver but runs
    entirely on the agent side against the local DBClient, eliminating
    the per-column round trips through Mars.
    """

    def __init__(self, client: DBClient, dialect: str):
        self.client = client
        self.sql = _DialectSQL(dialect)

    def resolve(
        self, entity_type: str, entity_id: str, target_tables: List[str]
    ) -> Dict[str, Any]:
        id_variants = self._generate_id_variants(entity_type, entity_id)
        matches: List[Dict[str, Any]] = []

        for table in target_tables:
            columns = self._get_columns(table)
            table_matches = self._discover_in_table(
                table, columns, entity_id, id_variants
            )
            matches.extend(table_matches)

        best_match = None
        if matches:
            best_match = max(
                matches, key=lambda m: (m["sample_hit_count"], m["confidence"])
            )

        return {
            "entity_type": entity_type,
            "entity_id": entity_id,
            "matches": matches,
            "best_match": best_match,
            "resolution_metadata": {
                "tables_searched": len(target_tables),
                "matches_found": len(matches),
                "id_variants_tested": len(id_variants),
                "dialect": self.sql.dialect,
            },
        }

    # ------------------------------------------------------------------
    # ID variant generation
    # ------------------------------------------------------------------

    @staticmethod
    def _generate_id_variants(entity_type: str, entity_id: str) -> List[str]:
        variants = [entity_id]
        if entity_type:
            variants.append(f"{entity_type}/{entity_id}")
            variants.append(f"{entity_type.lower()}/{entity_id}")
            variants.append(f"{entity_type.upper()}/{entity_id}")
        if '-' in entity_id and len(entity_id.replace('-', '')) == 32:
            no_hyphens = entity_id.replace('-', '')
            variants.append(no_hyphens)
            if entity_type:
                variants.append(f"{entity_type}/{no_hyphens}")
        seen: set = set()
        unique: List[str] = []
        for v in variants:
            if v not in seen:
                seen.add(v)
                unique.append(v)
        return unique

    # ------------------------------------------------------------------
    # Column discovery
    # ------------------------------------------------------------------

    def _get_columns(self, table: str) -> List[Dict[str, str]]:
        query = self.sql.schema_columns_query(table)
        if not query:
            return []
        try:
            rows = self.client.run_query(query)
            result = []
            for row in rows:
                col_name = row.get('column_name') or row.get('COLUMN_NAME')
                data_type = row.get('data_type') or row.get('DATA_TYPE') or 'unknown'
                if col_name:
                    result.append({'column': col_name, 'data_type': data_type.lower()})
            return result
        except Exception as exc:
            logger.debug("Failed to get columns for %s: %s", table, exc)
            return []

    # ------------------------------------------------------------------
    # Table-level discovery
    # ------------------------------------------------------------------

    def _discover_in_table(
        self,
        table: str,
        columns: List[Dict[str, str]],
        entity_id: str,
        id_variants: List[str],
    ) -> List[Dict[str, Any]]:
        matches: List[Dict[str, Any]] = []
        id_is_numeric = self._entity_id_is_numeric(entity_id)

        for col_info in columns:
            column = col_info['column']
            data_type = col_info['data_type']

            if self._is_definitely_not_entity(column, data_type):
                continue
            if self._is_numeric_type(data_type) and not id_is_numeric:
                continue

            match = self._probe_column(table, column, data_type, entity_id, id_variants)
            if match:
                matches.append(match)

        return matches

    # ------------------------------------------------------------------
    # Column probing
    # ------------------------------------------------------------------

    def _probe_column(
        self, table: str, column: str, data_type: str,
        entity_id: str, id_variants: List[str],
    ) -> Optional[Dict[str, Any]]:
        if 'json' in data_type:
            return self._probe_json_column(table, column, data_type, entity_id, id_variants)
        return self._probe_scalar_column(table, column, data_type, entity_id, id_variants)

    def _probe_scalar_column(
        self, table: str, column: str, data_type: str,
        entity_id: str, id_variants: List[str],
    ) -> Optional[Dict[str, Any]]:
        is_numeric = self._is_numeric_type(data_type)
        is_uuid = 'uuid' in data_type

        # CONTAINS
        if not is_numeric and not is_uuid:
            for variant in id_variants:
                where = self.sql.contains(column, variant)
                hits = self._count_hits(table, where)
                if hits > 0:
                    return self._match(
                        table, column, data_type, "contains", variant,
                        None, hits, 0.7 if variant == entity_id else 0.6,
                    )

        # EXACT
        if not is_uuid:
            variants_to_try = [entity_id] if is_numeric else id_variants
            for variant in variants_to_try:
                where = self.sql.exact(column, variant)
                hits = self._count_hits(table, where)
                if hits > 0:
                    return self._match(
                        table, column, data_type, "exact", variant, None, hits, 0.9,
                    )

        # CAST
        if not is_numeric:
            variants_to_try = (
                [entity_id, entity_id.replace('-', '')] if is_uuid else [entity_id]
            )
            for variant in variants_to_try:
                where = self.sql.cast_match(column, variant)
                hits = self._count_hits(table, where)
                if hits > 0:
                    return self._match(
                        table, column, data_type, "cast_compare", variant, None, hits, 0.8,
                    )

        return None

    def _probe_json_column(
        self, table: str, column: str, data_type: str,
        entity_id: str, id_variants: List[str],
    ) -> Optional[Dict[str, Any]]:
        for variant in id_variants:
            where = self.sql.contains(column, variant)
            hits = self._count_hits(table, where)
            if hits > 0:
                json_path = self._discover_json_path(table, column, entity_id, id_variants)
                if json_path:
                    for v2 in id_variants:
                        where2 = self.sql.json_match(column, json_path, v2)
                        path_hits = self._count_hits(table, where2)
                        if path_hits > 0:
                            return self._match(
                                table, column, data_type, "json_extract",
                                v2, json_path, path_hits, 0.85,
                            )
                return self._match(
                    table, column, data_type, "contains", variant, None, hits, 0.6,
                )
        return None

    def _discover_json_path(
        self, table: str, column: str,
        entity_id: str, id_variants: List[str],
    ) -> Optional[str]:
        keys_query = self.sql.json_keys_query(table, column)
        keys: List[str] = []
        if keys_query:
            try:
                rows = self.client.run_query(keys_query)
                keys = [r.get('key') for r in rows if r.get('key')]
            except Exception:
                pass
        if not keys:
            keys = self._sample_json_keys_clientside(table, column)
        for key in keys:
            for variant in id_variants:
                where = self.sql.json_match(column, key, variant)
                if self._count_hits(table, where) > 0:
                    return key
        return None

    def _sample_json_keys_clientside(self, table: str, column: str, limit: int = 20) -> List[str]:
        col = self.sql._quote(column)
        query = f"SELECT DISTINCT {col} FROM {table} WHERE {col} IS NOT NULL LIMIT {limit}"
        try:
            rows = self.client.run_query(query)
        except Exception:
            return []
        keys: set = set()
        for row in rows:
            val = row.get(column) or (list(row.values())[0] if row else None)
            if val:
                try:
                    obj = json.loads(val) if isinstance(val, str) else val
                    if isinstance(obj, dict):
                        keys.update(obj.keys())
                except (json.JSONDecodeError, TypeError):
                    pass
        return list(keys)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _count_hits(self, table: str, where_clause: str, limit: int = 10) -> int:
        query = f"SELECT 1 FROM {table} WHERE {where_clause} LIMIT {limit}"
        try:
            rows = self.client.run_query(query)
            return len(rows)
        except Exception:
            return 0

    @staticmethod
    def _match(
        table: str, column: str, data_type: str, strategy: str,
        value_format: str, json_path: Optional[str],
        sample_hit_count: int, confidence: float,
    ) -> Dict[str, Any]:
        return {
            "table": table,
            "column": column,
            "data_type": data_type,
            "strategy": strategy,
            "value_format": value_format,
            "json_path": json_path,
            "sample_hit_count": sample_hit_count,
            "confidence": confidence,
        }

    @staticmethod
    def _is_definitely_not_entity(column: str, data_type: str) -> bool:
        if any(t in data_type for t in _SKIP_TYPE_TOKENS):
            return True
        if column.lower() in ('created_at', 'updated_at', 'modified_at', 'deleted_at'):
            return True
        return False

    @staticmethod
    def _is_numeric_type(data_type: str) -> bool:
        base = data_type.split('(')[0].strip()
        return base in _NUMERIC_TYPE_EXACT

    @staticmethod
    def _entity_id_is_numeric(entity_id: str) -> bool:
        try:
            int(entity_id)
            return True
        except ValueError:
            return False
