"""Tests for SSO integration — WorkOS SAML/OIDC."""

from __future__ import annotations

import sys
from unittest.mock import AsyncMock, MagicMock, patch

from backend.dashboard.sso import (
    SSOConfig,
    SSOConfigInput,
    SSOUser,
    get_sso_config,
    save_sso_config,
)

_FS = "backend.shared.firestore"


def _make_sso_config_data() -> dict[str, object]:
    return {
        "customer_id": "cust-123",
        "enabled": True,
        "broker_connection_id": "conn_abc",
        "email_domain": "acme.com",
        "enforce_sso": False,
        "provider_type": "saml",
        "provider_name": "okta",
        "created_at": "2024-01-01T00:00:00Z",
        "updated_at": "2024-01-01T00:00:00Z",
        "created_by": "admin@acme.com",
        "group_mapping": {"AdminGroup": "admin"},
    }


class TestGetSSOConfig:
    async def test_get_sso_config_found(self) -> None:
        mock_doc = MagicMock()
        mock_doc.exists = True
        mock_doc.to_dict.return_value = _make_sso_config_data()

        mock_db = MagicMock()
        mock_db.collection.return_value.document.return_value.get = AsyncMock(
            return_value=mock_doc
        )

        with patch(f"{_FS}.get_db", return_value=mock_db):
            result = await get_sso_config("cust-123")

        assert result is not None
        assert isinstance(result, SSOConfig)
        assert result.customer_id == "cust-123"

    async def test_get_sso_config_not_found(self) -> None:
        mock_doc = MagicMock()
        mock_doc.exists = False

        mock_db = MagicMock()
        mock_db.collection.return_value.document.return_value.get = AsyncMock(
            return_value=mock_doc
        )

        with patch(f"{_FS}.get_db", return_value=mock_db):
            result = await get_sso_config("cust-999")

        assert result is None


class TestInitiateSSO:
    async def test_initiate_sso_login(self) -> None:
        sso_data = _make_sso_config_data()

        async def mock_stream():
            doc = MagicMock()
            doc.to_dict.return_value = sso_data
            yield doc

        mock_db = MagicMock()
        q = mock_db.collection.return_value.where
        chain = q.return_value.where.return_value.limit
        chain.return_value.stream = mock_stream

        mock_workos_mod = MagicMock()
        mock_client = MagicMock()
        mock_client.sso.get_authorization_url.return_value = (
            "https://idp.example.com/authorize?s=abc"
        )
        mock_workos_mod.WorkOSClient.return_value = mock_client

        with (
            patch(f"{_FS}.get_db", return_value=mock_db),
            patch(
                "backend.dashboard.sso._get_workos_api_key",
                return_value="sk-test",
            ),
            patch(
                "backend.dashboard.sso._get_workos_client_id",
                return_value="client-123",
            ),
            patch.dict(sys.modules, {"workos": mock_workos_mod}),
        ):
            from backend.dashboard.sso import initiate_sso_login

            url = await initiate_sso_login("acme.com", "state-123")

        assert "idp.example.com" in url


class TestHandleSSOCallback:
    async def test_handle_sso_callback(self) -> None:
        mock_profile = MagicMock()
        mock_profile.email = "user@acme.com"
        mock_profile.first_name = "Jane"
        mock_profile.last_name = "Doe"
        mock_profile.idp_id = "idp-user-1"
        mock_profile.connection_id = "conn_abc"

        mock_result = MagicMock()
        mock_result.profile = mock_profile

        mock_workos_mod = MagicMock()
        mock_client = MagicMock()
        mock_client.sso.get_profile_and_token.return_value = mock_result
        mock_workos_mod.WorkOSClient.return_value = mock_client

        with (
            patch(
                "backend.dashboard.sso._get_workos_api_key",
                return_value="sk-test",
            ),
            patch(
                "backend.dashboard.sso._get_workos_client_id",
                return_value="client-123",
            ),
            patch.dict(sys.modules, {"workos": mock_workos_mod}),
        ):
            from backend.dashboard.sso import handle_sso_callback

            user = await handle_sso_callback("auth-code-123")

        assert isinstance(user, SSOUser)
        assert user.email == "user@acme.com"
        assert user.connection_id == "conn_abc"


class TestSaveSSOConfig:
    async def test_save_sso_config_owner_only(self) -> None:
        mock_doc = MagicMock()
        mock_doc.exists = False

        mock_collection = MagicMock()
        mock_collection.document.return_value.get = AsyncMock(return_value=mock_doc)
        mock_collection.document.return_value.set = AsyncMock()
        mock_db = MagicMock()
        mock_db.collection.return_value = mock_collection

        with patch(f"{_FS}.get_db", return_value=mock_db):
            config = SSOConfigInput(
                enabled=True,
                email_domain="acme.com",
                enforce_sso=False,
                provider_type="saml",
                provider_name="okta",
            )
            await save_sso_config(
                "cust-123",
                config,
                "admin@acme.com",
                "conn_abc",
            )

        mock_collection.document.return_value.set.assert_called_once()


class TestEmailFirstLoginRouting:
    async def test_sso_enforced(self) -> None:
        sso_data = _make_sso_config_data()
        sso_data["enforce_sso"] = True
        config = SSOConfig(**sso_data)
        assert config.enforce_sso is True

    async def test_sso_optional(self) -> None:
        config = SSOConfig(**_make_sso_config_data())
        assert config.enforce_sso is False

    async def test_no_sso(self) -> None:
        import backend.dashboard.sso as sso_mod

        with patch.object(
            sso_mod,
            "get_sso_config_by_domain",
            new_callable=AsyncMock,
            return_value=None,
        ):
            result = await sso_mod.get_sso_config_by_domain("no-sso.com")
            assert result is None


class TestSSOCallbackSession:
    async def test_sso_callback_creates_session(self) -> None:
        from backend.dashboard.session import create_session

        mock_request = MagicMock()
        mock_request.client.host = "10.0.0.1"
        mock_request.headers.get.return_value = "test-agent"

        with patch(
            "backend.dashboard.session._get_encryption_key",
            return_value=b"\x00" * 32,
        ):
            _, _, session_doc, _ = create_session(
                email="user@acme.com",
                customer_id="cust-123",
                tier="team",
                request=mock_request,
            )
            session_doc["auth_method"] = "saml_sso"
            session_doc["idp_connection_id"] = "conn_abc"

        assert session_doc["auth_method"] == "saml_sso"
        assert session_doc["email"] == "user@acme.com"
