"""
Databricks Unity Catalog connector for AutoPIL.

Reads sensitivity labels, column tags, and table properties from Unity Catalog
and maps them to AutoPIL CatalogClassification records.

Install:
    pip install autopil[catalog-unity]

Required credentials (passed as 'credentials' dict when registering connection):
    host:       Databricks workspace URL  (e.g. https://my-workspace.azuredatabricks.net)
    token:      Personal access token or service principal OAuth token
    catalog:    Unity Catalog name to scan  (e.g. "main")
    schema:     (optional) Restrict scan to one schema  (e.g. "finance")

Sensitivity mapping:
    AutoPIL reads the following, in priority order:
      1. Column-level tag: "sensitivity_level" or "data_sensitivity" or "autopil_sensitivity"
      2. Table-level property: "sensitivity_level" or "data_sensitivity"
      3. Table-level tag: same keys
    If none found, defaults to "medium".

Unity Catalog write-back:
    Enforcement events are appended as table-level comments via the SDK.
    This is best-effort and never blocks evaluation.
"""

from __future__ import annotations

import json
import logging
import secrets
from datetime import datetime
from typing import TYPE_CHECKING

from autopil.catalog.base import CatalogConnector
from autopil.catalog.registry import CatalogRegistry
from autopil.models import CatalogClassification, CatalogConnection

if TYPE_CHECKING:
    from autopil.models import AuditEvent

logger = logging.getLogger(__name__)

# Keys AutoPIL looks for when extracting sensitivity from UC properties/tags
_SENSITIVITY_KEYS = ("sensitivity_level", "data_sensitivity", "autopil_sensitivity")


def _creds(connection: CatalogConnection) -> dict:
    """Decrypt and parse credentials JSON. Falls back to plaintext JSON."""
    raw = connection.credentials_enc
    enc_key = _get_encryption_key()
    if enc_key:
        try:
            from cryptography.fernet import Fernet
            return json.loads(Fernet(enc_key.encode()).decrypt(raw.encode()))
        except Exception:
            pass
    return json.loads(raw)


def _get_encryption_key() -> str | None:
    import os
    return os.environ.get("AUTOPIL_CATALOG_ENCRYPTION_KEY")


def _encrypt(data: dict) -> str:
    enc_key = _get_encryption_key()
    raw = json.dumps(data)
    if enc_key:
        try:
            from cryptography.fernet import Fernet
            return Fernet(enc_key.encode()).encrypt(raw.encode()).decode()
        except ImportError:
            logger.warning(
                "cryptography package not installed; storing catalog credentials as plaintext. "
                "Install with: pip install cryptography"
            )
    return raw


