"""
Unit tests for policy versioning: registry behavior, status transitions, and version CRUD.
"""

from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock

import pytest

from litellm.proxy.policy_engine.policy_registry import (
    PolicyRegistry,
    _row_to_policy_db_response,
    get_policy_registry,
)
from litellm.types.proxy.policy_engine import (
    Policy,
    PolicyCreateRequest,
    PolicyDBResponse,
    PolicyGuardrails,
    PolicyUpdateRequest,
)


def _make_row(
    policy_id="pid-1",
    policy_name="test-policy",
    version_number=1,
    version_status="production",
    parent_version_id=None,
    is_latest=True,
    published_at=None,
    production_at=None,
    inherit=None,
    description="desc",
    guardrails_add=None,
    guardrails_remove=None,
    condition=None,
    pipeline=None,
    created_at=None,
    updated_at=None,
    created_by=None,
    updated_by=None,
):
    row = MagicMock()
    row.policy_id = policy_id
    row.policy_name = policy_name
    row.version_number = version_number
    row.version_status = version_status
    row.parent_version_id = parent_version_id
    row.is_latest = is_latest
    row.published_at = published_at
    row.production_at = production_at
    row.inherit = inherit
    row.description = description
    row.guardrails_add = guardrails_add or []
    row.guardrails_remove = guardrails_remove or []
    row.condition = condition
    row.pipeline = pipeline
    row.created_at = created_at or datetime.now(timezone.utc)
    row.updated_at = updated_at or datetime.now(timezone.utc)
    row.created_by = created_by
    row.updated_by = updated_by
    return row


class TestRowToPolicyDBResponse:
    """Test _row_to_policy_db_response includes all version fields."""

    def test_includes_version_fields(self):
        row = _make_row(
            version_number=2,
            version_status="draft",
            parent_version_id="pid-0",
            is_latest=True,
            published_at=None,
            production_at=None,
        )
        resp = _row_to_policy_db_response(row)
        assert isinstance(resp, PolicyDBResponse)
        assert resp.policy_id == "pid-1"
        assert resp.policy_name == "test-policy"
        assert resp.version_number == 2
        assert resp.version_status == "draft"
        assert resp.parent_version_id == "pid-0"
        assert resp.is_latest is True
        assert resp.published_at is None
        assert resp.production_at is None

    def test_backward_compat_missing_version_attrs(self):
        row = _make_row()
        del row.version_number
        del row.version_status
        del row.parent_version_id
        del row.is_latest
        del row.published_at
        del row.production_at
        resp = _row_to_policy_db_response(row)
        assert resp.version_number == 1
        assert resp.version_status == "production"
        assert resp.parent_version_id is None
        assert resp.is_latest is True


class TestSyncPoliciesFromDbProductionOnly:
    """Test that sync_policies_from_db only loads production versions."""

    @pytest.mark.asyncio
    async def test_get_all_policies_with_version_status_calls_find_many_with_where(
        self,
    ):
        registry = PolicyRegistry()
        prisma = MagicMock()
        prod_row = _make_row(policy_id="prod-1", version_status="production")
        prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row])

        result = await registry.get_all_policies_from_db(
            prisma, version_status="production"
        )

        assert len(result) == 1
        assert result[0].version_status == "production"
        prisma.db.litellm_policytable.find_many.assert_called_once()
        call_kw = prisma.db.litellm_policytable.find_many.call_args[1]
        assert call_kw.get("where") == {"version_status": "production"}

    @pytest.mark.asyncio
    async def test_sync_policies_from_db_only_loads_production(self):
        registry = PolicyRegistry()
        prisma = MagicMock()
        prod_row = _make_row(
            policy_id="prod-1",
            policy_name="foo",
            version_status="production",
            guardrails_add=["g1"],
        )
        prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row])

        await registry.sync_policies_from_db(prisma)

        assert registry.has_policy("foo")
        policy = registry.get_policy("foo")
        assert policy is not None
        assert policy.guardrails.add == ["g1"]
        # find_many was called with version_status=production (via get_all_policies_from_db)
        find_many_calls = prisma.db.litellm_policytable.find_many.call_args_list
        assert len(find_many_calls) >= 1
        assert find_many_calls[0][1].get("where") == {"version_status": "production"}


