"""
ClassificationEnricher — enriches ContextRequests with catalog-sourced sensitivity metadata.

Sits between ContextGuard._evaluate_request() and PolicyEngine.evaluate().
It is strictly additive:
  - It can only RAISE sensitivity_level (never lower it)
  - If no catalog classification is found, the request passes through unchanged
  - If the catalog store is unavailable, the request passes through unchanged (fail-open)
"""

from __future__ import annotations

import logging
import threading
from typing import TYPE_CHECKING

from autopil.models import SENSITIVITY_ORDER, AuditEvent, CatalogClassification, ContextRequest, SensitivityLevel

if TYPE_CHECKING:
    from autopil.db.base import CatalogStoreBase

logger = logging.getLogger(__name__)


class ClassificationEnricher:
    """
    Enriches a ContextRequest with catalog-sourced sensitivity metadata before
    policy evaluation, and schedules catalog write-back after evaluation.

    Args:
        catalog_store: The CatalogStoreBase used to look up cached classifications.
        write_back_enabled: If True, write enforcement events back to the catalog
                            asynchronously after each evaluation. Default True.
    """

    def __init__(
        self,
        catalog_store: CatalogStoreBase,
        write_back_enabled: bool = True,
    ) -> None:
        self._store = catalog_store
        self._write_back_enabled = write_back_enabled

    def enrich(self, request: ContextRequest, tenant_id: str) -> ContextRequest:
        """
        Look up catalog classifications for request.source_id and return an enriched
        copy of the request with the sensitivity_level raised if the catalog says higher.

        The original request object is never mutated.
        Returns the original request unchanged on any lookup failure.
        """
        try:
            classifications = self._store.get_classifications_for_asset(
                tenant_id, request.source_id
            )
        except Exception:
            logger.warning(
                "Catalog lookup failed for source_id=%r (tenant=%s); "
                "continuing with caller-supplied sensitivity",
                request.source_id, tenant_id,
                exc_info=True,
            )
            return request

        if not classifications:
            return request

        # Pick the most restrictive classification among all matches
        best = _most_restrictive(classifications)
        catalog_level = best.sensitivity_label

        caller_idx  = SENSITIVITY_ORDER.index(request.sensitivity_level.value)
        catalog_idx = SENSITIVITY_ORDER.index(catalog_level) if catalog_level in SENSITIVITY_ORDER else caller_idx

        if catalog_idx <= caller_idx:
            # Catalog is same or less restrictive — no change needed
            return request

        # Upgrade: return a new ContextRequest with the higher sensitivity level
        upgraded = ContextRequest(
            query=request.query,
            agent_role=request.agent_role,
            user_id=request.user_id,
            source_id=request.source_id,
            sensitivity_level=SensitivityLevel(catalog_level),
            session_id=request.session_id,
            timestamp=request.timestamp,
            task_type=request.task_type,
            catalog_data_classes=list(best.data_classes),
        )
        logger.debug(
            "Catalog enrichment: source=%r sensitivity %s → %s (catalog=%s, asset=%s)",
            request.source_id,
            request.sensitivity_level.value,
            catalog_level,
            best.connection_id,
            best.asset_fqn,
        )
        return upgraded

    def write_back(
        self,
        event: AuditEvent,
        tenant_id: str,
        connection_id: str | None = None,
    ) -> None:
        """
        Fire-and-forget: write an enforcement event back to the catalog(s).
        Runs on a daemon thread so it never blocks the caller.
        """
        if not self._write_back_enabled:
            return
        t = threading.Thread(
            target=self._do_write_back,
            args=(event, tenant_id, connection_id),
            daemon=True,
            name=f"autopil-catalog-writeback-{event.event_id[:8]}",
        )
        t.start()

    def _do_write_back(
        self,
        event: AuditEvent,
        tenant_id: str,
        connection_id: str | None,
    ) -> None:
        try:
            from autopil.catalog.registry import CatalogRegistry

            connections = self._store.list_connections(tenant_id)
            if connection_id:
                connections = [c for c in connections if c.connection_id == connection_id]

            for conn in connections:
                if not conn.is_enabled:
                    continue
                try:
                    connector = CatalogRegistry.get_connector(conn)
                    # Try to find the matching classification for context
                    clss = self._store.get_classifications_for_asset(tenant_id, event.source_id)
                    cls = clss[0] if clss else None
                    connector.write_enforcement_event(event, cls)
                except Exception:
                    logger.debug(
                        "Catalog write-back failed for connection=%s", conn.connection_id,
                        exc_info=True,
                    )
        except Exception:
            logger.debug("Catalog write-back thread error", exc_info=True)


def _most_restrictive(classifications: list[CatalogClassification]) -> CatalogClassification:
    """Return the classification with the highest sensitivity_label."""
    def _rank(cls: CatalogClassification) -> int:
        try:
            return SENSITIVITY_ORDER.index(cls.sensitivity_label)
        except ValueError:
            return 0  # unknown labels treated as low

    return max(classifications, key=_rank)
