"""
Collibra Data Intelligence Cloud connector for AutoPIL.

Reads asset sensitivity classifications from Collibra via the REST 2.0 API
and maps them to AutoPIL CatalogClassification records.

Install:
    pip install autopil[catalog-collibra]

Required credentials (passed as 'credentials' dict when registering connection):
    username:     Collibra user with read access to the catalog
    password:     Collibra password  (or use client_id / client_secret for OAuth)
    domain_id:    (optional) Restrict scan to a specific Collibra domain UUID

Optional credentials:
    client_id:    OAuth2 client ID  (if using client credentials flow instead of basic auth)
    client_secret: OAuth2 client secret

Sensitivity mapping:
    AutoPIL reads the attribute type named "Sensitivity" or "Data Classification" on each asset.
    If neither is present, assets default to "medium".

Collibra write-back:
    Enforcement events are written as a custom attribute "AutoPIL Enforcement" on the asset.
    The attribute type must exist in Collibra; write-back is best-effort and fail-open.
"""

from __future__ import annotations

import json
import logging
import os
import secrets
from datetime import datetime, timedelta
from typing import TYPE_CHECKING

try:
    import requests
except ImportError as exc:
    raise ImportError(
        "requests is required for the Collibra connector. "
        "Install with: pip install autopil[catalog-collibra]"
    ) from exc

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

# Collibra attribute type names AutoPIL looks for (case-insensitive match)
_SENSITIVITY_ATTR_NAMES = ("sensitivity", "data classification", "data sensitivity", "security classification")

# Collibra-specific label overrides on top of the base mapper
_COLLIBRA_LABEL_MAP: dict[str, str] = {
    "public":               "low",
    "open":                 "low",
    "internal":             "medium",
    "internal use only":    "medium",
    "confidential":         "high",
    "confidential class 1": "high",
    "confidential class 2": "high",
    "restricted":           "restricted",
    "highly confidential":  "restricted",
    "sensitive":            "high",
    "pii":                  "high",
    "phi":                  "restricted",
    "financial":            "high",
}


