"""Tests for Identity CLI commands."""

import json
from unittest.mock import Mock, patch

import pytest
from typer.testing import CliRunner

from bedrock_agentcore_starter_toolkit.cli.identity.commands import identity_app
from bedrock_agentcore_starter_toolkit.utils.runtime.config import save_config
from bedrock_agentcore_starter_toolkit.utils.runtime.schema import (
    AWSConfig,
    BedrockAgentCoreAgentSchema,
    BedrockAgentCoreConfigSchema,
    BedrockAgentCoreDeploymentInfo,
    CredentialProviderInfo,
    IdentityConfig,
    NetworkConfiguration,
    ObservabilityConfig,
    WorkloadIdentityInfo,
)


@pytest.fixture
def runner():
    """Create CLI test runner."""
    return CliRunner()


@pytest.fixture
def test_config(tmp_path):
    """Create a test configuration file."""
    config_path = tmp_path / ".bedrock_agentcore.yaml"
    agent_config = BedrockAgentCoreAgentSchema(
        name="test-agent",
        entrypoint="test.py",
        aws=AWSConfig(
            region="us-west-2",
            network_configuration=NetworkConfiguration(),
            observability=ObservabilityConfig(),
        ),
        bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
    )
    project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
    save_config(project_config, config_path)
    return config_path


