"""
Immuta Data Security Platform connector for AutoPIL.

Reads sensitivity classifications from the Immuta REST API and maps them
to AutoPIL CatalogClassification records.

Install:
    pip install autopil[catalog-immuta]

Required credentials (passed as 'credentials' dict when registering connection):
    api_key:    Immuta API key (generated in Settings → API Keys)

Optional credentials:
    page_size:  Number of data sources per page (default: 100)

Sensitivity mapping:
    Immuta sensitivity levels mapped to AutoPIL scale:
        "Highly Sensitive" / "Critical"     → restricted
        "Sensitive" / "Confidential"        → high
        "Moderate" / "Internal"             → medium
        "Low" / "Public"                    → low

Write-back:
    Enforcement events are recorded as Immuta audit log entries via the
    /audit endpoint. 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 Immuta connector. "
        "Install with: pip install autopil[catalog-immuta]"
    ) 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__)

_SENSITIVITY_MAP: dict[str, str] = {
    "highly sensitive": "restricted",
    "critical": "restricted",
    "sensitive": "high",
    "confidential": "high",
    "moderate": "medium",
    "internal": "medium",
    "low": "low",
    "public": "low",
}


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 ImmutaConnector(CatalogConnector):
    """
    AutoPIL connector for Immuta Data Security Platform.

    Uses the Immuta REST API to enumerate data sources and their
    sensitivity classifications, then writes enforcement events back
    as audit log entries.
    """

    catalog_type = "immuta"

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

    def _session(self) -> requests.Session:
        creds = _creds(self.connection)
        s = requests.Session()
        s.headers.update({
            "Authorization": f"Bearer {creds['api_key']}",
            "Accept": "application/json",
            "Content-Type": "application/json",
        })
        return s

    def test_connection(self) -> tuple[bool, str]:
        try:
            creds = _creds(self.connection)
            if "api_key" not in creds:
                return False, "Missing required credential: 'api_key'"
            s = self._session()
            resp = s.get(f"{self._base_url()}/bots/sensitivityClassification", timeout=10)
            if resp.status_code == 404:
                # Try the health endpoint as fallback
                resp = s.get(f"{self._base_url()}/health", timeout=10)
            resp.raise_for_status()
            return True, ""
        except ImportError as exc:
            return False, str(exc)
        except Exception as exc:
            return False, f"Immuta connection failed: {exc}"

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

        # First try the dedicated sensitivity classification endpoint
        classifications = self._fetch_via_sensitivity_api(s, now, page_size, asset_fqns)
        if not classifications:
            # Fall back to enumerating data sources
            classifications = self._fetch_via_datasources(s, now, page_size, asset_fqns)

        logger.info("Immuta fetch complete: classifications=%d", len(classifications))
        return classifications

    def write_enforcement_event(
        self,
        event: AuditEvent,
        classification: CatalogClassification | None,
    ) -> bool:
        """
        Write enforcement event to Immuta audit log. Best-effort and fail-open.
        """
        if classification is None:
            return False
        try:
            s = self._session()
            payload = {
                "type": "autopil_enforcement",
                "eventId": event.event_id,
                "decision": event.decision.value,
                "agentRole": event.agent_role,
                "sourceId": classification.asset_fqn,
                "policyName": event.policy_name,
                "reason": getattr(event, "reason", ""),
                "timestamp": event.timestamp.isoformat(),
            }
            resp = s.post(f"{self._base_url()}/audit", json=payload, timeout=10)
            return resp.status_code in (200, 201, 204)
        except Exception as exc:
            logger.debug("Immuta write-back failed: %s", exc)
            return False

    def map_label_to_sensitivity(self, raw_label: str) -> str:
        normalized = raw_label.strip().lower()
        return _SENSITIVITY_MAP.get(normalized, super().map_label_to_sensitivity(raw_label))

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

    def _fetch_via_sensitivity_api(
        self,
        s: requests.Session,
        now: datetime,
        page_size: int,
        asset_fqns: list[str] | None,
    ) -> list[CatalogClassification]:
        """Fetch from /bots/sensitivityClassification (Immuta >=2022.2)."""
        classifications: list[CatalogClassification] = []
        offset = 0

        while True:
            try:
                resp = s.get(
                    f"{self._base_url()}/bots/sensitivityClassification",
                    params={"limit": page_size, "offset": offset},
                    timeout=30,
                )
                if resp.status_code == 404:
                    return []  # endpoint not available on this Immuta version
                resp.raise_for_status()
                data = resp.json()
            except Exception as exc:
                logger.warning("Immuta sensitivity API failed (offset=%d): %s", offset, exc)
                break

            items = data if isinstance(data, list) else data.get("hits", data.get("items", []))
            if not items:
                break

            for item in items:
                fqn = item.get("dataSourceName") or item.get("name") or item.get("id", "")
                if not fqn:
                    continue
                if asset_fqns and not any(f in fqn for f in asset_fqns):
                    continue

                raw_label = (
                    item.get("sensitivityLevel")
                    or item.get("classification")
                    or "unclassified"
                )
                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="table",
                    sensitivity_label=self.map_label_to_sensitivity(raw_label),
                    raw_label=raw_label,
                    data_classes=item.get("dataClasses", []),
                    is_certified=item.get("isCertified", False),
                    trust_score=None,
                    catalog_policy_ids=[str(p) for p in item.get("policyIds", [])],
                    fetched_at=now,
                    expires_at=now + timedelta(minutes=self.connection.sync_interval_minutes),
                ))

            total = data.get("count", data.get("total", len(items))) if isinstance(data, dict) else len(items)
            offset += len(items)
            if isinstance(data, list) or offset >= total or len(items) < page_size:
                break

        return classifications

    def _fetch_via_datasources(
        self,
        s: requests.Session,
        now: datetime,
        page_size: int,
        asset_fqns: list[str] | None,
    ) -> list[CatalogClassification]:
        """Fallback: enumerate /dataSource and infer sensitivity from tags/policies."""
        classifications: list[CatalogClassification] = []
        offset = 0

        while True:
            try:
                resp = s.get(
                    f"{self._base_url()}/dataSource",
                    params={"size": page_size, "offset": offset},
                    timeout=30,
                )
                resp.raise_for_status()
                data = resp.json()
            except Exception as exc:
                logger.warning("Immuta dataSource fetch failed (offset=%d): %s", offset, exc)
                break

            items = data.get("hits", data.get("dataSources", []))
            if not items:
                break

            for ds in items:
                fqn = (
                    ds.get("qualifiedName")
                    or ds.get("name")
                    or str(ds.get("id", ""))
                )
                if not fqn:
                    continue
                if asset_fqns and not any(f in fqn for f in asset_fqns):
                    continue

                # Infer sensitivity from tags
                tags = [t.get("name", "") for t in ds.get("tags", []) if t.get("name")]
                raw_label = next(
                    (t for t in tags if t.lower() in _SENSITIVITY_MAP),
                    "unclassified",
                )
                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="table",
                    sensitivity_label=self.map_label_to_sensitivity(raw_label),
                    raw_label=raw_label,
                    data_classes=[t for t in tags if t.lower() not in _SENSITIVITY_MAP],
                    is_certified=False,
                    trust_score=None,
                    catalog_policy_ids=[],
                    fetched_at=now,
                    expires_at=now + timedelta(minutes=self.connection.sync_interval_minutes),
                ))

            total = data.get("count", data.get("total", len(items)))
            offset += len(items)
            if offset >= total or len(items) < page_size:
                break

        return classifications


# Register on import
CatalogRegistry.register(ImmutaConnector)
