"""
Unit tests for Wave 2 catalog connectors: Informatica, Immuta, DataHub.

All tests run fully offline — external HTTP calls are mocked via unittest.mock.

Requires: pip install autopil[catalog]   (or autopil[dev] which includes requests)
"""

from __future__ import annotations

import json
from datetime import datetime
from unittest.mock import MagicMock, patch

import pytest

# Skip entire module if requests is not installed
pytest.importorskip("requests")

from autopil.models import AuditEvent, CatalogConnection, Decision


# ── shared helpers ────────────────────────────────────────────────────────────

def _make_connection(catalog_type: str, creds: dict, base_url: str = "https://test.example.com") -> CatalogConnection:
    now = datetime.utcnow()
    return CatalogConnection(
        connection_id="conn_test_01",
        tenant_id="test_tenant",
        catalog_type=catalog_type,
        display_name="Test",
        base_url=base_url,
        credentials_enc=json.dumps(creds),
        sync_interval_minutes=60,
        is_enabled=True,
        created_at=now,
        updated_at=now,
    )


def _make_event() -> AuditEvent:
    return AuditEvent(
        event_id="evt_001",
        session_id="sess_001",
        agent_role="test_agent",
        user_id="test_user",
        source_id="catalog.schema.table",
        query="SELECT * FROM table",
        decision=Decision.ALLOW,
        policy_name="test_policy",
        reason="unit test",
        timestamp=datetime.utcnow(),
    )


def _mock_response(status_code: int = 200, json_data: object = None) -> MagicMock:
    resp = MagicMock()
    resp.status_code = status_code
    resp.json.return_value = json_data or {}
    resp.raise_for_status.return_value = None
    if status_code >= 400:
        from requests.exceptions import HTTPError
        resp.raise_for_status.side_effect = HTTPError(response=resp)
    return resp


# ── Informatica ───────────────────────────────────────────────────────────────

class TestInformaticaConnector:

    def _connector(self, **cred_overrides):
        from autopil.catalog.connectors.informatica import InformaticaConnector
        creds = {"client_id": "test_id", "client_secret": "test_secret", **cred_overrides}
        conn = _make_connection("informatica", creds)
        return InformaticaConnector(conn)

    def _mock_login(self, mock_post, pod_url="https://dm-us.informaticacloud.com"):
        mock_post.return_value = _mock_response(200, {
            "userInfo": {"sessionId": "session_abc"},
            "products": [{"baseApiUrl": pod_url}],
        })

    def test_registered(self):
        from autopil.catalog.registry import CatalogRegistry
        from autopil.catalog.connectors.informatica import InformaticaConnector
        assert CatalogRegistry.get("informatica") is InformaticaConnector

    def test_test_connection_success(self):
        connector = self._connector()
        with patch("requests.post") as mock_post, \
             patch("requests.Session") as mock_session_cls:
            self._mock_login(mock_post)
            session = MagicMock()
            session.post.return_value = _mock_response(200, {"hits": [], "total": 0})
            mock_session_cls.return_value = session
            ok, msg = connector.test_connection()
        assert ok is True
        assert msg == ""

    def test_test_connection_missing_credential(self):
        from autopil.catalog.connectors.informatica import InformaticaConnector
        conn = _make_connection("informatica", {"client_id": "only_id"})
        connector = InformaticaConnector(conn)
        ok, msg = connector.test_connection()
        assert ok is False
        assert "client_secret" in msg

    def test_fetch_returns_list(self):
        connector = self._connector()
        assets = [
            {"qualifiedName": "org.schema.table1", "sensitivityClassification": "Confidential",
             "type": "table", "isCertified": False},
            {"qualifiedName": "org.schema.table2", "sensitivityClassification": "Highly Restricted",
             "type": "table", "isCertified": True},
        ]
        with patch("requests.post") as mock_post, \
             patch("requests.Session") as mock_session_cls:
            self._mock_login(mock_post)
            session = MagicMock()
            session.post.return_value = _mock_response(200, {"hits": assets, "total": 2})
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert len(results) == 2

    def test_sensitivity_mapping(self):
        connector = self._connector()
        assert connector.map_label_to_sensitivity("Highly Restricted") == "restricted"
        assert connector.map_label_to_sensitivity("Restricted") == "restricted"
        assert connector.map_label_to_sensitivity("Confidential") == "high"
        assert connector.map_label_to_sensitivity("Sensitive") == "high"
        assert connector.map_label_to_sensitivity("Internal Use Only") == "medium"
        assert connector.map_label_to_sensitivity("Internal") == "medium"
        assert connector.map_label_to_sensitivity("Public") == "low"

    def test_sensitivity_label_set_correctly(self):
        connector = self._connector()
        assets = [{"qualifiedName": "a.b.c", "sensitivityClassification": "Confidential", "type": "table"}]
        with patch("requests.post") as mock_post, \
             patch("requests.Session") as mock_session_cls:
            self._mock_login(mock_post)
            session = MagicMock()
            session.post.return_value = _mock_response(200, {"hits": assets, "total": 1})
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert results[0].sensitivity_label == "high"
        assert results[0].raw_label == "Confidential"

    def test_fetch_empty_catalog(self):
        connector = self._connector()
        with patch("requests.post") as mock_post, \
             patch("requests.Session") as mock_session_cls:
            self._mock_login(mock_post)
            session = MagicMock()
            session.post.return_value = _mock_response(200, {"hits": [], "total": 0})
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert results == []

    def test_write_back_returns_bool(self):
        connector = self._connector()
        from autopil.models import CatalogClassification
        now = datetime.utcnow()
        cls = CatalogClassification(
            classification_id="cls_01", connection_id="conn_test_01",
            tenant_id="test_tenant", asset_fqn="org.schema.table1",
            asset_type="table", sensitivity_label="high", raw_label="Confidential",
            data_classes=[], is_certified=False, trust_score=None,
            catalog_policy_ids=[], fetched_at=now, expires_at=now,
        )
        event = _make_event()
        with patch("requests.post") as mock_post, \
             patch("requests.Session") as mock_session_cls:
            self._mock_login(mock_post)
            session = MagicMock()
            session.patch.return_value = _mock_response(200)
            mock_session_cls.return_value = session
            result = connector.write_enforcement_event(event, cls)
        assert isinstance(result, bool)

    def test_write_back_none_classification(self):
        connector = self._connector()
        result = connector.write_enforcement_event(_make_event(), None)
        assert result is False

    def test_auth_failure_returns_empty(self):
        connector = self._connector()
        with patch("requests.post") as mock_post:
            mock_post.side_effect = Exception("network error")
            results = connector.fetch_classifications()
        assert results == []


