"""
Alation Data Catalog connector for AutoPIL.

Reads table/column trust scores and steward-defined sensitivity attributes
from the Alation REST API and maps them to AutoPIL CatalogClassification records.

Install:
    pip install autopil[catalog-alation]

Required credentials (passed as 'credentials' dict when registering connection):
    api_token:    Alation API token (generate in Settings → API Tokens)

Optional credentials:
    schema_ids:   Comma-separated list of schema IDs to restrict the scan  (e.g. "42,99")
    sensitivity_field_id:  Custom field ID used for sensitivity labels (integer)
                           If omitted, AutoPIL infers from trust score.

Sensitivity mapping:
    1. Steward-defined custom field named "Sensitivity" or "Data Classification" → base mapper
    2. Trust score (0–1): >= 0.9 → certified; score alone does not set sensitivity level
    3. If no explicit label, AutoPIL defaults to "medium"

Alation write-back:
    Enforcement events are written via the Custom Field Value API.
    The `custom_field_id` (for "AutoPIL Enforcement") must be set in credentials
    or discovered by name. 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 Alation connector. "
        "Install with: pip install autopil[catalog-alation]"
    ) 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__)

# Alation custom field names AutoPIL looks for (case-insensitive)
_SENSITIVITY_FIELD_NAMES = ("sensitivity", "data classification", "data sensitivity", "security level")


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 AlationConnector(CatalogConnector):
    """
    AutoPIL connector for Alation Data Catalog.

    Uses Alation's Integration API (v1) to enumerate tables and columns,
    read trust scores, and pull steward-defined sensitivity field values.
    """

    catalog_type = "alation"

    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({
            "TOKEN": creds["api_token"],
            "Accept": "application/json",
            "Content-Type": "application/json",
        })
        return s

    def test_connection(self) -> tuple[bool, str]:
        try:
            creds = _creds(self.connection)
            if "api_token" not in creds:
                return False, "Missing required credential: 'api_token'"
            s = self._session()
            resp = s.get(f"{self._base_url()}/integration/v1/datasource/", timeout=10)
            resp.raise_for_status()
            return True, ""
        except ImportError as exc:
            return False, str(exc)
        except Exception as exc:
            return False, f"Alation 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()
        schema_ids = [
            int(i.strip())
            for i in str(creds.get("schema_ids", "")).split(",")
            if i.strip().isdigit()
        ]
        # Discover the sensitivity custom field ID if not explicitly configured
        sensitivity_field_id = creds.get("sensitivity_field_id")
        if not sensitivity_field_id:
            sensitivity_field_id = self._find_sensitivity_field_id(s)

        classifications: list[CatalogClassification] = []

        # Fetch tables (paginated)
        tables = self._fetch_tables(s, schema_ids)
        for tbl in tables:
            fqn = self._table_fqn(tbl)
            if asset_fqns and not any(fqn_filter in fqn for fqn_filter in asset_fqns):
                continue

            trust_score = tbl.get("trust_score")
            is_certified = trust_score is not None and trust_score >= 0.9

            # Try to get sensitivity from custom field values
            sensitivity, raw_label = self._get_sensitivity(s, "table", tbl["id"], sensitivity_field_id)

            data_classes = self._get_data_classes(s, "table", tbl["id"])

            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=sensitivity,
                raw_label=raw_label,
                data_classes=data_classes,
                is_certified=is_certified,
                trust_score=float(trust_score) if trust_score is not None else None,
                catalog_policy_ids=[],
                fetched_at=now,
                expires_at=now,  # set by scheduler
            ))

        logger.info(
            "Alation fetch complete: schemas=%s tables=%d classifications=%d",
            schema_ids or "all", len(tables), len(classifications),
        )
        return classifications

    def write_enforcement_event(
        self,
        event: AuditEvent,
        classification: CatalogClassification | None,
    ) -> bool:
        """
        Write enforcement event back to Alation via the Custom Field Value API.

        Requires a custom field named "AutoPIL Enforcement" in Alation, or
        `enforcement_field_id` set in the connection credentials. Best-effort.
        """
        if classification is None:
            return False
        try:
            s    = self._session()
            creds = _creds(self.connection)

            field_id = creds.get("enforcement_field_id") or self._find_field_id(s, "AutoPIL Enforcement")
            if not field_id:
                logger.debug("Alation write-back skipped: 'AutoPIL Enforcement' custom field not found")
                return False

            # Resolve oid from FQN (table only for now)
            oid = self._find_table_oid(s, classification.asset_fqn)
            if not oid:
                return False

            payload = [{
                "field_id": int(field_id),
                "oid":      oid,
                "otype":    "table",
                "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()}/integration/v1/custom_field_value/",
                json=payload,
                timeout=10,
            )
            resp.raise_for_status()
            return True
        except Exception as exc:
            logger.debug("Alation write-back failed: %s", exc)
            return False

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

    def _fetch_tables(self, s: requests.Session, schema_ids: list[int]) -> list[dict]:
        """Paginate through Alation tables, optionally filtered by schema."""
        tables: list[dict] = []
        limit, skip = 100, 0

        params: dict = {"limit": limit, "skip": skip}
        if schema_ids:
            params["schema_id"] = schema_ids[0]  # Alation filters by single schema per request

        schemas_to_scan = schema_ids or [None]

        for schema_id in schemas_to_scan:
            skip = 0
            while True:
                p: dict = {"limit": limit, "skip": skip}
                if schema_id is not None:
                    p["schema_id"] = schema_id
                try:
                    resp = s.get(f"{self._base_url()}/integration/v1/table/", params=p, timeout=30)
                    resp.raise_for_status()
                    batch = resp.json()
                except Exception as exc:
                    logger.warning("Alation table fetch failed (schema=%s, skip=%d): %s", schema_id, skip, exc)
                    break

                if not batch:
                    break
                tables.extend(batch)
                if len(batch) < limit:
                    break
                skip += limit

        return tables

    def _table_fqn(self, tbl: dict) -> str:
        """Build a dot-separated FQN from Alation table metadata."""
        parts = []
        if tbl.get("ds_id"):
            parts.append(str(tbl["ds_id"]))
        if tbl.get("schema_name"):
            parts.append(tbl["schema_name"])
        if tbl.get("name"):
            parts.append(tbl["name"])
        return ".".join(parts) if parts else str(tbl.get("id", "unknown"))

    def _find_sensitivity_field_id(self, s: requests.Session) -> int | None:
        """Discover the custom field ID for sensitivity by name."""
        return self._find_field_id(s, *_SENSITIVITY_FIELD_NAMES)

    def _find_field_id(self, s: requests.Session, *names: str) -> int | None:
        """Return the ID of the first custom field whose name matches any of `names`."""
        try:
            resp = s.get(f"{self._base_url()}/integration/v1/custom_field/", timeout=10)
            resp.raise_for_status()
            fields = resp.json()
            for f in fields:
                if f.get("name_singular", "").lower() in {n.lower() for n in names}:
                    return f["id"]
        except Exception:
            pass
        return None

    def _get_sensitivity(
        self, s: requests.Session, otype: str, oid: int, field_id: int | None
    ) -> tuple[str, str]:
        """Get sensitivity label from a custom field value. Returns (level, raw_label)."""
        if not field_id:
            return "medium", "unclassified"
        try:
            resp = s.get(
                f"{self._base_url()}/integration/v1/custom_field_value/",
                params={"field_id": field_id, "oid": oid, "otype": otype},
                timeout=10,
            )
            resp.raise_for_status()
            values = resp.json()
            for v in values:
                raw = str(v.get("value") or "").strip()
                if raw:
                    return self.map_label_to_sensitivity(raw), raw
        except Exception:
            pass
        return "medium", "unclassified"

    def _get_data_classes(
        self, s: requests.Session, otype: str, oid: int
    ) -> list[str]:
        """Collect all non-sensitivity custom field values as data class tags."""
        try:
            resp = s.get(
                f"{self._base_url()}/integration/v1/custom_field_value/",
                params={"oid": oid, "otype": otype},
                timeout=10,
            )
            resp.raise_for_status()
            classes = []
            for v in resp.json():
                name = (v.get("field_name") or "").lower()
                if name in ("data class", "information type", "data type"):
                    val = str(v.get("value") or "").strip()
                    if val:
                        classes.append(val)
            return classes
        except Exception:
            return []

    def _find_table_oid(self, s: requests.Session, asset_fqn: str) -> int | None:
        """Find a table OID by its name segment."""
        try:
            name = asset_fqn.split(".")[-1]
            resp = s.get(
                f"{self._base_url()}/integration/v1/table/",
                params={"name": name, "limit": 1},
                timeout=10,
            )
            resp.raise_for_status()
            results = resp.json()
            return results[0]["id"] if results else None
        except Exception:
            return None


# Register on import
CatalogRegistry.register(AlationConnector)