def _creds(connection: CatalogConnection) -> dict:
    raw = connection.credentials_enc
    enc_key = os.environ.get("AUTOPIL_CATALOG_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)


class CollibraConnector(CatalogConnector):
    """
    AutoPIL connector for Collibra Data Intelligence Cloud.

    Uses the Collibra REST 2.0 API with either basic auth (username/password)
    or OAuth2 client credentials. Scans asset sensitivity attributes and
    maps them to AutoPIL sensitivity levels.
    """

    catalog_type = "collibra"

    def _base_url(self) -> str:
        return self.connection.base_url.rstrip("/")

    def _session(self) -> requests.Session:
        creds = _creds(self.connection)
        session = requests.Session()
        session.headers.update({"Accept": "application/json", "Content-Type": "application/json"})

        if "client_id" in creds and "client_secret" in creds:
            # OAuth2 client credentials flow
            token_url = f"{self._base_url()}/rest/2.0/auth/oauth/token"
            resp = session.post(token_url, data={
                "grant_type":    "client_credentials",
                "client_id":     creds["client_id"],
                "client_secret": creds["client_secret"],
            }, timeout=10)
            resp.raise_for_status()
            token = resp.json()["access_token"]
            session.headers["Authorization"] = f"Bearer {token}"
        else:
            # Basic auth
            session.auth = (creds["username"], creds["password"])

        return session

    def test_connection(self) -> tuple[bool, str]:
        try:
            creds = _creds(self.connection)
            required = ("username", "password") if "client_id" not in creds else ("client_id", "client_secret")
            missing = [k for k in required if k not in creds]
            if missing:
                return False, f"Missing required credentials: {missing}"

            s = self._session()
            url = f"{self._base_url()}/rest/2.0/ping"
            resp = s.get(url, timeout=10)
            if resp.status_code == 404:
                # Older Collibra instances may not have /ping — try /currentUser
                resp = s.get(f"{self._base_url()}/rest/2.0/users/currentUser", timeout=10)
            resp.raise_for_status()
            return True, ""
        except ImportError as exc:
            return False, str(exc)
        except Exception as exc:
            return False, f"Collibra connection failed: {exc}"

    def fetch_classifications(
        self,
        asset_fqns: list[str] | None = None,
    ) -> list[CatalogClassification]:
        creds  = _creds(self.connection)
        s      = self._session()
        now    = datetime.utcnow()
        domain_id = creds.get("domain_id")

        classifications: list[CatalogClassification] = []

        # Step 1: find assets (paginated)
        assets = self._fetch_assets(s, domain_id, asset_fqns)

        # Step 2: for each asset, get sensitivity attribute
        for asset in assets:
            asset_id   = asset.get("id", "")
            asset_name = asset.get("name", "")
            domain     = (asset.get("domain") or {}).get("name", "")
            fqn        = f"{domain}.{asset_name}" if domain else asset_name

            sensitivity, raw_label, data_classes = self._get_sensitivity(s, asset_id, asset)

            classifications.append(CatalogClassification(
                classification_id="cls_" + secrets.token_hex(8),
                connection_id=self.connection.connection_id,
                tenant_id=self.connection.tenant_id,
                asset_fqn=fqn,
                asset_type=self._map_asset_type(asset),
                sensitivity_label=sensitivity,
                raw_label=raw_label,
                data_classes=data_classes,
                is_certified=asset.get("status", {}).get("name", "").lower() == "accepted",
                trust_score=None,
                catalog_policy_ids=[],
                fetched_at=now,
                expires_at=now,  # set by scheduler
            ))

        logger.info(
            "Collibra fetch complete: domain=%s assets=%d classifications=%d",
            domain_id or "all", len(assets), len(classifications),
        )
        return classifications

    def write_enforcement_event(
        self,
        event: AuditEvent,
        classification: CatalogClassification | None,
    ) -> bool:
        """
        Write an AutoPIL enforcement record back to Collibra as an asset attribute.

        The attribute type "AutoPIL Enforcement" must exist in Collibra. If it
        doesn't, this is a no-op (returns False). Best-effort — never raises.
        """
        if classification is None:
            return False
        try:
            s = self._session()
            # Look up the asset ID by FQN
            asset_id = self._find_asset_id(s, classification.asset_fqn)
            if not asset_id:
                return False

            # Look up the attribute type ID for "AutoPIL Enforcement"
            attr_type_id = self._find_attribute_type_id(s, "AutoPIL Enforcement")
            if not attr_type_id:
                logger.debug(
                    "Collibra write-back skipped: 'AutoPIL Enforcement' attribute type not found"
                )
                return False

            payload = {
                "assetId":       asset_id,
                "typeId":        attr_type_id,
                "value":         json.dumps({
                    "event_id":    event.event_id,
                    "decision":    event.decision.value,
                    "agent_role":  event.agent_role,
                    "policy_name": event.policy_name,
                    "timestamp":   event.timestamp.isoformat(),
                }),
            }
            resp = s.post(
                f"{self._base_url()}/rest/2.0/attributes",
                json=payload,
                timeout=10,
            )
            resp.raise_for_status()
            return True
        except Exception as exc:
            logger.debug("Collibra write-back failed: %s", exc)
            return False

    def map_label_to_sensitivity(self, raw_label: str) -> str:
        normalised = raw_label.lower().strip().replace("-", " ").replace("_", " ")
        return _COLLIBRA_LABEL_MAP.get(normalised, super().map_label_to_sensitivity(raw_label))

    # ── internal helpers ──────────────────────────────────────────────────────

    def _fetch_assets(
        self,
        s: requests.Session,
        domain_id: str | None,
        asset_fqns: list[str] | None,
    ) -> list[dict]:
        """Paginate through Collibra assets, optionally filtered by domain or FQN list."""
        assets: list[dict] = []
        offset, limit = 0, 100

        while True:
            params: dict = {"offset": offset, "limit": limit, "sortOrder": "ASC"}
            if domain_id:
                params["domainId"] = domain_id

            try:
                resp = s.get(f"{self._base_url()}/rest/2.0/assets", params=params, timeout=30)
                resp.raise_for_status()
                data = resp.json()
            except Exception as exc:
                logger.warning("Collibra asset fetch failed at offset %d: %s", offset, exc)
                break

            batch = data.get("results", [])
            if not batch:
                break

            # If caller wants specific FQNs, filter here (client-side)
            if asset_fqns:
                batch = [
                    a for a in batch
                    if any(fqn in f"{(a.get('domain') or {}).get('name','')}.{a.get('name','')}"
                           for fqn in asset_fqns)
                ]
            assets.extend(batch)

            if len(data.get("results", [])) < limit:
                break
            offset += limit

        return assets

    def _get_sensitivity(
        self, s: requests.Session, asset_id: str, asset: dict
    ) -> tuple[str, str, list[str]]:
        """Fetch sensitivity attribute for a single asset. Falls back to medium."""
        try:
            resp = s.get(
                f"{self._base_url()}/rest/2.0/attributes",
                params={"assetId": asset_id, "limit": 50},
                timeout=10,
            )
            resp.raise_for_status()
            attrs = resp.json().get("results", [])
        except Exception:
            return "medium", "unclassified", []

        data_classes: list[str] = []
        for attr in attrs:
            type_name = (attr.get("type") or {}).get("name", "").lower().strip()
            if type_name in _SENSITIVITY_ATTR_NAMES:
                raw = str(attr.get("value") or "").strip()
                if raw:
                    return self.map_label_to_sensitivity(raw), raw, data_classes
            elif type_name in ("data class", "data classification type", "information type"):
                val = str(attr.get("value") or "").strip()
                if val:
                    data_classes.append(val)

        return "medium", "unclassified", data_classes

    def _find_asset_id(self, s: requests.Session, asset_fqn: str) -> str | None:
        """Look up a Collibra asset ID by name (best-effort, returns first match)."""
        try:
            name = asset_fqn.split(".")[-1]  # last segment as the asset name
            resp = s.get(
                f"{self._base_url()}/rest/2.0/assets",
                params={"name": name, "nameMatchMode": "EXACT", "limit": 1},
                timeout=10,
            )
            resp.raise_for_status()
            results = resp.json().get("results", [])
            return results[0]["id"] if results else None
        except Exception:
            return None

    def _find_attribute_type_id(self, s: requests.Session, type_name: str) -> str | None:
        """Look up a Collibra attribute type ID by name."""
        try:
            resp = s.get(
                f"{self._base_url()}/rest/2.0/attributeTypes",
                params={"name": type_name, "nameMatchMode": "EXACT", "limit": 1},
                timeout=10,
            )
            resp.raise_for_status()
            results = resp.json().get("results", [])
            return results[0]["id"] if results else None
        except Exception:
            return None

    def _map_asset_type(self, asset: dict) -> str:
        type_name = (asset.get("type") or {}).get("name", "").lower()
        if "table" in type_name:
            return "table"
        if "column" in type_name or "field" in type_name:
            return "column"
        if "schema" in type_name or "database" in type_name:
            return "schema"
        return "asset"


# Register on import
CatalogRegistry.register(CollibraConnector)
