from types import SimpleNamespace
from unittest.mock import AsyncMock

import httpx
import pytest

from api.services import quota_service
from api.services.configuration.registry import ServiceProviders
from api.services.managed_model_services import MPS_CORRELATION_ID_CONTEXT_KEY
from api.services.quota_service import QuotaCheckResult

_UNSET = object()


def _dograh_config(
    api_key: str = "mps_sk_12345678",
    *,
    managed_service_version: int = 1,
):
    return SimpleNamespace(
        managed_service_version=managed_service_version,
        llm=SimpleNamespace(provider=ServiceProviders.DOGRAH, api_key=api_key),
        stt=None,
        tts=None,
        embeddings=None,
    )


def _byok_config():
    return SimpleNamespace(
        managed_service_version=2,
        llm=SimpleNamespace(provider="openai", api_key="sk-openai"),
        stt=None,
        tts=None,
        embeddings=None,
    )


def _workflow():
    return SimpleNamespace(
        id=7,
        user_id=123,
        organization_id=42,
        workflow_configurations={"model_overrides": {}},
    )


def _workflow_owner():
    return SimpleNamespace(
        id=123,
        provider_id="provider-123",
    )


def _actor():
    return SimpleNamespace(
        id=456,
        provider_id="actor-456",
        selected_organization_id=42,
    )


def _patch_workflow_context(monkeypatch, *, workflow=_UNSET, owner=None):
    workflow_value = _workflow() if workflow is _UNSET else workflow
    monkeypatch.setattr(
        quota_service.db_client,
        "get_workflow",
        AsyncMock(return_value=workflow_value),
    )
    monkeypatch.setattr(
        quota_service.db_client,
        "get_workflow_by_id",
        AsyncMock(side_effect=AssertionError("quota must not use unscoped workflow")),
    )
    monkeypatch.setattr(
        quota_service.db_client,
        "get_user_by_id",
        AsyncMock(return_value=owner or _workflow_owner()),
    )


@pytest.mark.asyncio
async def test_authorize_workflow_run_uses_workflow_org_for_hosted_v2(
    monkeypatch,
):
    get_config = AsyncMock(return_value=_dograh_config())
    authorize = AsyncMock(
        return_value={
            "allowed": True,
            "billing_mode": "v2",
            "remaining_credits": "25.0000",
        }
    )
    check_usage = AsyncMock()

    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service,
        "get_effective_ai_model_configuration_for_workflow",
        get_config,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "authorize_workflow_run_start",
        authorize,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "check_service_key_usage",
        check_usage,
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
    )

    assert result.has_quota is True
    quota_service.db_client.get_workflow.assert_awaited_once_with(7, organization_id=42)
    get_config.assert_awaited_once_with(
        organization_id=42,
        workflow_configurations={"model_overrides": {}},
    )
    authorize.assert_awaited_once_with(
        organization_id=42,
        workflow_run_id=None,
        service_key=None,
        require_correlation_id=False,
        minimum_credits=quota_service.MINIMUM_DOGRAH_CREDITS_FOR_CALL,
        created_by="provider-123",
        metadata={"dograh_user_id": "123", "workflow_id": 7},
    )
    check_usage.assert_not_awaited()


@pytest.mark.asyncio
async def test_authorize_workflow_run_v2_insufficient_credits_prompts_billing(
    monkeypatch,
):
    get_config = AsyncMock(return_value=_byok_config())
    authorize = AsyncMock(
        return_value={
            "allowed": False,
            "billing_mode": "v2",
            "remaining_credits": "0.0000",
            "error": "insufficient_credits",
        }
    )
    check_usage = AsyncMock()

    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service,
        "get_effective_ai_model_configuration_for_workflow",
        get_config,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "authorize_workflow_run_start",
        authorize,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "check_service_key_usage",
        check_usage,
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
    )

    assert result.has_quota is False
    assert result.error_code == "insufficient_credits"
    assert "/billing" in result.error_message
    assert "founders@dograh.com" not in result.error_message
    authorize.assert_awaited_once()
    check_usage.assert_not_awaited()


@pytest.mark.asyncio
async def test_authorize_workflow_run_oss_exhausted_key_blocks_run(
    monkeypatch,
):
    api_key = "mps_sk_12345678"
    get_config = AsyncMock(return_value=_dograh_config(api_key))
    check_usage = AsyncMock(
        return_value={"total_credits_used": 500.0, "remaining_credits": 0.0}
    )

    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service,
        "get_effective_ai_model_configuration_for_workflow",
        get_config,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "check_service_key_usage",
        check_usage,
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
    )

    assert result.has_quota is False
    assert result.error_code == "quota_exceeded"
    assert "app.dograh.com" in result.error_message
    assert "/billing" not in result.error_message
    check_usage.assert_awaited_once_with(api_key)