class TestCreateProvider:
    """Test create-credential-provider command."""

    def test_create_cognito_provider_success(self, runner, tmp_path, monkeypatch):
        """Test successful Cognito provider creation."""
        monkeypatch.chdir(tmp_path)

        # Create initial config
        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        # Mock IdentityClient at its source
        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_oauth2_credential_provider.return_value = {
                "credentialProviderArn": "arn:aws:identity:us-west-2:123456789012:provider/MyCognito",
                "callbackUrl": "https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            }
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-credential-provider",
                        "--name",
                        "MyCognito",
                        "--type",
                        "cognito",
                        "--client-id",
                        "abc123",
                        "--client-secret",
                        "xyz789",
                        "--discovery-url",
                        "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_xxx/.well-known/openid-configuration",
                    ],
                )

        assert result.exit_code == 0
        assert "Credential Provider Created" in result.stdout
        assert "MyCognito" in result.stdout

        # Verify client was called
        mock_identity.create_oauth2_credential_provider.assert_called_once()
        call_args = mock_identity.create_oauth2_credential_provider.call_args[0][0]
        assert call_args["name"] == "MyCognito"
        assert call_args["credentialProviderVendor"] == "CustomOauth2"

        # Verify config was saved
        from bedrock_agentcore_starter_toolkit.utils.runtime.config import load_config

        updated_config = load_config(config_path)
        updated_agent = updated_config.get_agent_config()
        assert updated_agent.identity is not None
        assert len(updated_agent.identity.credential_providers) == 1
        assert updated_agent.identity.credential_providers[0].name == "MyCognito"

    def test_create_github_provider_success(self, runner, tmp_path, monkeypatch):
        """Test GitHub provider creation."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_oauth2_credential_provider.return_value = {
                "credentialProviderArn": "arn:aws:identity:us-west-2:123456789012:provider/MyGitHub",
                "callbackUrl": "https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            }
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-credential-provider",
                        "--name",
                        "MyGitHub",
                        "--type",
                        "github",
                        "--client-id",
                        "github123",
                        "--client-secret",
                        "githubsecret",
                    ],
                )

        assert result.exit_code == 0
        assert "MyGitHub" in result.stdout

        call_args = mock_identity.create_oauth2_credential_provider.call_args[0][0]
        assert call_args["credentialProviderVendor"] == "GithubOauth2"

    def test_create_google_provider(self, runner, tmp_path, monkeypatch):
        """Test Google provider creation."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_oauth2_credential_provider.return_value = {
                "credentialProviderArn": "arn:aws:identity:us-west-2:123456789012:provider/MyGoogle",
                "callbackUrl": "https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            }
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-credential-provider",
                        "--name",
                        "MyGoogle",
                        "--type",
                        "google",
                        "--client-id",
                        "google123",
                        "--client-secret",
                        "googlesecret",
                    ],
                )

        assert result.exit_code == 0
        call_args = mock_identity.create_oauth2_credential_provider.call_args[0][0]
        assert call_args["credentialProviderVendor"] == "GoogleOauth2"

    def test_create_salesforce_provider(self, runner, tmp_path, monkeypatch):
        """Test Salesforce provider creation."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_oauth2_credential_provider.return_value = {
                "credentialProviderArn": "arn:aws:identity:us-west-2:123456789012:provider/MySalesforce",
                "callbackUrl": "https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            }
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-credential-provider",
                        "--name",
                        "MySalesforce",
                        "--type",
                        "salesforce",
                        "--client-id",
                        "sf123",
                        "--client-secret",
                        "sfsecret",
                    ],
                )

        assert result.exit_code == 0
        call_args = mock_identity.create_oauth2_credential_provider.call_args[0][0]
        assert call_args["credentialProviderVendor"] == "SalesforceOauth2"

    def test_create_provider_with_cognito_auto_update(self, runner, tmp_path, monkeypatch):
        """Test provider creation with automatic Cognito callback URL update."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_oauth2_credential_provider.return_value = {
                "credentialProviderArn": "arn:aws:identity:us-west-2:123456789012:provider/MyCognito",
                "callbackUrl": "https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            }
            mock_identity_class.return_value = mock_identity

            with (
                patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
                patch(
                    "bedrock_agentcore_starter_toolkit.cli.identity.commands.update_cognito_callback_urls"
                ) as mock_update,
            ):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-credential-provider",
                        "--name",
                        "MyCognito",
                        "--type",
                        "cognito",
                        "--client-id",
                        "abc123",
                        "--client-secret",
                        "xyz789",
                        "--discovery-url",
                        "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_xxx/.well-known/openid-configuration",
                        "--cognito-pool-id",
                        "us-west-2_testpool",
                    ],
                )

        assert result.exit_code == 0
        mock_update.assert_called_once_with(
            pool_id="us-west-2_testpool",
            client_id="abc123",
            callback_url="https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            region="us-west-2",
        )

    def test_create_provider_cognito_auto_update_failure(self, runner, tmp_path, monkeypatch):
        """Test provider creation when Cognito auto-update fails (should still succeed)."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_oauth2_credential_provider.return_value = {
                "credentialProviderArn": "arn:aws:identity:us-west-2:123456789012:provider/MyCognito",
                "callbackUrl": "https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            }
            mock_identity_class.return_value = mock_identity

            with (
                patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
                patch(
                    "bedrock_agentcore_starter_toolkit.cli.identity.commands.update_cognito_callback_urls",
                    side_effect=Exception("Update failed"),
                ),
            ):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-credential-provider",
                        "--name",
                        "MyCognito",
                        "--type",
                        "cognito",
                        "--client-id",
                        "abc123",
                        "--client-secret",
                        "xyz789",
                        "--discovery-url",
                        "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_xxx/.well-known/openid-configuration",
                        "--cognito-pool-id",
                        "us-west-2_testpool",
                    ],
                )

        # Should still succeed with warning
        assert result.exit_code == 0
        assert "manually add this callback URL" in result.stdout

    def test_create_provider_missing_discovery_url(self, runner, tmp_path, monkeypatch):
        """Test error when discovery URL missing for Cognito."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
            result = runner.invoke(
                identity_app,
                [
                    "create-credential-provider",
                    "--name",
                    "MyCognito",
                    "--type",
                    "cognito",
                    "--client-id",
                    "abc123",
                    "--client-secret",
                    "xyz789",
                ],
            )

        assert result.exit_code != 0

    def test_create_provider_unsupported_type(self, runner, tmp_path, monkeypatch):
        """Test error with unsupported provider type."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
            result = runner.invoke(
                identity_app,
                [
                    "create-credential-provider",
                    "--name",
                    "MyProvider",
                    "--type",
                    "unsupported",
                    "--client-id",
                    "abc123",
                    "--client-secret",
                    "xyz789",
                ],
            )

        assert result.exit_code != 0

    def test_create_provider_no_callback_url(self, runner, tmp_path, monkeypatch):
        """Test provider creation when no callback URL is returned."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_oauth2_credential_provider.return_value = {
                "credentialProviderArn": "arn:aws:identity:us-west-2:123456789012:provider/MyProvider",
                # No callbackUrl
            }
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-credential-provider",
                        "--name",
                        "MyGitHub",
                        "--type",
                        "github",
                        "--client-id",
                        "abc123",
                        "--client-secret",
                        "xyz789",
                    ],
                )

        assert result.exit_code == 0

    def test_create_provider_api_error(self, runner, tmp_path, monkeypatch):
        """Test error handling when API call fails."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_oauth2_credential_provider.side_effect = Exception("API Error")
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-credential-provider",
                        "--name",
                        "MyProvider",
                        "--type",
                        "github",
                        "--client-id",
                        "abc123",
                        "--client-secret",
                        "xyz789",
                    ],
                )

        assert result.exit_code != 0


class TestCreateWorkload:
    """Test create-workload-identity command."""

    def test_create_workload_with_name_and_urls(self, runner, tmp_path, monkeypatch):
        """Test workload creation with explicit name and return URLs."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_workload_identity.return_value = {
                "workloadIdentityArn": "arn:aws:identity:us-west-2:123456789012:workload/MyAgent"
            }
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-workload-identity",
                        "--name",
                        "MyAgent",
                        "--return-urls",
                        "http://localhost:8081/oauth2/callback,https://prod.example.com/callback",
                    ],
                )

        assert result.exit_code == 0
        assert "Workload Identity Created" in result.stdout
        assert "MyAgent" in result.stdout

        # Verify client was called correctly
        mock_identity.create_workload_identity.assert_called_once_with(
            name="MyAgent",
            allowed_resource_oauth_2_return_urls=[
                "http://localhost:8081/oauth2/callback",
                "https://prod.example.com/callback",
            ],
        )

    def test_create_workload_auto_generated_name(self, runner, tmp_path, monkeypatch):
        """Test workload creation with auto-generated name from config."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_workload_identity.return_value = {
                "workloadIdentityArn": "arn:aws:identity:us-west-2:123456789012:workload/test-agent-workload"
            }
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-workload-identity",
                        "--return-urls",
                        "http://localhost:8081/oauth2/callback",
                    ],
                )

        assert result.exit_code == 0

        # Verify name was auto-generated from config
        call_args = mock_identity.create_workload_identity.call_args[1]
        assert call_args["name"] == "test-agent-workload"

    def test_create_workload_no_config_generates_uuid(self, runner, tmp_path, monkeypatch):
        """Test workload creation generates UUID when no config exists."""
        monkeypatch.chdir(tmp_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_workload_identity.return_value = {
                "workloadIdentityArn": "arn:aws:identity:us-west-2:123456789012:workload/workload-abc123"
            }
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-workload-identity",
                        "--return-urls",
                        "http://localhost:8081/oauth2/callback",
                    ],
                )

        assert result.exit_code == 0
        call_args = mock_identity.create_workload_identity.call_args[1]
        assert call_args["name"].startswith("workload-")

    def test_create_workload_api_error(self, runner, tmp_path, monkeypatch):
        """Test error handling when API call fails."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.create_workload_identity.side_effect = Exception("API Error")
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "create-workload-identity",
                        "--name",
                        "MyAgent",
                        "--return-urls",
                        "http://localhost:8081/callback",
                    ],
                )

        assert result.exit_code != 0