# ── Immuta ────────────────────────────────────────────────────────────────────

class TestImmutaConnector:

    def _connector(self, **cred_overrides):
        from autopil.catalog.connectors.immuta import ImmutaConnector
        creds = {"api_key": "test_api_key", **cred_overrides}
        conn = _make_connection("immuta", creds, base_url="https://immuta.example.com")
        return ImmutaConnector(conn)

    def test_registered(self):
        from autopil.catalog.registry import CatalogRegistry
        from autopil.catalog.connectors.immuta import ImmutaConnector
        assert CatalogRegistry.get("immuta") is ImmutaConnector

    def test_test_connection_success(self):
        connector = self._connector()
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.get.return_value = _mock_response(200, [])
            mock_session_cls.return_value = session
            ok, msg = connector.test_connection()
        assert ok is True

    def test_test_connection_missing_credential(self):
        from autopil.catalog.connectors.immuta import ImmutaConnector
        conn = _make_connection("immuta", {})
        connector = ImmutaConnector(conn)
        ok, msg = connector.test_connection()
        assert ok is False
        assert "api_key" in msg

    def test_fetch_via_sensitivity_api(self):
        connector = self._connector()
        items = [
            {"dataSourceName": "db.schema.table1", "sensitivityLevel": "Sensitive", "dataClasses": ["PII"]},
            {"dataSourceName": "db.schema.table2", "sensitivityLevel": "Public", "dataClasses": []},
        ]
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.get.return_value = _mock_response(200, items)
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert len(results) == 2
        assert results[0].sensitivity_label == "high"
        assert results[1].sensitivity_label == "low"

    def test_fetch_falls_back_to_datasources(self):
        connector = self._connector()
        datasources = {
            "hits": [
                {"qualifiedName": "ds.schema.tbl", "tags": [{"name": "Confidential"}, {"name": "PII"}]},
            ],
            "count": 1,
        }
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            # First call (sensitivity API) returns 404, second (datasource) returns data
            session.get.side_effect = [
                _mock_response(404),
                _mock_response(200, datasources),
                _mock_response(200, {"hits": [], "count": 0}),  # next page empty
            ]
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert len(results) == 1
        assert results[0].sensitivity_label == "high"  # "Confidential" → high

    def test_sensitivity_mapping(self):
        connector = self._connector()
        assert connector.map_label_to_sensitivity("Highly Sensitive") == "restricted"
        assert connector.map_label_to_sensitivity("Critical") == "restricted"
        assert connector.map_label_to_sensitivity("Sensitive") == "high"
        assert connector.map_label_to_sensitivity("Confidential") == "high"
        assert connector.map_label_to_sensitivity("Moderate") == "medium"
        assert connector.map_label_to_sensitivity("Internal") == "medium"
        assert connector.map_label_to_sensitivity("Low") == "low"
        assert connector.map_label_to_sensitivity("Public") == "low"

    def test_write_back_returns_bool(self):
        connector = self._connector()
        from autopil.models import CatalogClassification
        now = datetime.utcnow()
        cls = CatalogClassification(
            classification_id="cls_01", connection_id="conn_test_01",
            tenant_id="test_tenant", asset_fqn="ds.schema.table",
            asset_type="table", sensitivity_label="high", raw_label="Sensitive",
            data_classes=["PII"], is_certified=False, trust_score=None,
            catalog_policy_ids=[], fetched_at=now, expires_at=now,
        )
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.post.return_value = _mock_response(201)
            mock_session_cls.return_value = session
            result = connector.write_enforcement_event(_make_event(), cls)
        assert isinstance(result, bool)

    def test_write_back_none_classification(self):
        connector = self._connector()
        assert connector.write_enforcement_event(_make_event(), None) is False

    def test_data_classes_populated(self):
        connector = self._connector()
        items = [{"dataSourceName": "a.b.c", "sensitivityLevel": "Sensitive", "dataClasses": ["PII", "PHI"]}]
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.get.return_value = _mock_response(200, items)
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert "PII" in results[0].data_classes
        assert "PHI" in results[0].data_classes