class TestUpdatePolicyDraftOnly:
    """Test that update_policy_in_db only allows draft versions."""

    @pytest.mark.asyncio
    async def test_update_production_raises(self):
        registry = PolicyRegistry()
        prisma = MagicMock()
        prod_row = _make_row(policy_id="pid-1", version_status="production")
        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row)

        with pytest.raises(Exception, match='Error updating policy in DB: Only draft versions can be') as exc_info:
            await registry.update_policy_in_db(
                policy_id="pid-1",
                policy_request=PolicyUpdateRequest(description="new"),
                prisma_client=prisma,
            )
        assert (
            "Only draft" in str(exc_info.value)
            or "draft" in str(exc_info.value).lower()
        )
        prisma.db.litellm_policytable.update.assert_not_called()

    @pytest.mark.asyncio
    async def test_update_draft_succeeds_and_does_not_update_registry(self):
        registry = PolicyRegistry()
        registry.add_policy("test-policy", MagicMock())  # in-memory state
        prisma = MagicMock()
        draft_row = _make_row(
            policy_id="draft-1",
            policy_name="test-policy",
            version_status="draft",
            description="old",
        )
        updated_row = _make_row(
            policy_id="draft-1",
            policy_name="test-policy",
            version_status="draft",
            description="new",
        )
        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft_row)
        prisma.db.litellm_policytable.update = AsyncMock(return_value=updated_row)

        result = await registry.update_policy_in_db(
            policy_id="draft-1",
            policy_request=PolicyUpdateRequest(description="new"),
            prisma_client=prisma,
        )

        assert result.description == "new"
        prisma.db.litellm_policytable.update.assert_called_once()
        # Registry still has old in-memory policy (drafts are not in registry; we don't add)
        assert registry.has_policy("test-policy")


class TestDeletePolicyFromDb:
    """Test delete_policy_from_db removes production from registry and returns warning."""

    @pytest.mark.asyncio
    async def test_delete_production_removes_from_registry_and_returns_warning(self):
        registry = PolicyRegistry()
        registry.add_policy("deleted-policy", MagicMock())
        prisma = MagicMock()
        prod_row = _make_row(
            policy_id="prod-1",
            policy_name="deleted-policy",
            version_status="production",
        )
        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row)
        prisma.db.litellm_policytable.delete = AsyncMock()

        result = await registry.delete_policy_from_db(
            policy_id="prod-1",
            prisma_client=prisma,
        )

        assert result["message"]
        assert "warning" in result
        assert "Production" in result["warning"] or "production" in result["warning"]
        assert not registry.has_policy("deleted-policy")

    @pytest.mark.asyncio
    async def test_delete_draft_does_not_remove_from_registry_no_warning(self):
        registry = PolicyRegistry()
        registry.add_policy("my-policy", MagicMock())
        prisma = MagicMock()
        draft_row = _make_row(
            policy_id="draft-1",
            policy_name="my-policy",
            version_status="draft",
        )
        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft_row)
        prisma.db.litellm_policytable.delete = AsyncMock()

        result = await registry.delete_policy_from_db(
            policy_id="draft-1",
            prisma_client=prisma,
        )

        assert "warning" not in result
        assert registry.has_policy("my-policy")