class TestUpdateWorkload:
    """Test update-workload-identity command."""

    def test_update_workload_add_urls(self, runner, tmp_path, monkeypatch):
        """Test adding return URLs to existing workload."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.get_workload_identity.return_value = {
                "workloadIdentityArn": "arn:aws:identity:us-west-2:123456789012:workload/MyAgent",
                "allowedResourceOauth2ReturnUrls": ["http://localhost:8081/oauth2/callback"],
            }
            mock_identity.update_workload_identity.return_value = {}
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "update-workload-identity",
                        "--name",
                        "MyAgent",
                        "--add-return-urls",
                        "https://prod.example.com/callback",
                    ],
                )

        assert result.exit_code == 0
        assert "Workload Identity Updated" in result.stdout

        # Verify update was called with combined URLs
        call_args = mock_identity.update_workload_identity.call_args[1]
        assert set(call_args["allowed_resource_oauth_2_return_urls"]) == {
            "http://localhost:8081/oauth2/callback",
            "https://prod.example.com/callback",
        }

    def test_update_workload_set_urls(self, runner, tmp_path, monkeypatch):
        """Test replacing all return URLs."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.get_workload_identity.return_value = {
                "workloadIdentityArn": "arn:aws:identity:us-west-2:123456789012:workload/MyAgent",
                "allowedResourceOauth2ReturnUrls": ["http://localhost:8081/oauth2/callback"],
            }
            mock_identity.update_workload_identity.return_value = {}
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "update-workload-identity",
                        "--name",
                        "MyAgent",
                        "--set-return-urls",
                        "https://new1.example.com/callback,https://new2.example.com/callback",
                    ],
                )

        assert result.exit_code == 0

        # Verify URLs were replaced
        call_args = mock_identity.update_workload_identity.call_args[1]
        assert set(call_args["allowed_resource_oauth_2_return_urls"]) == {
            "https://new1.example.com/callback",
            "https://new2.example.com/callback",
        }

    def test_update_workload_no_options_error(self, runner, tmp_path, monkeypatch):
        """Test error when neither add nor set options provided."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
            result = runner.invoke(identity_app, ["update-workload-identity", "--name", "MyAgent"])

        assert result.exit_code != 0

    def test_update_workload_api_error(self, runner, tmp_path, monkeypatch):
        """Test error handling when API call fails."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.get_workload_identity.side_effect = Exception("API Error")
            mock_identity_class.return_value = mock_identity

            with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
                result = runner.invoke(
                    identity_app,
                    [
                        "update-workload-identity",
                        "--name",
                        "MyAgent",
                        "--add-return-urls",
                        "https://example.com/callback",
                    ],
                )

        assert result.exit_code != 0