# ── DataHub ───────────────────────────────────────────────────────────────────

class TestDataHubConnector:

    def _connector(self, **cred_overrides):
        from autopil.catalog.connectors.datahub import DataHubConnector
        creds = {"token": "test_token", **cred_overrides}
        conn = _make_connection("datahub", creds, base_url="https://datahub.example.com")
        return DataHubConnector(conn)

    def _graphql_hit(self, fqn: str, terms: list[str], tags: list[str] | None = None) -> dict:
        """Build a minimal DataHub GraphQL search hit."""
        return {
            "entity": {
                "urn": f"urn:li:dataset:(urn:li:dataPlatform:hive,{fqn},PROD)",
                "type": "DATASET",
                "name": fqn.split(".")[-1],
                "properties": {"qualifiedName": fqn, "customProperties": []},
                "glossaryTerms": {
                    "terms": [
                        {"term": {"urn": f"urn:li:glossaryTerm:{t}", "properties": {"name": t}}}
                        for t in terms
                    ]
                },
                "tags": {
                    "tags": [
                        {"tag": {"urn": f"urn:li:tag:{t}", "properties": {"name": t}}}
                        for t in (tags or [])
                    ]
                },
                "subTypes": {"typeNames": []},
            }
        }

    def test_registered(self):
        from autopil.catalog.registry import CatalogRegistry
        from autopil.catalog.connectors.datahub import DataHubConnector
        assert CatalogRegistry.get("datahub") is DataHubConnector

    def test_test_connection_success(self):
        connector = self._connector()
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.post.return_value = _mock_response(200, {"data": {"__typename": "Query"}})
            mock_session_cls.return_value = session
            ok, msg = connector.test_connection()
        assert ok is True

    def test_test_connection_missing_token(self):
        from autopil.catalog.connectors.datahub import DataHubConnector
        conn = _make_connection("datahub", {})
        connector = DataHubConnector(conn)
        ok, msg = connector.test_connection()
        assert ok is False
        assert "token" in msg

    def test_fetch_via_graphql(self):
        connector = self._connector()
        hits = [
            self._graphql_hit("catalog.schema.table1", ["Confidential"]),
            self._graphql_hit("catalog.schema.table2", ["Public"]),
        ]
        graphql_resp = {"data": {"searchAcrossEntities": {"total": 2, "searchResults": hits}}}
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.post.return_value = _mock_response(200, graphql_resp)
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert len(results) == 2

    def test_sensitivity_from_glossary_term(self):
        connector = self._connector()
        hits = [self._graphql_hit("a.b.c", ["Confidential"])]
        graphql_resp = {"data": {"searchAcrossEntities": {"total": 1, "searchResults": hits}}}
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.post.return_value = _mock_response(200, graphql_resp)
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert results[0].sensitivity_label == "high"

    def test_sensitivity_from_tag(self):
        connector = self._connector()
        hits = [self._graphql_hit("a.b.c", [], tags=["restricted"])]
        graphql_resp = {"data": {"searchAcrossEntities": {"total": 1, "searchResults": hits}}}
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.post.return_value = _mock_response(200, graphql_resp)
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert results[0].sensitivity_label == "restricted"

    def test_most_restrictive_term_wins(self):
        connector = self._connector()
        # Entity has both "Public" and "Confidential" terms — should pick highest
        hits = [self._graphql_hit("a.b.c", ["Public", "Confidential"])]
        graphql_resp = {"data": {"searchAcrossEntities": {"total": 1, "searchResults": hits}}}
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.post.return_value = _mock_response(200, graphql_resp)
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert results[0].sensitivity_label == "high"

    def test_sensitivity_mapping(self):
        connector = self._connector()
        assert connector.map_label_to_sensitivity("restricted") == "restricted"
        assert connector.map_label_to_sensitivity("Highly Confidential") == "restricted"
        assert connector.map_label_to_sensitivity("confidential") == "high"
        assert connector.map_label_to_sensitivity("PII") == "high"
        assert connector.map_label_to_sensitivity("internal") == "medium"
        assert connector.map_label_to_sensitivity("public") == "low"
        assert connector.map_label_to_sensitivity("open") == "low"

    def test_graphql_falls_back_to_rest(self):
        connector = self._connector()
        rest_resp = {
            "entities": [{"name": "schema.table", "urn": "urn:li:dataset:hive,schema.table,PROD"}],
            "total": 1,
        }
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            # GraphQL returns 404, REST returns data
            session.post.side_effect = [
                _mock_response(404),        # test_connection probe → 404
            ]
            session.get.side_effect = [
                _mock_response(200, rest_resp),                      # /entities page 1
                _mock_response(200, {"entities": [], "total": 0}),   # /entities page 2 (empty → stop)
            ]
            mock_session_cls.return_value = session
            # Directly test _fetch_via_rest
            from datetime import datetime
            results = connector._fetch_via_rest(session, datetime.utcnow(), 200, None)
        assert len(results) == 1

    def test_write_back_graphql_success(self):
        connector = self._connector()
        from autopil.models import CatalogClassification
        now = datetime.utcnow()
        cls = CatalogClassification(
            classification_id="cls_01", connection_id="conn_test_01",
            tenant_id="test_tenant", asset_fqn="catalog.schema.table",
            asset_type="table", sensitivity_label="high", raw_label="Confidential",
            data_classes=[], is_certified=False, trust_score=None,
            catalog_policy_ids=[], fetched_at=now, expires_at=now,
        )
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.post.return_value = _mock_response(
                200, {"data": {"upsertCustomProperties": {"urn": "urn:li:dataset:test"}}}
            )
            mock_session_cls.return_value = session
            result = connector.write_enforcement_event(_make_event(), cls)
        assert isinstance(result, bool)

    def test_write_back_none_classification(self):
        connector = self._connector()
        assert connector.write_enforcement_event(_make_event(), None) is False

    def test_fqn_to_urn_passthrough(self):
        connector = self._connector()
        urn = "urn:li:dataset:(urn:li:dataPlatform:hive,a.b.c,PROD)"
        assert connector._fqn_to_urn(urn) == urn

    def test_fqn_to_urn_conversion(self):
        connector = self._connector()
        urn = connector._fqn_to_urn("hive.schema.table")
        assert urn.startswith("urn:li:dataset:")
        assert "hive" in urn

    def test_fetch_empty_returns_list(self):
        connector = self._connector()
        graphql_resp = {"data": {"searchAcrossEntities": {"total": 0, "searchResults": []}}}
        with patch("requests.Session") as mock_session_cls:
            session = MagicMock()
            session.post.return_value = _mock_response(200, graphql_resp)
            mock_session_cls.return_value = session
            results = connector.fetch_classifications()
        assert results == []