class TestCreateNewVersion:
    """Test create_new_version copies all fields and sets draft."""

    @pytest.mark.asyncio
    async def test_create_new_version_from_production_increments_version(self):
        registry = PolicyRegistry()
        prisma = MagicMock()
        prod = _make_row(
            policy_id="prod-1",
            policy_name="foo",
            version_number=1,
            version_status="production",
            guardrails_add=["g1"],
            description="base",
            inherit=None,
            pipeline={"mode": "pre_call", "steps": []},
        )
        # find_first for production
        prisma.db.litellm_policytable.find_first = AsyncMock(return_value=prod)
        # find_first for latest version number
        prisma.db.litellm_policytable.find_first.side_effect = [
            prod,  # production lookup
            prod,  # latest version_number lookup
        ]
        # update_many for is_latest=False
        prisma.db.litellm_policytable.update_many = AsyncMock()
        new_row = _make_row(
            policy_id="new-id",
            policy_name="foo",
            version_number=2,
            version_status="draft",
            parent_version_id="prod-1",
            is_latest=True,
            guardrails_add=["g1"],
            description="base",
            pipeline={"mode": "pre_call", "steps": []},
        )
        prisma.db.litellm_policytable.create = AsyncMock(return_value=new_row)

        result = await registry.create_new_version(
            policy_name="foo",
            prisma_client=prisma,
            source_policy_id=None,
            created_by="user",
        )

        assert result.version_number == 2
        assert result.version_status == "draft"
        assert result.parent_version_id == "prod-1"
        assert result.guardrails_add == ["g1"]
        assert result.description == "base"
        create_call = prisma.db.litellm_policytable.create.call_args[1]["data"]
        assert create_call["version_number"] == 2
        assert create_call["version_status"] == "draft"
        assert create_call["parent_version_id"] == "prod-1"
        assert create_call["guardrails_add"] == ["g1"]


class TestUpdateVersionStatus:
    """Test status transitions: valid succeed, invalid return error."""

    @pytest.mark.asyncio
    async def test_draft_to_published_sets_published_at(self):
        registry = PolicyRegistry()
        prisma = MagicMock()
        draft = _make_row(policy_id="d-1", version_status="draft")
        updated = _make_row(
            policy_id="d-1",
            version_status="published",
            published_at=datetime.now(timezone.utc),
        )
        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft)
        prisma.db.litellm_policytable.update = AsyncMock(return_value=updated)

        result = await registry.update_version_status(
            policy_id="d-1",
            new_status="published",
            prisma_client=prisma,
        )

        assert result.version_status == "published"
        update_data = prisma.db.litellm_policytable.update.call_args[1]["data"]
        assert update_data["version_status"] == "published"
        assert "published_at" in update_data

    @pytest.mark.asyncio
    async def test_draft_to_production_raises(self):
        registry = PolicyRegistry()
        prisma = MagicMock()
        draft = _make_row(policy_id="d-1", version_status="draft")
        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft)

        with pytest.raises(Exception, match='Error updating version status: Cannot promote draft') as exc_info:
            await registry.update_version_status(
                policy_id="d-1",
                new_status="production",
                prisma_client=prisma,
            )
        assert (
            "publish" in str(exc_info.value).lower()
            or "draft" in str(exc_info.value).lower()
        )

    @pytest.mark.asyncio
    async def test_published_to_production_demotes_old_and_updates_registry(self):
        registry = PolicyRegistry()
        prisma = MagicMock()
        published_row = _make_row(
            policy_id="pub-1",
            policy_name="foo",
            version_status="published",
        )
        updated_row = _make_row(
            policy_id="pub-1",
            policy_name="foo",
            version_status="production",
            production_at=datetime.now(timezone.utc),
        )
        prisma.db.litellm_policytable.find_unique = AsyncMock(
            return_value=published_row
        )
        prisma.db.litellm_policytable.update_many = AsyncMock()
        prisma.db.litellm_policytable.update = AsyncMock(return_value=updated_row)

        result = await registry.update_version_status(
            policy_id="pub-1",
            new_status="production",
            prisma_client=prisma,
        )

        assert result.version_status == "production"
        # update_many should have been called to demote current production
        assert prisma.db.litellm_policytable.update_many.called
        # Registry should have been updated with new production
        assert registry.has_policy("foo")


