"""
Microsoft Purview connector for AutoPIL.

Reads sensitivity labels and classifications from Microsoft Purview
via the Apache Atlas REST API and maps them to AutoPIL CatalogClassification records.

Install:
    pip install autopil[catalog-purview]

Required credentials (passed as 'credentials' dict when registering connection):
    tenant_id:     Azure AD tenant ID
    client_id:     Service principal / app registration client ID
    client_secret: Service principal client secret
    account_name:  Purview account name (the subdomain of purview.azure.com)

Optional credentials:
    collection_name: Restrict scan to a specific Purview collection

Sensitivity mapping:
    AutoPIL reads Microsoft Information Protection (MIP) sensitivity labels:
        Highly Confidential → restricted
        Confidential        → high
        General             → medium
        Public              → low
    It also reads Apache Atlas classifications (tags) and maps them via the base mapper.

Purview write-back:
    Enforcement events are written as custom attributes on the Atlas entity.
    The attribute "autopil_enforcement" is added to the entity's attributes.
    Write-back is best-effort and fail-open.
"""

from __future__ import annotations

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

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

try:
    import msal
except ImportError as exc:
    raise ImportError(
        "msal is required for the Purview connector. "
        "Install with: pip install autopil[catalog-purview]"
    ) 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__)

# Azure Purview resource scope for MSAL token acquisition
_PURVIEW_SCOPE = "https://purview.azure.com/.default"

# Microsoft Information Protection label → AutoPIL level
_MIP_LABEL_MAP: dict[str, str] = {
    "highly confidential": "restricted",
    "confidential":        "high",
    "general":             "medium",
    "public":              "low",
    "personal":            "medium",
    "non-business":        "low",
}