class TestGetToken:
    """Test get-cognito-inbound-token command."""

    def test_get_token_user_flow_without_secret(self, runner):
        """Test getting token from Cognito using USER flow without client secret."""
        with (
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_access_token") as mock_get_token,
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
        ):
            mock_get_token.return_value = "test-access-token-12345"

            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "user",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    "--username",
                    "testuser",
                    "--password",
                    "Pass123!",
                ],
            )

        assert result.exit_code == 0
        assert "test-access-token-12345" in result.stdout

        mock_get_token.assert_called_once_with(
            pool_id="us-west-2_testpool",
            client_id="abc123",
            username="testuser",
            password="Pass123!",
            client_secret=None,
            region="us-west-2",
        )

    def test_get_token_user_flow_with_secret(self, runner):
        """Test getting token with client secret (USER flow)."""
        with (
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_access_token") as mock_get_token,
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
        ):
            mock_get_token.return_value = "test-access-token-with-secret"

            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "user",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    "--username",
                    "testuser",
                    "--password",
                    "Pass123!",
                    "--client-secret",
                    "mysecret",
                ],
            )

        assert result.exit_code == 0
        assert "test-access-token-with-secret" in result.stdout

        call_args = mock_get_token.call_args[1]
        assert call_args["client_secret"] == "mysecret"

    def test_get_token_user_flow_default(self, runner):
        """Test USER flow is default when --auth-flow not specified."""
        with (
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_access_token") as mock_get_token,
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
        ):
            mock_get_token.return_value = "default-flow-token"

            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    "--username",
                    "testuser",
                    "--password",
                    "Pass123!",
                ],
            )

        assert result.exit_code == 0
        mock_get_token.assert_called_once()

    def test_get_token_m2m_flow_success(self, runner):
        """Test getting token using M2M flow."""
        with (
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_m2m_token") as mock_m2m_token,
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
        ):
            mock_m2m_token.return_value = "m2m-access-token-xyz"

            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "m2m",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    "--client-secret",
                    "secret789",
                ],
            )

        assert result.exit_code == 0
        assert "m2m-access-token-xyz" in result.stdout

        mock_m2m_token.assert_called_once_with(
            pool_id="us-west-2_testpool",
            client_id="abc123",
            client_secret="secret789",
            region="us-west-2",
        )

    def test_get_token_user_flow_missing_username(self, runner):
        """Test USER flow error when username missing."""
        with (
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_access_token") as mock_get_token,
            patch.dict("os.environ", {}, clear=True),  # Clear all environment variables
        ):
            mock_get_token.return_value = "should-not-be-called"

            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "user",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    "--password",
                    "Pass123!",
                    # Missing --username
                ],
            )

        print(f"Exit code: {result.exit_code}")
        print(f"Output: {result.stdout}")
        assert result.exit_code != 0
        assert "Username required for USER flow" in result.stdout

    def test_get_token_user_flow_missing_password(self, runner):
        """Test USER flow error when password missing."""
        with (
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_access_token") as mock_get_token,
            patch.dict("os.environ", {}, clear=True),  # Clear all environment variables
        ):
            mock_get_token.return_value = "should-not-be-called"

            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "user",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    "--username",
                    "testuser",
                    # Missing --password
                ],
            )

        print(f"Exit code: {result.exit_code}")
        print(f"Output: {result.stdout}")
        assert result.exit_code != 0
        assert "Password required for USER flow" in result.stdout

    def test_get_token_m2m_flow_missing_secret(self, runner):
        """Test M2M flow error when client secret missing."""
        with (
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_m2m_token") as mock_m2m_token,
            patch.dict("os.environ", {}, clear=True),  # Clear all environment variables
        ):
            mock_m2m_token.return_value = "should-not-be-called"

            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "m2m",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    # Missing --client-secret
                ],
            )

        print(f"Exit code: {result.exit_code}")
        print(f"Output: {result.stdout}")
        assert result.exit_code != 0
        assert "Client secret required for M2M flow" in result.stdout

    def test_get_token_invalid_auth_flow(self, runner):
        """Test error with invalid auth flow type."""
        with patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"):
            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "invalid",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                ],
            )

        assert result.exit_code != 0
        assert "--auth-flow must be 'user' or 'm2m'" in result.stdout

    def test_get_token_user_flow_error(self, runner):
        """Test error handling when token retrieval fails (USER flow)."""
        with (
            patch(
                "bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_access_token",
                side_effect=Exception("Auth failed"),
            ),
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
        ):
            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "user",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    "--username",
                    "testuser",
                    "--password",
                    "Pass123!",
                ],
            )

        assert result.exit_code != 0

    def test_get_token_m2m_flow_error(self, runner):
        """Test error handling when M2M token retrieval fails."""
        with (
            patch(
                "bedrock_agentcore_starter_toolkit.cli.identity.commands.get_cognito_m2m_token",
                side_effect=Exception("M2M auth failed"),
            ),
            patch("bedrock_agentcore_starter_toolkit.cli.identity.commands.get_region", return_value="us-west-2"),
        ):
            result = runner.invoke(
                identity_app,
                [
                    "get-cognito-inbound-token",
                    "--auth-flow",
                    "m2m",
                    "--pool-id",
                    "us-west-2_testpool",
                    "--client-id",
                    "abc123",
                    "--client-secret",
                    "secret789",
                ],
            )

        assert result.exit_code != 0