@pytest.mark.asyncio
async def test_authorize_workflow_run_managed_v2_stores_hosted_correlation(
    monkeypatch,
):
    api_key = "mps_sk_12345678"
    workflow_run = SimpleNamespace(initial_context={"existing": "value"})
    get_config = AsyncMock(
        return_value=_dograh_config(api_key, managed_service_version=2)
    )
    authorize = AsyncMock(
        return_value={
            "allowed": True,
            "billing_mode": "v2",
            "remaining_credits": "25.0000",
            "correlation_id": "mps-corr-123",
        }
    )
    update_workflow_run = AsyncMock()

    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service.db_client,
        "get_workflow_run_by_id",
        AsyncMock(return_value=workflow_run),
    )
    monkeypatch.setattr(
        quota_service.db_client,
        "update_workflow_run",
        update_workflow_run,
    )
    monkeypatch.setattr(
        quota_service,
        "get_effective_ai_model_configuration_for_workflow",
        get_config,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "authorize_workflow_run_start",
        authorize,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "check_service_key_usage",
        AsyncMock(),
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
        workflow_run_id=88,
    )

    assert result.has_quota is True
    authorize.assert_awaited_once_with(
        organization_id=42,
        workflow_run_id=88,
        service_key=api_key,
        require_correlation_id=True,
        minimum_credits=quota_service.MINIMUM_DOGRAH_CREDITS_FOR_CALL,
        created_by="provider-123",
        metadata={"dograh_user_id": "123", "workflow_id": 7},
    )
    update_workflow_run.assert_awaited_once_with(
        88,
        initial_context={
            "existing": "value",
            MPS_CORRELATION_ID_CONTEXT_KEY: "mps-corr-123",
        },
    )


@pytest.mark.asyncio
async def test_authorize_workflow_run_service_token_from_wrong_org_prompts_new_token(
    monkeypatch,
):
    api_key = "mps_sk_12345678"
    get_config = AsyncMock(
        return_value=_dograh_config(api_key, managed_service_version=2)
    )
    request = httpx.Request(
        "POST",
        "http://localhost:8004/api/v1/billing/accounts/42/run-authorization",
    )
    response = httpx.Response(
        403,
        json={"detail": "Service key organization mismatch"},
        request=request,
    )
    authorize = AsyncMock(
        side_effect=httpx.HTTPStatusError(
            "Failed to authorize MPS workflow run start",
            request=request,
            response=response,
        )
    )

    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service,
        "get_effective_ai_model_configuration_for_workflow",
        get_config,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "authorize_workflow_run_start",
        authorize,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "check_service_key_usage",
        AsyncMock(),
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
        workflow_run_id=88,
    )

    assert result.has_quota is False
    assert result.error_code == "service_key_org_mismatch"
    assert result.error_message == quota_service.SERVICE_TOKEN_ORG_MISMATCH_MESSAGE
    assert "new service token from the Developers tab" in result.error_message
    authorize.assert_awaited_once_with(
        organization_id=42,
        workflow_run_id=88,
        service_key=api_key,
        require_correlation_id=True,
        minimum_credits=quota_service.MINIMUM_DOGRAH_CREDITS_FOR_CALL,
        created_by="provider-123",
        metadata={"dograh_user_id": "123", "workflow_id": 7},
    )