# Atlas classification names → AutoPIL level overrides
_ATLAS_CLASS_MAP: dict[str, str] = {
    "microsoft.financial":      "high",
    "microsoft.government":     "high",
    "microsoft.personal":       "high",
    "microsoft.privacy":        "high",
    "microsoft.security":       "restricted",
    "microsoft.enterprise":     "medium",
}


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 PurviewConnector(CatalogConnector):
    """
    AutoPIL connector for Microsoft Purview.

    Authenticates via MSAL service principal (client credentials flow),
    then calls the Purview Apache Atlas API to enumerate entities and read
    MIP sensitivity labels and Atlas classifications.
    """

    catalog_type = "purview"

    def _creds(self) -> dict:
        return _creds(self.connection)

    def _atlas_endpoint(self) -> str:
        creds = self._creds()
        account = creds["account_name"]
        return f"https://{account}.purview.azure.com/catalog/api"

    def _get_token(self) -> str:
        """Acquire an MSAL access token using client credentials."""
        creds = self._creds()
        app = msal.ConfidentialClientApplication(
            client_id=creds["client_id"],
            client_credential=creds["client_secret"],
            authority=f"https://login.microsoftonline.com/{creds['tenant_id']}",
        )
        result = app.acquire_token_for_client(scopes=[_PURVIEW_SCOPE])
        if "access_token" not in result:
            error = result.get("error_description", result.get("error", "Unknown MSAL error"))
            raise RuntimeError(f"MSAL token acquisition failed: {error}")
        return result["access_token"]

    def _session(self) -> requests.Session:
        token = self._get_token()
        s = requests.Session()
        s.headers.update({
            "Authorization": f"Bearer {token}",
            "Accept":        "application/json",
            "Content-Type":  "application/json",
        })
        return s

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

            s = self._session()
            resp = s.get(
                f"{self._atlas_endpoint()}/atlas/v2/types/typedefs",
                params={"type": "classification"},
                timeout=15,
            )
            resp.raise_for_status()
            return True, ""
        except ImportError as exc:
            return False, str(exc)
        except Exception as exc:
            return False, f"Purview connection failed: {exc}"

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

        classifications: list[CatalogClassification] = []

        # If specific FQNs requested, look them up individually
        if asset_fqns:
            for fqn in asset_fqns:
                cls = self._fetch_by_fqn(s, fqn, now)
                if cls:
                    classifications.append(cls)
            return classifications

        # Full scan via Atlas search
        entities = self._search_entities(s, collection)
        for entity in entities:
            cls = self._entity_to_classification(s, entity, now)
            if cls:
                classifications.append(cls)

        logger.info(
            "Purview fetch complete: collection=%s entities=%d classifications=%d",
            collection or "all", len(entities), len(classifications),
        )
        return classifications

    def write_enforcement_event(
        self,
        event: AuditEvent,
        classification: CatalogClassification | None,
    ) -> bool:
        """
        Write an AutoPIL enforcement record back to the Purview entity as a custom attribute.

        Looks up the entity GUID by qualified name, then patches the entity with
        an `autopil_enforcement` attribute. Best-effort — returns False on any failure.
        """
        if classification is None:
            return False
        try:
            s = self._session()
            guid = self._find_entity_guid(s, classification.asset_fqn)
            if not guid:
                return False

            enforcement_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(),
            })

            # Patch entity attributes
            payload = {
                "entity": {
                    "guid": guid,
                    "attributes": {
                        "autopil_last_event_id": event.event_id,
                        "autopil_last_decision": event.decision.value,
                        "autopil_enforcement":   enforcement_value,
                    },
                }
            }
            resp = s.post(
                f"{self._atlas_endpoint()}/atlas/v2/entity",
                json=payload,
                timeout=15,
            )
            resp.raise_for_status()
            return True
        except Exception as exc:
            logger.debug("Purview write-back failed: %s", exc)
            return False

    def map_label_to_sensitivity(self, raw_label: str) -> str:
        normalised = raw_label.lower().strip().replace("-", " ").replace("_", " ")
        # MIP labels take precedence
        if normalised in _MIP_LABEL_MAP:
            return _MIP_LABEL_MAP[normalised]
        # Atlas classification names
        for pattern, level in _ATLAS_CLASS_MAP.items():
            if pattern in normalised:
                return level
        return super().map_label_to_sensitivity(raw_label)

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

    def _search_entities(self, s: requests.Session, collection: str | None) -> list[dict]:
        """Search Purview entities using the Atlas quick search API (paginated)."""
        entities: list[dict] = []
        offset, limit = 0, 100

        while True:
            payload: dict = {
                "limit":  limit,
                "offset": offset,
                "typeName": "DataSet",   # broad type covering tables, files, etc.
            }
            if collection:
                payload["attributes"] = {"collectionId": collection}

            try:
                resp = s.post(
                    f"{self._atlas_endpoint()}/atlas/v2/search/advanced",
                    json=payload,
                    timeout=30,
                )
                resp.raise_for_status()
                data = resp.json()
            except Exception as exc:
                logger.warning("Purview entity search failed at offset %d: %s", offset, exc)
                break

            batch = data.get("value", [])
            if not batch:
                break
            entities.extend(batch)
            if len(batch) < limit:
                break
            offset += limit

        return entities

    def _fetch_by_fqn(
        self, s: requests.Session, fqn: str, now: datetime
    ) -> CatalogClassification | None:
        """Fetch a single entity by its qualified name."""
        try:
            resp = s.get(
                f"{self._atlas_endpoint()}/atlas/v2/entity/uniqueAttribute/type/DataSet",
                params={"attr:qualifiedName": fqn},
                timeout=15,
            )
            if resp.status_code == 404:
                return None
            resp.raise_for_status()
            entity = resp.json().get("entity", {})
            return self._entity_to_classification(s, entity, now)
        except Exception as exc:
            logger.debug("Purview FQN fetch failed for %s: %s", fqn, exc)
            return None

    def _entity_to_classification(
        self, s: requests.Session, entity: dict, now: datetime
    ) -> CatalogClassification | None:
        """Convert a Purview Atlas entity dict to a CatalogClassification."""
        guid = entity.get("guid") or entity.get("id", "")
        attrs = entity.get("attributes") or {}
        fqn   = attrs.get("qualifiedName") or attrs.get("name") or guid

        if not fqn:
            return None

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

        return 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_entity_type(entity),
            sensitivity_label=sensitivity,
            raw_label=raw_label,
            data_classes=data_classes,
            is_certified=None,
            trust_score=None,
            catalog_policy_ids=[],
            fetched_at=now,
            expires_at=now,  # set by scheduler
        )

    def _extract_sensitivity(self, entity: dict) -> tuple[str, str, list[str]]:
        """
        Extract sensitivity from MIP labels first, then Atlas classifications.
        Returns (sensitivity_level, raw_label, data_classes).
        """
        attrs = entity.get("attributes") or {}
        data_classes: list[str] = []

        # 1. MIP sensitivity label (Purview-native)
        mip_label = attrs.get("label") or attrs.get("sensitivityLabel") or ""
        if mip_label:
            return self.map_label_to_sensitivity(mip_label), mip_label, data_classes

        # 2. Atlas classifications (tags applied to the entity)
        classifications = entity.get("classifications") or []
        highest_level = "medium"
        raw = "unclassified"

        for cls_dict in classifications:
            type_name = cls_dict.get("typeName", "")
            level = self.map_label_to_sensitivity(type_name)
            from autopil.models import SENSITIVITY_ORDER
            if SENSITIVITY_ORDER.index(level) > SENSITIVITY_ORDER.index(highest_level):
                highest_level = level
                raw = type_name
            # Collect as data class tags
            data_classes.append(type_name)

        return highest_level, raw, data_classes

    def _map_entity_type(self, entity: dict) -> str:
        type_name = (entity.get("typeName") or "").lower()
        if "table" in type_name or "dataset" in type_name:
            return "table"
        if "column" in type_name or "field" in type_name:
            return "column"
        if "schema" in type_name:
            return "schema"
        return "asset"

    def _find_entity_guid(self, s: requests.Session, fqn: str) -> str | None:
        """Look up an entity GUID by qualified name."""
        try:
            resp = s.get(
                f"{self._atlas_endpoint()}/atlas/v2/entity/uniqueAttribute/type/DataSet",
                params={"attr:qualifiedName": fqn},
                timeout=10,
            )
            if resp.ok:
                return resp.json().get("entity", {}).get("guid")
        except Exception:
            pass
        return None


# Register on import
CatalogRegistry.register(PurviewConnector)