class TestListProviders:
    """Test list-credential-providers command."""

    def test_list_providers_success(self, runner, tmp_path, monkeypatch):
        """Test listing configured providers."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"

        # Create identity config with providers
        identity_config = IdentityConfig()
        identity_config.credential_providers = [
            CredentialProviderInfo(
                name="MyCognito",
                arn="arn:aws:identity:us-west-2:123456789012:provider/MyCognito",
                type="cognito",
                callback_url="https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            ),
            CredentialProviderInfo(
                name="MyGitHub",
                arn="arn:aws:identity:us-west-2:123456789012:provider/MyGitHub",
                type="github",
                callback_url="https://bedrock-agentcore.us-west-2.amazonaws.com/callback2",
            ),
        ]

        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
            identity=identity_config,
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        result = runner.invoke(identity_app, ["list-credential-providers"])

        assert result.exit_code == 0, f"Expected exit code 0, got {result.exit_code}. Output: {result.stdout}"
        assert "Configured Credential Providers" in result.stdout
        assert "MyCognito" in result.stdout
        assert "MyGitHub" in result.stdout
        assert "cognito" in result.stdout
        assert "github" in result.stdout

    def test_list_providers_with_workload(self, runner, tmp_path, monkeypatch):
        """Test listing providers when workload identity is configured."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"

        identity_config = IdentityConfig()
        identity_config.credential_providers = [
            CredentialProviderInfo(
                name="MyCognito",
                arn="arn:aws:identity:us-west-2:123456789012:provider/MyCognito",
                type="cognito",
                callback_url="https://bedrock-agentcore.us-west-2.amazonaws.com/callback",
            ),
        ]
        identity_config.workload = WorkloadIdentityInfo(
            name="test-agent-workload",
            arn="arn:aws:identity:us-west-2:123456789012:workload/test-agent-workload",
            return_urls=["http://localhost:8081/oauth2/callback"],
        )

        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
            identity=identity_config,
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        result = runner.invoke(identity_app, ["list-credential-providers"])

        assert result.exit_code == 0
        assert "Workload Identity:" in result.stdout
        assert "test-agent-workload" in result.stdout
        assert "App Return URLs" in result.stdout

    def test_list_providers_no_config(self, runner, tmp_path, monkeypatch):
        """Test list-credential-providers when no config file exists."""
        monkeypatch.chdir(tmp_path)

        result = runner.invoke(identity_app, ["list-credential-providers"])

        assert result.exit_code == 1
        assert "No .bedrock_agentcore.yaml found" in result.stdout

    def test_list_providers_empty(self, runner, tmp_path, monkeypatch):
        """Test list-credential-providers when no providers configured."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"

        # Create identity config with explicitly empty credential providers
        identity_config = IdentityConfig()
        identity_config.credential_providers = []

        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
            identity=identity_config,
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        result = runner.invoke(identity_app, ["list-credential-providers"])

        # Command shows helpful message when no providers configured
        # Note: Currently exits with code 1 (typer.Exit(0) caught by error handler)
        # This is acceptable as it indicates "no results" rather than success with results
        assert result.exit_code == 1
        assert "No credential providers configured" in result.stdout
        assert "agentcore identity create-credential-provider" in result.stdout

    def test_list_providers_no_identity_attribute(self, runner, tmp_path, monkeypatch):
        """Test list-credential-providers when identity attribute is not set at all."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"

        # Create agent config without identity attribute
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
            # No identity attribute set
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        result = runner.invoke(identity_app, ["list-credential-providers"])

        # Should handle missing identity attribute gracefully
        expected_exit_code_msg = (
            f"Expected exit code 1, got {result.exit_code}. "
            f"Output: {result.stdout}\n"
            f"Error: {result.stderr if hasattr(result, 'stderr') else 'N/A'}"
        )
        assert result.exit_code == 1, expected_exit_code_msg
        assert (
            "No credential providers configured" in result.stdout or "No .bedrock_agentcore.yaml found" in result.stdout
        )