class TestCompareVersions:
    """Test compare_versions returns correct field diffs."""

    @pytest.mark.asyncio
    async def test_compare_versions_returns_diffs(self):
        registry = PolicyRegistry()
        prisma = MagicMock()
        a = _make_row(
            policy_id="a",
            policy_name="p",
            description="desc A",
            guardrails_add=["g1"],
        )
        b = _make_row(
            policy_id="b",
            policy_name="p",
            description="desc B",
            guardrails_add=["g1", "g2"],
        )
        prisma.db.litellm_policytable.find_unique = AsyncMock(side_effect=[a, b])

        result = await registry.compare_versions(
            policy_id_a="a",
            policy_id_b="b",
            prisma_client=prisma,
        )

        assert result.version_a.policy_id == "a"
        assert result.version_b.policy_id == "b"
        assert "description" in result.field_diffs
        assert result.field_diffs["description"]["version_a"] == "desc A"
        assert result.field_diffs["description"]["version_b"] == "desc B"
        assert "guardrails_add" in result.field_diffs


class TestResolveGuardrailsProductionOnly:
    """Test that resolve_guardrails_from_db uses only production versions."""

    @pytest.mark.asyncio
    async def test_resolve_guardrails_calls_get_all_with_production_filter(self):
        registry = PolicyRegistry()
        prisma = MagicMock()
        prod_row = _make_row(
            policy_name="base",
            version_status="production",
            guardrails_add=["g1"],
        )
        prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row])

        result = await registry.resolve_guardrails_from_db(
            policy_name="base",
            prisma_client=prisma,
        )

        assert "g1" in result
        call_kw = prisma.db.litellm_policytable.find_many.call_args[1]
        assert call_kw.get("where") == {"version_status": "production"}


class TestGetPolicyRegistrySingleton:
    """Test get_policy_registry returns same instance."""

    def test_returns_singleton(self):
        a = get_policy_registry()
        b = get_policy_registry()
        assert a is b


def _prisma_with_policy_rows(production_rows, non_production_rows=None):
    prisma = MagicMock()
    prisma.db.litellm_policytable.find_many = AsyncMock(side_effect=[production_rows, non_production_rows or []])
    return prisma


class TestConfigPoliciesPreservedAcrossDbSync:
    """Config-defined policies must survive sync_policies_from_db (regression for issue #35255)."""

    @pytest.mark.asyncio
    async def test_sync_with_empty_db_preserves_config_policies(self):
        registry = PolicyRegistry()
        registry.load_policies({"config-policy": {"description": "from config", "guardrails": {"add": ["tooling"]}}})

        await registry.sync_policies_from_db(_prisma_with_policy_rows([]))

        assert registry.has_policy("config-policy")
        policy = registry.get_policy("config-policy")
        assert policy is not None
        assert policy.guardrails.add == ["tooling"]
        assert registry.get_source("config-policy") == "config"

    @pytest.mark.asyncio
    async def test_sync_merges_db_policies_with_config_policies(self):
        registry = PolicyRegistry()
        registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
        db_row = _make_row(policy_id="db-1", policy_name="db-policy", guardrails_add=["db-guard"])

        await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))

        assert registry.has_policy("config-policy")
        assert registry.has_policy("db-policy")
        assert registry.get_source("config-policy") == "config"
        assert registry.get_source("db-policy") == "db"

    @pytest.mark.asyncio
    async def test_db_wins_on_policy_name_conflict(self):
        registry = PolicyRegistry()
        registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
        db_row = _make_row(policy_id="db-1", policy_name="shared-name", guardrails_add=["db-guard"])

        await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))

        policy = registry.get_policy("shared-name")
        assert policy is not None
        assert policy.guardrails.add == ["db-guard"]
        assert registry.get_source("shared-name") == "db"

    @pytest.mark.asyncio
    async def test_config_policy_restored_after_conflicting_db_row_deleted(self):
        registry = PolicyRegistry()
        registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
        db_row = _make_row(policy_id="db-1", policy_name="shared-name", guardrails_add=["db-guard"])

        await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))
        await registry.sync_policies_from_db(_prisma_with_policy_rows([]))

        policy = registry.get_policy("shared-name")
        assert policy is not None
        assert policy.guardrails.add == ["config-guard"]
        assert registry.get_source("shared-name") == "config"

    @pytest.mark.asyncio
    async def test_config_policy_resolves_guardrails_after_sync(self):
        from litellm.proxy.policy_engine.policy_resolver import PolicyResolver

        registry = PolicyRegistry()
        registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})

        await registry.sync_policies_from_db(_prisma_with_policy_rows([]))

        resolved = PolicyResolver.resolve_policy_guardrails(
            policy_name="config-policy",
            policies=registry.get_all_policies(),
            context=None,
        )
        assert resolved.guardrails == ["tooling"]

    @pytest.mark.asyncio
    async def test_add_policy_with_config_source_survives_sync(self):
        registry = PolicyRegistry()
        registry.add_policy(
            "late-config-policy",
            Policy(guardrails=PolicyGuardrails(add=["tooling"])),
            source="config",
        )

        await registry.sync_policies_from_db(_prisma_with_policy_rows([]))

        assert registry.has_policy("late-config-policy")
        assert registry.get_source("late-config-policy") == "config"

    @pytest.mark.asyncio
    async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self):
        registry = PolicyRegistry()
        registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})

        registry.clear()
        await registry.sync_policies_from_db(_prisma_with_policy_rows([]))

        assert not registry.has_policy("config-policy")
        assert registry.get_source("config-policy") is None


