from __future__ import annotations

import logging
from typing import Any, Dict, 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)

        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 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
