"""
DataHub (OSS and Cloud) connector for AutoPIL.

Reads sensitivity classifications from DataHub using the GraphQL API
(recommended) or the REST GMS API (fallback for older OSS deployments).

Install:
    pip install autopil[catalog-datahub]

Required credentials (passed as 'credentials' dict when registering connection):
    token:    DataHub personal access token or service account token

Optional credentials:
    graphql_url:  Full GraphQL endpoint URL
                  (default: <base_url>/api/graphql)
    gms_url:      GMS REST endpoint URL
                  (default: <base_url>/api/v2)
    page_size:    Assets per GraphQL query page (default: 200)

Sensitivity mapping:
    DataHub glossary terms and tags mapped to AutoPIL sensitivity scale:
    Glossary term URNs containing:
        "restricted" / "highly_confidential"   → restricted
        "confidential" / "sensitive" / "pii"   → high
        "internal" / "moderate"                → medium
        "public" / "open"                      → low

Write-back:
    Enforcement events are written as custom properties on the DataHub entity
    via the GraphQL mutation or the REST ingest proposal 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 DataHub connector. "
        "Install with: pip install autopil[catalog-datahub]"
    ) 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__)

# GraphQL query: search for datasets and subTypes with sensitivity glossary terms
_SEARCH_QUERY = """
query searchAssets($input: SearchAcrossEntitiesInput!) {
  searchAcrossEntities(input: $input) {
    total
    searchResults {
      entity {
        urn
        type
        ... on Dataset {
          name
          platform { name }
          glossaryTerms {
            terms {
              term {
                urn
                properties { name }
              }
            }
          }
          tags {
            tags {
              tag {
                urn
                properties { name }
              }
            }
          }
          subTypes { typeNames }
          domain { domain { urn properties { name } } }
          properties { qualifiedName customProperties { key value } }
        }
      }
    }
  }
}
"""

_UPSERT_PROPERTIES_MUTATION = """
mutation upsertCustomProperties($input: UpsertCustomPropertiesInput!) {
  upsertCustomProperties(input: $input) {
    urn
  }
}
"""

_SENSITIVITY_KEYWORDS: list[tuple[str, str]] = [
    ("restricted", "restricted"),
    ("highly_confidential", "restricted"),
    ("highly confidential", "restricted"),
    ("confidential", "high"),
    ("sensitive", "high"),
    ("pii", "high"),
    ("phi", "high"),
    ("pci", "high"),
    ("internal", "medium"),
    ("moderate", "medium"),
    ("low", "low"),
    ("public", "low"),
    ("open", "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 DataHubConnector(CatalogConnector):
    """
    AutoPIL connector for DataHub (OSS and Cloud).

    Queries DataHub's GraphQL API to enumerate datasets and their
    glossary term / tag-based sensitivity classifications.
    Falls back to the REST GMS API for older OSS deployments.
    """

    catalog_type = "datahub"

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

    def _graphql_url(self) -> str:
        creds = _creds(self.connection)
        return creds.get("graphql_url", f"{self._base_url()}/api/graphql")

    def _gms_url(self) -> str:
        creds = _creds(self.connection)
        return creds.get("gms_url", f"{self._base_url()}/api/v2")

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

    def test_connection(self) -> tuple[bool, str]:
        try:
            creds = _creds(self.connection)
            if "token" not in creds:
                return False, "Missing required credential: 'token'"
            s = self._session()
            # Probe with a minimal GraphQL introspection query
            resp = s.post(
                self._graphql_url(),
                json={"query": "{ __typename }"},
                timeout=10,
            )
            if resp.status_code == 404:
                # Fallback: probe REST health endpoint
                resp = s.get(f"{self._gms_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"DataHub 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", 200))
        now = datetime.utcnow()
        s = self._session()

        # Try GraphQL first
        classifications = self._fetch_via_graphql(s, now, page_size, asset_fqns)
        if classifications is None:
            # Fall back to REST
            classifications = self._fetch_via_rest(s, now, page_size, asset_fqns)

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

    def write_enforcement_event(
        self,
        event: AuditEvent,
        classification: CatalogClassification | None,
    ) -> bool:
        """
        Write enforcement event to DataHub as custom properties on the entity.
        Best-effort and fail-open.
        """
        if classification is None:
            return False
        try:
            s = self._session()
            # Try GraphQL mutation first
            entity_urn = self._fqn_to_urn(classification.asset_fqn)
            variables = {
                "input": {
                    "urn": entity_urn,
                    "properties": [
                        {"key": "autopil_last_event_id",  "value": event.event_id},
                        {"key": "autopil_last_decision",  "value": event.decision.value},
                        {"key": "autopil_agent_role",     "value": event.agent_role},
                        {"key": "autopil_policy_name",    "value": event.policy_name},
                        {"key": "autopil_timestamp",      "value": event.timestamp.isoformat()},
                    ],
                }
            }
            resp = s.post(
                self._graphql_url(),
                json={"query": _UPSERT_PROPERTIES_MUTATION, "variables": variables},
                timeout=10,
            )
            if resp.status_code == 200 and "errors" not in resp.json():
                return True

            # Fallback: REST ingest proposal
            return self._write_via_rest(s, entity_urn, event)
        except Exception as exc:
            logger.debug("DataHub write-back failed: %s", exc)
            return False

    def map_label_to_sensitivity(self, raw_label: str) -> str:
        normalized = raw_label.strip().lower().replace("-", "_").replace(" ", "_")
        for keyword, level in _SENSITIVITY_KEYWORDS:
            if keyword.replace(" ", "_") in normalized:
                return level
        return super().map_label_to_sensitivity(raw_label)

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

    def _fetch_via_graphql(
        self,
        s: requests.Session,
        now: datetime,
        page_size: int,
        asset_fqns: list[str] | None,
    ) -> list[CatalogClassification] | None:
        """Search DataHub via GraphQL. Returns None if GraphQL is unavailable."""
        classifications: list[CatalogClassification] = []
        start = 0

        while True:
            variables = {
                "input": {
                    "types": ["DATASET"],
                    "query": "*",
                    "start": start,
                    "count": page_size,
                }
            }
            if asset_fqns:
                variables["input"]["query"] = " OR ".join(f'"{fqn}"' for fqn in asset_fqns)

            try:
                resp = s.post(
                    self._graphql_url(),
                    json={"query": _SEARCH_QUERY, "variables": variables},
                    timeout=30,
                )
                if resp.status_code == 404:
                    return None  # GraphQL not available
                resp.raise_for_status()
                body = resp.json()
            except Exception as exc:
                logger.warning("DataHub GraphQL search failed (start=%d): %s", start, exc)
                return None

            if "errors" in body:
                logger.warning("DataHub GraphQL errors: %s", body["errors"])
                return None

            search_data = body.get("data", {}).get("searchAcrossEntities", {})
            results = search_data.get("searchResults", [])
            total = search_data.get("total", 0)

            if not results:
                break

            for hit in results:
                entity = hit.get("entity", {})
                cls = self._entity_to_classification(entity, now)
                if cls:
                    classifications.append(cls)

            start += len(results)
            if start >= total or len(results) < page_size:
                break

        return classifications

    def _fetch_via_rest(
        self,
        s: requests.Session,
        now: datetime,
        page_size: int,
        asset_fqns: list[str] | None,
    ) -> list[CatalogClassification]:
        """Fallback REST enumeration via /entities endpoint."""
        classifications: list[CatalogClassification] = []
        start = 0

        while True:
            try:
                resp = s.get(
                    f"{self._gms_url()}/entities",
                    params={"entityName": "dataset", "start": start, "count": page_size},
                    timeout=30,
                )
                resp.raise_for_status()
                data = resp.json()
            except Exception as exc:
                logger.warning("DataHub REST entities failed (start=%d): %s", start, exc)
                break

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

            for entity in entities:
                fqn = entity.get("name") or entity.get("urn", "")
                if not fqn:
                    continue
                if asset_fqns and not any(f in fqn for f in asset_fqns):
                    continue

                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="medium",
                    raw_label="unclassified",
                    data_classes=[],
                    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("total", len(entities))
            start += len(entities)
            if start >= total or len(entities) < page_size:
                break

        return classifications

    def _entity_to_classification(
        self, entity: dict, now: datetime
    ) -> CatalogClassification | None:
        """Convert a DataHub GraphQL entity to a CatalogClassification."""
        try:
            urn = entity.get("urn", "")
            props = entity.get("properties") or {}
            fqn = props.get("qualifiedName") or entity.get("name") or urn
            if not fqn:
                return None

            # Collect all term/tag names for sensitivity inference
            term_names: list[str] = []
            for term_edge in (entity.get("glossaryTerms") or {}).get("terms", []):
                term_props = (term_edge.get("term") or {}).get("properties") or {}
                name = term_props.get("name", "")
                if name:
                    term_names.append(name)

            tag_names: list[str] = []
            for tag_edge in (entity.get("tags") or {}).get("tags", []):
                tag_props = (tag_edge.get("tag") or {}).get("properties") or {}
                name = tag_props.get("name", "")
                if name:
                    tag_names.append(name)

            all_labels = term_names + tag_names
            sensitivity, raw_label = self._infer_sensitivity(all_labels)

            # Data classes: terms/tags that aren't sensitivity-related
            data_classes = [
                n for n in all_labels
                if self.map_label_to_sensitivity(n) == "medium"  # didn't match a sensitivity keyword
                and n.lower() not in ("internal", "moderate", "medium")
            ]

            # Custom properties
            custom_props = {
                cp["key"]: cp["value"]
                for cp in (props.get("customProperties") or [])
                if cp.get("key")
            }
            if "sensitivity" in custom_props:
                raw_label = custom_props["sensitivity"]
                sensitivity = self.map_label_to_sensitivity(raw_label)

            sub_types = (entity.get("subTypes") or {}).get("typeNames", [])
            asset_type = "column" if "COLUMN" in sub_types else "table"

            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=asset_type,
                sensitivity_label=sensitivity,
                raw_label=raw_label,
                data_classes=data_classes,
                is_certified=False,
                trust_score=None,
                catalog_policy_ids=[],
                fetched_at=now,
                expires_at=now + timedelta(minutes=self.connection.sync_interval_minutes),
            )
        except Exception as exc:
            logger.debug("Skipping DataHub entity (parse error): %s", exc)
            return None

    def _infer_sensitivity(self, labels: list[str]) -> tuple[str, str]:
        """Pick the highest sensitivity from a list of term/tag names."""
        from autopil.models import SENSITIVITY_ORDER
        best_level = "medium"
        best_raw = "unclassified"
        for label in labels:
            level = self.map_label_to_sensitivity(label)
            if SENSITIVITY_ORDER.index(level) > SENSITIVITY_ORDER.index(best_level):
                best_level = level
                best_raw = label
        return best_level, best_raw

    def _fqn_to_urn(self, fqn: str) -> str:
        """Convert a dot-separated FQN to a DataHub dataset URN (best-effort)."""
        if fqn.startswith("urn:li:"):
            return fqn
        parts = fqn.split(".")
        platform = parts[0] if len(parts) > 1 else "unknown"
        return f"urn:li:dataset:(urn:li:dataPlatform:{platform},{fqn},PROD)"

    def _write_via_rest(
        self, s: requests.Session, entity_urn: str, event: AuditEvent
    ) -> bool:
        """Fallback: write enforcement event via REST ingest proposal."""
        try:
            proposal = {
                "entityType": "dataset",
                "entityUrn": entity_urn,
                "aspectName": "datasetCustomProperties",
                "aspect": {
                    "__type": "DatasetCustomProperties",
                    "customProperties": {
                        "autopil_last_event_id": event.event_id,
                        "autopil_last_decision": event.decision.value,
                        "autopil_agent_role": event.agent_role,
                        "autopil_policy_name": event.policy_name,
                        "autopil_timestamp": event.timestamp.isoformat(),
                    },
                },
            }
            resp = s.post(
                f"{self._gms_url()}/entities?action=ingestProposal",
                json={"proposal": proposal},
                timeout=10,
            )
            return resp.status_code in (200, 201, 204)
        except Exception:
            return False


# Register on import
CatalogRegistry.register(DataHubConnector)