class TestRemovePolicyRestoresConfigFallback:
    """Deleting a same-named DB override must re-activate the config policy immediately, not at the next sync."""

    def test_remove_policy_restores_config_version_immediately(self):
        registry = PolicyRegistry()
        registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
        registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")

        assert registry.remove_policy("shared-name") is True

        policy = registry.get_policy("shared-name")
        assert policy is not None
        assert policy.guardrails.add == ["config-guard"]
        assert registry.get_source("shared-name") == "config"

    def test_remove_policy_without_config_fallback_removes_entirely(self):
        registry = PolicyRegistry()
        registry.add_policy("db-only", Policy(guardrails=PolicyGuardrails(add=["db-guard"])))

        assert registry.remove_policy("db-only") is True

        assert not registry.has_policy("db-only")
        assert registry.get_source("db-only") is None

    def test_remove_missing_policy_returns_false(self):
        registry = PolicyRegistry()

        assert registry.remove_policy("missing") is False

    @pytest.mark.asyncio
    async def test_delete_production_override_reactivates_config_policy_and_says_so(self):
        registry = PolicyRegistry()
        registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
        registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
        prisma = MagicMock()
        prod_row = _make_row(policy_id="prod-1", policy_name="shared-name", version_status="production")
        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row)
        prisma.db.litellm_policytable.delete = AsyncMock()

        result = await registry.delete_policy_from_db(policy_id="prod-1", prisma_client=prisma)

        assert "config" in result["warning"]
        policy = registry.get_policy("shared-name")
        assert policy is not None
        assert policy.guardrails.add == ["config-guard"]
        assert registry.get_source("shared-name") == "config"

    @pytest.mark.asyncio
    async def test_delete_all_versions_reactivates_config_policy(self):
        registry = PolicyRegistry()
        registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
        registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
        prisma = MagicMock()
        prisma.db.litellm_policytable.delete_many = AsyncMock()

        result = await registry.delete_all_versions(policy_name="shared-name", prisma_client=prisma)

        assert registry.get_source("shared-name") == "config"
        policy = registry.get_policy("shared-name")
        assert policy is not None
        assert policy.guardrails.add == ["config-guard"]
        assert "config" in result["warning"]

    async def test_delete_all_versions_without_config_twin_has_no_warning(self):
        registry = PolicyRegistry()
        registry.add_policy("db-only", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
        prisma = MagicMock()
        prisma.db.litellm_policytable.delete_many = AsyncMock()

        result = await registry.delete_all_versions(policy_name="db-only", prisma_client=prisma)

        assert registry.get_policy("db-only") is None
        assert "warning" not in result