from unittest.mock import Mock
import pytest
from pr_review.terraform_infra import (
    AtlantisPRWaiter,
    TerraformInfraPromptBuilder,
    TerraformInfraPRAnalyzer,
)


class TestAtlantisPRWaiter:
    def test_wait_for_analysis_ready_success(self):
        """Test successful wait for Atlantis plan."""
        mock_github_client = Mock()
        pr_data = {
            'tf_plans': ['Plan output'],
            'diff': [],
            'last_commit_date': '2024-01-01T00:00:00Z',
        }
        mock_github_client.get_pr_details.return_value = pr_data
        mock_github_client.get_pr_status.side_effect = [
            {'state': 'pending'},
            {
                'state': 'pending',
                'statuses': [
                    {
                        'context': 'atlantis/plan',
                        'state': 'pending',
                    }
                ],
            },
            {
                'state': 'success',
                'statuses': [
                    {
                        'context': 'atlantis/plan',
                        'state': 'success',
                    }
                ],
            },
        ]

        waiter = AtlantisPRWaiter()
        result = waiter.wait_for_analysis_ready(mock_github_client, 123, 600, 0)

        assert result == pr_data
        mock_github_client.get_pr_status.assert_called_with(123)
        mock_github_client.get_pr_details.assert_called_with(123)

    def test_wait_for_analysis_ready_timeout(self):
        """Test timeout waiting for Atlantis plan."""
        mock_github_client = Mock()
        mock_github_client.get_pr_status.return_value = {
            'state': 'pending',
            'statuses': [
                {
                    'context': 'atlantis/plan',
                    'state': 'pending',
                }
            ],
        }

        waiter = AtlantisPRWaiter()

        with pytest.raises(TimeoutError) as e:
            waiter.wait_for_analysis_ready(mock_github_client, 123, 0, 0)

        assert str(e.value) == 'Timed out waiting for Atlantis plan status check'


class TestTerraformInfraPromptBuilder:
    def test_build_prompt(self):
        """Test building Terraform-specific prompt."""
        builder = TerraformInfraPromptBuilder()

        pr_data = {
            'pr_title': 'Add S3 bucket',
            'pr_body': 'Adding new S3 bucket for storage',
            'diff': [{'filename': 'main.tf', 'patch': 'added bucket'}],
            'tf_plans': ['Plan: +1 to add'],
        }

        result = builder.build_prompt(pr_data)

        assert 'Add S3 bucket' in result
        assert 'Adding new S3 bucket for storage' in result
        assert 'Plan: +1 to add' in result
        assert 'Terraform Pull Request Review Prompt' in result


class TestTerraformInfraPRAnalyzer:
    def test_post_process_analysis(self):
        """Test header de-emphasis."""
        mock_github_client = Mock()
        mock_llm_client = Mock()
        mock_prompt_builder = Mock()
        mock_pr_waiter = Mock()

        analyzer = TerraformInfraPRAnalyzer(
            mock_github_client, mock_llm_client, mock_prompt_builder, mock_pr_waiter
        )

        analysis = '# Header 1\n## Header 2\n### Header 3\n#### Header 4'
        result = analyzer.post_process_analysis(analysis)

        expected = '#### Header 1\n#### Header 2\n#### Header 3\n#### Header 4'
        assert result == expected