class TestSetupCognito:
    """Test setup-cognito command."""

    def test_setup_cognito_user_flow_success(self, runner, tmp_path, monkeypatch):
        """Test successful Cognito pool setup with user flow."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        mock_result = {
            "runtime": {
                "pool_id": "us-west-2_runtime123",
                "client_id": "runtime_client_123",
                "discovery_url": (
                    "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_runtime123/.well-known/openid-configuration"
                ),
                "username": "testuser1234",
                "password": "TestPass123!@#",
            },
            "identity": {
                "pool_id": "us-west-2_identity456",
                "client_id": "identity_client_456",
                "client_secret": "identity_secret_789",
                "discovery_url": (
                    "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_identity456/.well-known/openid-configuration"
                ),
                "username": "externaluser5678",
                "password": "ExtPass456!@#",
            },
        }

        with patch(
            "bedrock_agentcore_starter_toolkit.cli.identity.commands.IdentityCognitoManager"
        ) as mock_manager_class:
            mock_manager = Mock()
            mock_manager.create_user_federation_pools.return_value = mock_result
            mock_manager_class.return_value = mock_manager

            result = runner.invoke(identity_app, ["setup-cognito", "--region", "us-west-2", "--auth-flow", "user"])

        assert result.exit_code == 0
        assert "Cognito pools created successfully" in result.stdout
        assert "Runtime Pool (Inbound Auth)" in result.stdout
        assert "Identity Pool" in result.stdout
        assert "us-west-2_runtime123" in result.stdout
        assert "us-west-2_identity456" in result.stdout

        # Verify files were created with correct naming
        assert (tmp_path / ".agentcore_identity_cognito_user.json").exists()
        assert (tmp_path / ".agentcore_identity_user.env").exists()

        # Verify JSON file content
        with open(tmp_path / ".agentcore_identity_cognito_user.json") as f:
            saved_config = json.load(f)
            assert saved_config == mock_result

    def test_setup_cognito_m2m_flow_success(self, runner, tmp_path, monkeypatch):
        """Test successful Cognito pool setup with m2m flow."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        mock_result = {
            "runtime": {
                "pool_id": "us-west-2_runtime123",
                "client_id": "runtime_client_123",
                "discovery_url": (
                    "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_runtime123/.well-known/openid-configuration"
                ),
                "username": "testuser1234",
                "password": "TestPass123!@#",
            },
            "identity": {
                "pool_id": "us-west-2_identity456",
                "client_id": "identity_client_456",
                "client_secret": "identity_secret_789",
                "token_endpoint": "https://agentcore-identity-abc123.auth.us-west-2.amazoncognito.com/oauth2/token",
                "resource_server_identifier": "https://api.example.com",
            },
        }

        with patch(
            "bedrock_agentcore_starter_toolkit.cli.identity.commands.IdentityCognitoManager"
        ) as mock_manager_class:
            mock_manager = Mock()
            mock_manager.create_m2m_pools.return_value = mock_result
            mock_manager_class.return_value = mock_manager

            result = runner.invoke(identity_app, ["setup-cognito", "--region", "us-west-2", "--auth-flow", "m2m"])

        assert result.exit_code == 0
        assert "Cognito pools created successfully" in result.stdout
        assert "M2M" in result.stdout or "m2m" in result.stdout.lower()

        # Verify files were created with correct naming
        assert (tmp_path / ".agentcore_identity_cognito_m2m.json").exists()
        assert (tmp_path / ".agentcore_identity_m2m.env").exists()

    def test_setup_cognito_uses_config_region(self, runner, tmp_path, monkeypatch):
        """Test setup-cognito uses region from config when not specified."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="eu-west-1",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        mock_result = {
            "runtime": {
                "pool_id": "eu-west-1_runtime",
                "client_id": "client1",
                "discovery_url": "https://example.com",
                "username": "user1",
                "password": "pass1",
            },
            "identity": {
                "pool_id": "eu-west-1_identity",
                "client_id": "client2",
                "client_secret": "secret2",
                "discovery_url": "https://example.com",
                "username": "user2",
                "password": "pass2",
            },
        }

        with patch(
            "bedrock_agentcore_starter_toolkit.cli.identity.commands.IdentityCognitoManager"
        ) as mock_manager_class:
            mock_manager = Mock()
            mock_manager.create_user_federation_pools.return_value = mock_result
            mock_manager_class.return_value = mock_manager

            result = runner.invoke(identity_app, ["setup-cognito"])

        assert result.exit_code == 0
        # Verify manager was created with eu-west-1
        mock_manager_class.assert_called_once_with("eu-west-1")

    def test_setup_cognito_fallback_region(self, runner, tmp_path, monkeypatch):
        """Test setup-cognito falls back to boto3 session region."""
        monkeypatch.chdir(tmp_path)

        mock_result = {
            "runtime": {
                "pool_id": "us-east-1_runtime",
                "client_id": "client1",
                "discovery_url": "https://example.com",
                "username": "user1",
                "password": "pass1",
            },
            "identity": {
                "pool_id": "us-east-1_identity",
                "client_id": "client2",
                "client_secret": "secret2",
                "discovery_url": "https://example.com",
                "username": "user2",
                "password": "pass2",
            },
        }

        with (
            patch(
                "bedrock_agentcore_starter_toolkit.cli.identity.commands.IdentityCognitoManager"
            ) as mock_manager_class,
            patch("boto3.Session") as mock_session_class,
        ):
            mock_manager = Mock()
            mock_manager.create_user_federation_pools.return_value = mock_result
            mock_manager_class.return_value = mock_manager

            mock_session = Mock()
            mock_session.region_name = "us-east-1"
            mock_session_class.return_value = mock_session

            result = runner.invoke(identity_app, ["setup-cognito"])

        assert result.exit_code == 0
        mock_manager_class.assert_called_once_with("us-east-1")

    def test_setup_cognito_invalid_auth_flow(self, runner, tmp_path, monkeypatch):
        """Test setup-cognito with invalid auth flow."""
        monkeypatch.chdir(tmp_path)

        result = runner.invoke(identity_app, ["setup-cognito", "--auth-flow", "invalid"])

        assert result.exit_code == 1
        assert "--auth-flow must be 'user' or 'm2m'" in result.stdout

    def test_setup_cognito_error(self, runner, tmp_path, monkeypatch):
        """Test error handling when setup fails."""
        monkeypatch.chdir(tmp_path)

        with patch(
            "bedrock_agentcore_starter_toolkit.cli.identity.commands.IdentityCognitoManager"
        ) as mock_manager_class:
            mock_manager = Mock()
            mock_manager.create_user_federation_pools.side_effect = Exception("Setup failed")
            mock_manager_class.return_value = mock_manager

            result = runner.invoke(identity_app, ["setup-cognito", "--region", "us-west-2"])

        assert result.exit_code != 0


class TestCleanup:
    """Test cleanup command."""

    def test_cleanup_success_with_force(self, runner, tmp_path, monkeypatch):
        """Test successful cleanup with force flag."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        identity_config = IdentityConfig()
        identity_config.credential_providers = [
            CredentialProviderInfo(
                name="TestProvider",
                arn="arn:aws:identity:us-west-2:123456789012:provider/TestProvider",
                type="cognito",
                callback_url="https://example.com/callback",
            ),
        ]
        identity_config.workload = WorkloadIdentityInfo(
            name="test-workload",
            arn="arn:aws:identity:us-west-2:123456789012:workload/test-workload",
            return_urls=["http://localhost:8081/callback"],
        )

        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
            identity=identity_config,
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        # Create Cognito config files for both flows
        for flow in ["user", "m2m"]:
            cognito_config = {
                "runtime": {"pool_id": f"us-west-2_runtime_{flow}"},
                "identity": {"pool_id": f"us-west-2_identity_{flow}"},
            }
            cognito_config_path = tmp_path / f".agentcore_identity_cognito_{flow}.json"
            with open(cognito_config_path, "w") as f:
                json.dump(cognito_config, f)

            env_file_path = tmp_path / f".agentcore_identity_{flow}.env"
            env_file_path.write_text(f"export TEST_{flow.upper()}=1")

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_identity.cp_client = Mock()
            mock_identity.identity_client = Mock()
            mock_identity_class.return_value = mock_identity

            with patch(
                "bedrock_agentcore_starter_toolkit.cli.identity.commands.IdentityCognitoManager"
            ) as mock_manager_class:
                mock_manager = Mock()
                mock_manager_class.return_value = mock_manager

                result = runner.invoke(identity_app, ["cleanup", "--force"])

        assert result.exit_code == 0
        assert "Identity cleanup complete" in result.stdout

        # Verify deletions were called
        mock_identity.cp_client.delete_oauth2_credential_provider.assert_called_once_with(name="TestProvider")
        mock_identity.identity_client.delete_workload_identity.assert_called_once_with(name="test-workload")

        # Verify Cognito cleanup was called for each flow
        assert mock_manager.cleanup_cognito_pools.call_count == 2

        # Verify Cognito config files were deleted
        assert not (tmp_path / ".agentcore_identity_cognito_user.json").exists()
        assert not (tmp_path / ".agentcore_identity_cognito_m2m.json").exists()
        assert not (tmp_path / ".agentcore_identity_user.env").exists()
        assert not (tmp_path / ".agentcore_identity_m2m.env").exists()

    def test_cleanup_without_force_cancelled(self, runner, tmp_path, monkeypatch):
        """Test cleanup cancelled when user declines confirmation."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        identity_config = IdentityConfig()
        identity_config.credential_providers = [
            CredentialProviderInfo(
                name="TestProvider",
                arn="arn:aws:identity:us-west-2:123456789012:provider/TestProvider",
                type="cognito",
                callback_url="https://example.com/callback",
            ),
        ]

        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
            identity=identity_config,
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        # Simulate user declining confirmation
        result = runner.invoke(identity_app, ["cleanup"], input="n\n")

        assert result.exit_code == 0
        assert "Cancelled" in result.stdout

    def test_cleanup_no_config_error(self, runner, tmp_path, monkeypatch):
        """Test cleanup fails when no config file exists."""
        monkeypatch.chdir(tmp_path)

        result = runner.invoke(identity_app, ["cleanup", "--force"])

        assert result.exit_code == 1
        assert "No .bedrock_agentcore.yaml found" in result.stdout

    def test_cleanup_provider_deletion_error(self, runner, tmp_path, monkeypatch):
        """Test cleanup continues when provider deletion fails."""
        monkeypatch.chdir(tmp_path)

        config_path = tmp_path / ".bedrock_agentcore.yaml"
        identity_config = IdentityConfig()
        identity_config.credential_providers = [
            CredentialProviderInfo(
                name="TestProvider",
                arn="arn:aws:identity:us-west-2:123456789012:provider/TestProvider",
                type="cognito",
                callback_url="https://example.com/callback",
            ),
        ]

        agent_config = BedrockAgentCoreAgentSchema(
            name="test-agent",
            entrypoint="test.py",
            aws=AWSConfig(
                region="us-west-2",
                network_configuration=NetworkConfiguration(),
                observability=ObservabilityConfig(),
            ),
            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
            identity=identity_config,
        )
        project_config = BedrockAgentCoreConfigSchema(default_agent="test-agent", agents={"test-agent": agent_config})
        save_config(project_config, config_path)

        with patch("bedrock_agentcore.services.identity.IdentityClient") as mock_identity_class:
            mock_identity = Mock()
            mock_cp_client = Mock()

            # Create a mock exception class for ResourceNotFoundException
            mock_exceptions = Mock()
            mock_exceptions.ResourceNotFoundException = type("ResourceNotFoundException", (Exception,), {})
            mock_cp_client.exceptions = mock_exceptions

            # Set up the deletion to raise a generic exception (not ResourceNotFoundException)
            mock_cp_client.delete_oauth2_credential_provider.side_effect = Exception("Deletion failed")

            mock_identity.cp_client = mock_cp_client
            mock_identity.identity_client = Mock()
            mock_identity_class.return_value = mock_identity

            result = runner.invoke(identity_app, ["cleanup", "--force"])

        # Should complete despite error (shows warning but continues)
        assert result.exit_code == 0
        assert "Error:" in result.stdout or "⚠️" in result.stdout


class TestBuildProviderConfig:
    """Test _build_provider_config helper function."""

    def test_build_cognito_config(self):
        """Test building Cognito provider config."""
        from bedrock_agentcore_starter_toolkit.cli.identity.commands import _build_provider_config

        config = _build_provider_config(
            provider_type="cognito",
            name="MyCognito",
            client_id="abc123",
            client_secret="xyz789",
            discovery_url=(
                "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_xxx/.well-known/openid-configuration"
            ),
        )

        assert config["name"] == "MyCognito"
        assert config["credentialProviderVendor"] == "CustomOauth2"
        assert config["oauth2ProviderConfigInput"]["customOauth2ProviderConfig"]["clientId"] == "abc123"
        assert config["oauth2ProviderConfigInput"]["customOauth2ProviderConfig"]["clientSecret"] == "xyz789"
        assert (
            config["oauth2ProviderConfigInput"]["customOauth2ProviderConfig"]["oauthDiscovery"]["discoveryUrl"]
            == "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_xxx/.well-known/openid-configuration"
        )

    def test_build_github_config(self):
        """Test building GitHub provider config."""
        from bedrock_agentcore_starter_toolkit.cli.identity.commands import _build_provider_config

        config = _build_provider_config(
            provider_type="github",
            name="MyGitHub",
            client_id="github123",
            client_secret="githubsecret",
            discovery_url=None,
        )

        assert config["name"] == "MyGitHub"
        assert config["credentialProviderVendor"] == "GithubOauth2"
        assert config["oauth2ProviderConfigInput"]["githubOauth2ProviderConfig"]["clientId"] == "github123"
        assert config["oauth2ProviderConfigInput"]["githubOauth2ProviderConfig"]["clientSecret"] == "githubsecret"

    def test_build_google_config(self):
        """Test building Google provider config."""
        from bedrock_agentcore_starter_toolkit.cli.identity.commands import _build_provider_config

        config = _build_provider_config(
            provider_type="google",
            name="MyGoogle",
            client_id="google123",
            client_secret="googlesecret",
            discovery_url=None,
        )

        assert config["credentialProviderVendor"] == "GoogleOauth2"
        assert config["oauth2ProviderConfigInput"]["googleOauth2ProviderConfig"]["clientId"] == "google123"

    def test_build_salesforce_config(self):
        """Test building Salesforce provider config."""
        from bedrock_agentcore_starter_toolkit.cli.identity.commands import _build_provider_config

        config = _build_provider_config(
            provider_type="salesforce",
            name="MySalesforce",
            client_id="sf123",
            client_secret="sfsecret",
            discovery_url=None,
        )

        assert config["credentialProviderVendor"] == "SalesforceOauth2"
        assert config["oauth2ProviderConfigInput"]["salesforceOauth2ProviderConfig"]["clientId"] == "sf123"
