from __future__ import annotations

import logging
import time
from typing import Any, Dict, Optional

from actor.connectors.llm import LLMClient
from actor.observability import ObservabilityService, ComponentStatus
from actor.runtime import RuntimeConfig
from actor.authentication.config import SecretStoreConfig
from actor.authentication.secrets import build_secret_provider, CompositeSecretProvider
from actor.authentication.llm import build_llm_config
from actor.runtime.db_tasks import DBTaskExecutor

logger = logging.getLogger(__name__)


def _categorize_error(task_type: str, exc: Exception) -> str:
    """Categorize an exception for metrics tracking."""
    exc_str = str(exc).lower()

    # Check for timeout first
    if isinstance(exc, TimeoutError) or "timeout" in exc_str:
        return "timeout"

    # Categorize by task type
    if task_type and task_type.startswith("db_"):
        return "db_error"
    if task_type in ("llm_validate", "classify_fields"):
        return "llm_error"
    if task_type == "pipeline_validate":
        return "db_error" if "database" in exc_str else "llm_error"

    # Fallback based on exception content
    if "llm" in exc_str or "openai" in exc_str or "completion" in exc_str:
        return "llm_error"
    if "desensitiz" in exc_str:
        return "desensitization_error"
    if "database" in exc_str or "sql" in exc_str or "connection" in exc_str:
        return "db_error"

    return "unknown_error"


class TaskRouter:
    """Routes queue tasks to the appropriate connector."""

    def __init__(self, runtime_config: RuntimeConfig, observability: ObservabilityService):
        self.runtime_config = runtime_config
        self.observability = observability
        self._db_executor: Optional[DBTaskExecutor] = None
        self._llm_client: Optional[LLMClient] = None

    def _ensure_llm(self) -> Optional[LLMClient]:
        """Return LLM client if configured, None otherwise (not an error)."""
        if self._llm_client:
            return self._llm_client
        cfg = self.runtime_config.get_llm()
        if not cfg:
            return None
        overrides = self._materialize_api_secrets(cfg, build_llm_config)
        self._llm_client = LLMClient(
            endpoint=overrides.get("endpoint", cfg.endpoint),
            auth_type=cfg.auth_type,
            api_key=overrides.get("api_key", cfg.api_key),
            username=cfg.username,
            password=cfg.password,
            api_version=overrides.get("api_version", cfg.api_version),
            deployment_name=cfg.deployment_name,
            model=overrides.get("model", cfg.model),
            observability=self.observability,
        )
        return self._llm_client

    def _ensure_db(self) -> DBTaskExecutor:
        """Create or return cached DBTaskExecutor with optional LLM client."""
        if self._db_executor:
            return self._db_executor
        cfg = self.runtime_config.get_database()
        if not cfg:
            raise RuntimeError("Database config not available from SaaS")
        llm_client = self._ensure_llm()
        self._db_executor = DBTaskExecutor(
            runtime_config=self.runtime_config,
            observability=self.observability,
            llm_client=llm_client,
        )
        return self._db_executor

    @staticmethod
    def _materialize_api_secrets(cfg, build_fn) -> Dict[str, Any]:
        """Resolve secrets via SecretProvider if secret_store + mapping present.

        Returns a dict of overrides for endpoint/api_key/model/api_version.
        """
        overrides: Dict[str, Any] = {}
        if cfg.secret_store and cfg.secret_mapping:
            store = SecretStoreConfig.parse_obj(cfg.secret_store)
            base_provider = build_secret_provider(store)
            literal_values: Dict[str, Any] = {}
            mapping = dict(cfg.secret_mapping)
            for key in ("endpoint", "model", "api_version", "deployment_name"):
                val = getattr(cfg, key, None)
                if key not in mapping and val:
                    mapping[key] = key
                    literal_values[key] = val
            provider = CompositeSecretProvider(base_provider, literal_values)
            mat = build_fn(cfg.auth_type, mapping, provider)
            for src, dst in (("base_url", "endpoint"), ("api_key", "api_key"), ("model", "model"), ("api_version", "api_version")):
                if src in mat:
                    overrides[dst] = mat[src]
        return overrides

    def shutdown(self) -> None:
        """Release any held resources before process exit."""
        if self._db_executor:
            try:
                self._db_executor.dispose()
            finally:
                self._db_executor = None

    def handle(self, task: Dict[str, Any]) -> Dict[str, Any]:
        ttype = task.get("task_type")
        task_id = task.get("task_id", "unknown")
        payload = task.get("payload") or {}

        # Record task start for metrics
        self.observability.record_task_start(task_id, ttype or "unknown")
        start_time = time.time()

        try:
            result = self._execute_task(ttype, payload)
            # Record success with execution time
            execution_time_ms = int((time.time() - start_time) * 1000)
            self.observability.record_task_success(task_id, execution_time_ms)
            try:
                self.observability.send_metrics_report()
            except Exception:
                logger.debug("Metrics publish failed after task success", exc_info=True)
            return result
        except Exception as exc:
            # Categorize and record the failure
            error_type = _categorize_error(ttype, exc)
            self.observability.record_task_failure(task_id, error_type, str(exc))
            try:
                self.observability.send_metrics_report()
            except Exception:
                logger.debug("Metrics publish failed after task failure", exc_info=True)
            raise

    def _execute_task(self, ttype: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a task and return the result."""
        # DB tasks (db_validate, db_schemas, db_tables, db_columns, db_sample, db_query)
        if ttype and ttype.startswith("db_"):
            return self._ensure_db().handle({"task_type": ttype, "payload": payload})

        # classify_fields delegates to DBTaskExecutor (needs DB for sampling + LLM for classification)
        if ttype == "classify_fields":
            return self._ensure_db().handle({"task_type": ttype, "payload": payload})

        if ttype == "llm_validate":
            llm = self._ensure_llm()
            if not llm:
                raise RuntimeError("LLM config not available from SaaS")
            start = time.time()
            llm.validate()
            latency_ms = int((time.time() - start) * 1000)
            return {
                "status": "ok",
                "model": llm.model,
                "endpoint": llm.endpoint,
                "latency_ms": latency_ms,
            }

        if ttype == "pipeline_validate":
            return self._handle_pipeline_validate()

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

    def _handle_pipeline_validate(self) -> Dict[str, Any]:
        """End-to-end DB+LLM check."""
        result: Dict[str, Any] = {}

        # Validate DB
        db_cfg = self.runtime_config.get_database()
        if db_cfg:
            try:
                executor = self._ensure_db()
                executor.validate()
                result["database"] = {"status": "ok", "database_type": getattr(db_cfg, "database_type", None)}
            except Exception as exc:
                result["database"] = {"status": "error", "error": str(exc)}
        else:
            result["database"] = {"status": "not_configured"}

        # Validate LLM (optional)
        llm_client = self._ensure_llm()
        if llm_client:
            try:
                start = time.time()
                llm_client.validate()
                latency_ms = int((time.time() - start) * 1000)
                result["llm"] = {"status": "ok", "model": llm_client.model, "latency_ms": latency_ms}
            except Exception as exc:
                result["llm"] = {"status": "error", "error": str(exc)}
        else:
            result["llm"] = {"status": "not_configured"}

        # Desensitization status
        if llm_client:
            result["desensitization"] = "active"
        elif self.runtime_config.llm_passthrough:
            result["desensitization"] = "passthrough"
        else:
            result["desensitization"] = "blocked"

        return result