class UnityCatalogConnector(CatalogConnector):
    """
    AutoPIL connector for Databricks Unity Catalog.

    Uses the databricks-sdk WorkspaceClient to enumerate tables and read
    sensitivity tags and table properties.
    """

    catalog_type = "unity_catalog"

    def _client(self):
        try:
            from databricks.sdk import WorkspaceClient
        except ImportError as exc:
            raise ImportError(
                "databricks-sdk is required for the Unity Catalog connector. "
                "Install with: pip install autopil[catalog-unity]"
            ) from exc
        creds = _creds(self.connection)
        return WorkspaceClient(host=creds["host"], token=creds["token"])

    def test_connection(self) -> tuple[bool, str]:
        try:
            creds = _creds(self.connection)
            if "host" not in creds or "token" not in creds:
                return False, "Missing required credentials: 'host' and 'token'"
            if "catalog" not in creds:
                return False, "Missing required credential: 'catalog'"
            w = self._client()
            # Try listing catalogs to verify auth
            list(w.catalogs.list())
            return True, ""
        except ImportError as exc:
            return False, str(exc)
        except Exception as exc:
            return False, f"Connection failed: {exc}"

    def fetch_classifications(
        self,
        asset_fqns: list[str] | None = None,
    ) -> list[CatalogClassification]:
        creds   = _creds(self.connection)
        catalog = creds["catalog"]
        schema  = creds.get("schema")          # optional filter
        w       = self._client()

        classifications: list[CatalogClassification] = []
        now = datetime.utcnow()

        try:
            schemas_to_scan = [schema] if schema else self._list_schemas(w, catalog)
        except Exception as exc:
            raise RuntimeError(f"Failed to list schemas in catalog '{catalog}': {exc}") from exc

        for schema_name in schemas_to_scan:
            try:
                tables = list(w.tables.list(catalog_name=catalog, schema_name=schema_name))
            except Exception as exc:
                logger.warning("Could not list tables in %s.%s: %s", catalog, schema_name, exc)
                continue

            for tbl in tables:
                full_name = f"{catalog}.{schema_name}.{tbl.name}"

                if asset_fqns and not any(fqn in full_name for fqn in asset_fqns):
                    continue

                sensitivity, raw_label, data_classes = self._extract_sensitivity(tbl)

                cls = CatalogClassification(
                    classification_id="cls_" + secrets.token_hex(8),
                    connection_id=self.connection.connection_id,
                    tenant_id=self.connection.tenant_id,
                    asset_fqn=full_name,
                    asset_type="table",
                    sensitivity_label=sensitivity,
                    raw_label=raw_label,
                    data_classes=data_classes,
                    is_certified=getattr(tbl, "is_managed_external_table", None) is not None,
                    trust_score=None,
                    catalog_policy_ids=[],
                    fetched_at=now,
                    expires_at=now,  # set by scheduler
                )
                classifications.append(cls)

                # Column-level classifications for higher-sensitivity columns
                col_clss = self._column_classifications(tbl, full_name, now)
                classifications.extend(col_clss)

        logger.info(
            "Unity Catalog fetch complete: catalog=%s tables=%d classifications=%d",
            catalog, len(schemas_to_scan), len(classifications),
        )
        return classifications

    def write_enforcement_event(
        self,
        event: AuditEvent,
        classification: CatalogClassification | None,
    ) -> bool:
        """
        Append an AutoPIL enforcement record to the table's comments field via SDK.
        Best-effort only — returns False on any failure.
        """
        try:
            creds = _creds(self.connection)
            w     = self._client()
            # Build a compact enforcement record
            record = {
                "autopil_event_id": event.event_id,
                "decision":         event.decision.value,
                "agent_role":       event.agent_role,
                "source_id":        event.source_id,
                "policy_name":      event.policy_name,
                "timestamp":        event.timestamp.isoformat(),
            }
            # Unity Catalog doesn't have a native audit write-back API per se;
            # we tag the table with the last enforcement event_id as a property.
            if classification:
                parts = classification.asset_fqn.split(".")
                if len(parts) == 3:
                    w.tables.update(
                        full_name=classification.asset_fqn,
                        properties={"autopil_last_event_id": event.event_id,
                                    "autopil_last_decision": event.decision.value},
                    )
            return True
        except Exception as exc:
            logger.debug("Unity Catalog write-back failed: %s", exc)
            return False

    # ── helpers ───────────────────────────────────────────────────────────────

    def _list_schemas(self, w, catalog: str) -> list[str]:
        return [s.name for s in w.schemas.list(catalog_name=catalog) if s.name]

    def _extract_sensitivity(self, tbl) -> tuple[str, str, list[str]]:
        """
        Extract sensitivity level, raw label, and data classes from a table object.
        Checks table properties then falls back to the base mapper.
        """
        props = getattr(tbl, "properties", None) or {}
        data_classes: list[str] = []

        for key in _SENSITIVITY_KEYS:
            if key in props:
                raw = props[key]
                mapped = self.map_label_to_sensitivity(raw)
                # Also look for data_class or data_classification property
                dc = props.get("data_class") or props.get("data_classification") or ""
                if dc:
                    data_classes = [c.strip() for c in dc.split(",") if c.strip()]
                return mapped, raw, data_classes

        # No explicit label — check table_type for hints
        table_type = str(getattr(tbl, "table_type", "") or "").upper()
        if table_type in ("EXTERNAL",):
            return "medium", "external_table", []
        return "medium", "unclassified", []

    def _column_classifications(
        self, tbl, table_fqn: str, now: datetime
    ) -> list[CatalogClassification]:
        """
        Generate column-level classifications for columns with explicit sensitivity tags.
        Only emits a classification if the column sensitivity > the table level.
        """
        result = []
        columns = getattr(tbl, "columns", None) or []
        for col in columns:
            col_props = getattr(col, "metadata", None) or {}
            if not col_props:
                continue
            for key in _SENSITIVITY_KEYS:
                if key in col_props:
                    raw = col_props[key]
                    mapped = self.map_label_to_sensitivity(raw)
                    col_fqn = f"{table_fqn}.{col.name}"
                    result.append(CatalogClassification(
                        classification_id="cls_" + secrets.token_hex(8),
                        connection_id=self.connection.connection_id,
                        tenant_id=self.connection.tenant_id,
                        asset_fqn=col_fqn,
                        asset_type="column",
                        sensitivity_label=mapped,
                        raw_label=raw,
                        data_classes=[],
                        is_certified=None,
                        trust_score=None,
                        catalog_policy_ids=[],
                        fetched_at=now,
                        expires_at=now,
                    ))
                    break
        return result


# Register on import
CatalogRegistry.register(UnityCatalogConnector)


# ── Credential helper for API layer ──────────────────────────────────────────

def encrypt_credentials(credentials: dict) -> str:
    """Encrypt a credentials dict for storage. Used by the API when registering a connection."""
    return _encrypt(credentials)