@pytest.mark.asyncio
async def test_authorize_workflow_run_oss_uses_key_paths_not_workflow_org(
    monkeypatch,
):
    api_key = "mps_sk_12345678"
    workflow_run = SimpleNamespace(initial_context={})
    get_config = AsyncMock(
        return_value=_dograh_config(api_key, managed_service_version=2)
    )
    hosted_authorize = AsyncMock()
    check_usage = AsyncMock(
        return_value={"total_credits_used": 1.0, "remaining_credits": 499.0}
    )
    create_correlation = AsyncMock(return_value={"correlation_id": "oss-corr-123"})
    update_workflow_run = AsyncMock()

    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service.db_client,
        "get_workflow_run_by_id",
        AsyncMock(return_value=workflow_run),
    )
    monkeypatch.setattr(
        quota_service.db_client,
        "update_workflow_run",
        update_workflow_run,
    )
    monkeypatch.setattr(
        quota_service,
        "get_effective_ai_model_configuration_for_workflow",
        get_config,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "authorize_workflow_run_start",
        hosted_authorize,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "check_service_key_usage",
        check_usage,
    )
    monkeypatch.setattr(
        quota_service.mps_service_key_client,
        "create_correlation_id",
        create_correlation,
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
        workflow_run_id=88,
    )

    assert result.has_quota is True
    hosted_authorize.assert_not_awaited()
    check_usage.assert_awaited_once_with(api_key)
    create_correlation.assert_awaited_once_with(
        service_key=api_key,
        workflow_run_id=88,
    )
    update_workflow_run.assert_awaited_once_with(
        88,
        initial_context={MPS_CORRELATION_ID_CONTEXT_KEY: "oss-corr-123"},
    )


@pytest.mark.asyncio
async def test_authorize_workflow_run_rejects_actor_not_a_member(monkeypatch):
    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service.db_client,
        "is_user_member_of_organization",
        AsyncMock(return_value=False),
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
        actor_user=SimpleNamespace(id=456, selected_organization_id=999),
    )

    assert result.has_quota is False
    assert result.error_code == "workflow_not_found"


@pytest.mark.asyncio
async def test_authorize_workflow_run_membership_lookup_error_fails_closed(monkeypatch):
    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service.db_client,
        "is_user_member_of_organization",
        AsyncMock(side_effect=RuntimeError("database unavailable")),
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
        actor_user=SimpleNamespace(id=456, selected_organization_id=42),
    )

    assert result.has_quota is False
    assert result.error_code == "workflow_not_found"
    quota_service.db_client.get_user_by_id.assert_not_awaited()


@pytest.mark.asyncio
async def test_authorize_workflow_run_allows_invited_member(monkeypatch):
    """User invited to an org can start workflows belonging to that org."""
    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch)
    monkeypatch.setattr(
        quota_service.db_client,
        "is_user_member_of_organization",
        AsyncMock(return_value=True),
    )
    monkeypatch.setattr(
        quota_service,
        "get_effective_ai_model_configuration_for_workflow",
        AsyncMock(return_value=_byok_config()),
    )
    hosted_authorize = AsyncMock(return_value=QuotaCheckResult(has_quota=True))
    monkeypatch.setattr(
        quota_service,
        "_authorize_hosted_workflow_run_start",
        hosted_authorize,
    )

    # actor_user.selected_organization_id=999 differs from workflow.organization_id=42,
    # but is_user_member_of_organization returns True so the run should be allowed.
    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=42,
        actor_user=SimpleNamespace(id=456, selected_organization_id=999),
    )

    assert result.has_quota is True


@pytest.mark.asyncio
async def test_authorize_workflow_run_requires_organization_scope(monkeypatch):
    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch)
    is_member_mock = AsyncMock()
    monkeypatch.setattr(
        quota_service.db_client,
        "is_user_member_of_organization",
        is_member_mock,
    )
    monkeypatch.setattr(
        quota_service,
        "get_effective_ai_model_configuration_for_workflow",
        AsyncMock(return_value=_byok_config()),
    )
    monkeypatch.setattr(
        quota_service,
        "_authorize_hosted_workflow_run_start",
        AsyncMock(return_value=QuotaCheckResult(has_quota=True)),
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=None,
        actor_user=SimpleNamespace(id=456),
    )

    assert result.has_quota is False
    assert result.error_code == "workflow_not_found"
    quota_service.db_client.get_workflow.assert_not_awaited()
    is_member_mock.assert_not_awaited()


@pytest.mark.asyncio
async def test_authorize_workflow_run_rejects_workflow_outside_org(monkeypatch):
    monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
    _patch_workflow_context(monkeypatch, workflow=None)
    is_member_mock = AsyncMock()
    monkeypatch.setattr(
        quota_service.db_client,
        "is_user_member_of_organization",
        is_member_mock,
    )

    result = await quota_service.authorize_workflow_run_start(
        workflow_id=7,
        organization_id=99,
        actor_user=SimpleNamespace(id=456),
    )

    assert result.has_quota is False
    assert result.error_code == "workflow_not_found"
    quota_service.db_client.get_workflow.assert_awaited_once_with(7, organization_id=99)
    is_member_mock.assert_not_awaited()
    quota_service.db_client.get_user_by_id.assert_not_awaited()