"""
Unit tests for generate_summary.py module
"""
import pytest
import os
import json
from unittest.mock import Mock, MagicMock, patch, call
from pydantic import ValidationError

# Import the functions and models to test
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from scripts.generate_summary import (
    general_summary,
    custom_summary,
    BasicSummary,
    AdvancedSummary,
    KeyPoint,
    Summary,
    Purpose,
    Chapters,
    Outcomes,
    ActionItemsPerUser
)


class TestPydanticModels:
    """Test Pydantic model validation"""
    
    def test_key_point_model(self):
        """Test KeyPoint model validation"""
        key_point = KeyPoint(text="Test point", timestamp=10.5)
        assert key_point.text == "Test point"
        assert key_point.timestamp == 10.5
    
    def test_summary_model(self):
        """Test Summary model validation"""
        summary = Summary(text="Test summary", duration_minutes=15.5)
        assert summary.text == "Test summary"
        assert summary.duration_minutes == 15.5
    
    def test_basic_summary_model(self, sample_basic_summary_response):
        """Test BasicSummary model validation"""
        basic_summary = BasicSummary(**sample_basic_summary_response)
        assert len(basic_summary.Key_Points) == 2
        assert basic_summary.Summary.text is not None
        assert basic_summary.Summary.duration_minutes > 0
    
    def test_advanced_summary_model(self, sample_advanced_summary_response):
        """Test AdvancedSummary model validation"""
        advanced_summary = AdvancedSummary(**sample_advanced_summary_response)
        assert advanced_summary.Purpose.text is not None
        assert len(advanced_summary.Chapters.content) > 0
        assert len(advanced_summary.Outcomes.content) > 0
        assert len(advanced_summary.Action_Items_Per_User) > 0


class TestGeneralSummary:
    """Test general_summary function"""
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_general_summary_freemium_plan(
        self, 
        mock_prompt_template, 
        mock_chat_anthropic,
        sample_transcription_dict,
        sample_basic_summary_response
    ):
        """Test general_summary with freemium plan"""
        # Setup mocks
        mock_model = MagicMock()
        mock_chain = MagicMock()
        mock_structured_model = MagicMock()
        
        # Create a mock BasicSummary instance
        from scripts.generate_summary import BasicSummary
        mock_result = BasicSummary(**sample_basic_summary_response)
        
        mock_structured_model.invoke.return_value = mock_result
        mock_model.with_structured_output.return_value = mock_structured_model
        mock_prompt_template.from_messages.return_value = MagicMock()
        mock_chat_anthropic.return_value = mock_model
        
        # Mock the pipe operator
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        mock_prompt_instance.__or__ = lambda self, other: mock_chain
        mock_chain.invoke.return_value = mock_result
        
        # Call the function
        result = general_summary(sample_transcription_dict, plan_tier="freemium")
        
        # Verify results
        assert result is not None
        assert "Key_Points" in result
        assert "Summary" in result
        assert len(result["Key_Points"]) > 0
        
        # Verify model was initialized with correct parameters
        mock_chat_anthropic.assert_called_once()
        call_args = mock_chat_anthropic.call_args
        assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
        assert call_args.kwargs["max_tokens"] == 2000
        assert call_args.kwargs["temperature"] == 0.2
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_general_summary_pro_plan(
        self,
        mock_prompt_template,
        mock_chat_anthropic,
        sample_transcription_dict,
        sample_advanced_summary_response
    ):
        """Test general_summary with pro plan"""
        # Setup mocks
        mock_model = MagicMock()
        mock_chain = MagicMock()
        
        # Create a mock AdvancedSummary instance
        from scripts.generate_summary import AdvancedSummary
        mock_result = AdvancedSummary(**sample_advanced_summary_response)
        
        mock_structured_model = MagicMock()
        mock_structured_model.invoke.return_value = mock_result
        mock_model.with_structured_output.return_value = mock_structured_model
        mock_prompt_template.from_messages.return_value = MagicMock()
        mock_chat_anthropic.return_value = mock_model
        
        # Mock the pipe operator
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        mock_prompt_instance.__or__ = lambda self, other: mock_chain
        mock_chain.invoke.return_value = mock_result
        
        # Call the function
        result = general_summary(sample_transcription_dict, plan_tier="pro")
        
        # Verify results
        assert result is not None
        assert "Purpose" in result
        assert "Chapters" in result
        assert "Outcomes" in result
        assert "Action_Items_Per_User" in result
        
        # Verify model was initialized with correct parameters
        mock_chat_anthropic.assert_called_once()
        call_args = mock_chat_anthropic.call_args
        assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
        assert call_args.kwargs["max_tokens"] == 4000
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_general_summary_with_string_transcription(
        self,
        mock_prompt_template,
        mock_chat_anthropic,
        sample_basic_summary_response
    ):
        """Test general_summary with string transcription"""
        transcription_str = json.dumps({"sentences": []})
        
        mock_model = MagicMock()
        mock_chain = MagicMock()
        from scripts.generate_summary import BasicSummary
        mock_result = BasicSummary(**sample_basic_summary_response)
        
        mock_structured_model = MagicMock()
        mock_structured_model.invoke.return_value = mock_result
        mock_model.with_structured_output.return_value = mock_structured_model
        mock_prompt_template.from_messages.return_value = MagicMock()
        mock_chat_anthropic.return_value = mock_model
        
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        mock_prompt_instance.__or__ = lambda self, other: mock_chain
        mock_chain.invoke.return_value = mock_result
        
        result = general_summary(transcription_str, plan_tier="freemium")
        
        assert result is not None
        assert "Key_Points" in result
    
    def test_general_summary_missing_api_key(self, sample_transcription_dict):
        """Test general_summary raises error when API key is missing"""
        with patch.dict(os.environ, {}, clear=True):
            with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"):
                general_summary(sample_transcription_dict, plan_tier="freemium")
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_general_summary_fallback_on_error(
        self,
        mock_prompt_template,
        mock_chat_anthropic,
        sample_transcription_dict,
        sample_basic_summary_response
    ):
        """Test general_summary falls back to non-structured output on error"""
        mock_model = MagicMock()
        mock_chain = MagicMock()
        mock_fallback_chain = MagicMock()
        mock_response = MagicMock()
        mock_response.content = json.dumps(sample_basic_summary_response)
        
        # First call (structured) raises error
        mock_structured_model = MagicMock()
        mock_structured_model.invoke.side_effect = Exception("Structured output failed")
        mock_model.with_structured_output.return_value = mock_structured_model
        
        # Set up the chain so that structured chain raises exception
        # and fallback chain returns the mock_response
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        
        # When __or__ is called with structured_model, return chain that raises exception
        # When __or__ is called with model (fallback), return fallback_chain that succeeds
        def or_handler(self, other):
            if other == mock_structured_model:
                # Structured chain - should raise exception
                mock_chain.invoke.side_effect = Exception("Structured output failed")
                return mock_chain
            elif other == mock_model:
                # Fallback chain - should succeed
                mock_fallback_chain.invoke.return_value = mock_response
                return mock_fallback_chain
            return mock_chain
        
        mock_prompt_instance.__or__ = or_handler
        mock_chat_anthropic.return_value = mock_model
        
        result = general_summary(sample_transcription_dict, plan_tier="freemium")
        
        assert result is not None
        assert "Key_Points" in result


class TestCustomSummary:
    """Test custom_summary function"""
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_custom_summary_success(
        self,
        mock_prompt_template,
        mock_chat_anthropic,
        sample_transcription_dict,
        sample_template
    ):
        """Test custom_summary with successful response"""
        mock_model = MagicMock()
        mock_chain = MagicMock()
        mock_response = MagicMock()
        
        expected_result = {
            "Key_Points": {"content": []},
            "Summary": {"content": []},
            "Next_Steps": {"content": []}
        }
        mock_response.content = json.dumps(expected_result)
        
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        mock_prompt_instance.__or__ = lambda self, other: mock_chain
        mock_chain.invoke.return_value = mock_response
        mock_chat_anthropic.return_value = mock_model
        
        result = custom_summary(sample_template, sample_transcription_dict)
        
        assert result is not None
        assert "Key_Points" in result or "Summary" in result
        
        # Verify model was initialized correctly
        mock_chat_anthropic.assert_called_once()
        call_args = mock_chat_anthropic.call_args
        assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
        assert call_args.kwargs["max_tokens"] == 8000
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_custom_summary_with_markdown_wrapper(
        self,
        mock_prompt_template,
        mock_chat_anthropic,
        sample_transcription_dict,
        sample_template
    ):
        """Test custom_summary handles markdown-wrapped JSON"""
        mock_model = MagicMock()
        mock_chain = MagicMock()
        mock_response = MagicMock()
        
        expected_result = {"result": "test"}
        wrapped_json = f"```json\n{json.dumps(expected_result)}\n```"
        mock_response.content = wrapped_json
        
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        mock_prompt_instance.__or__ = lambda self, other: mock_chain
        mock_chain.invoke.return_value = mock_response
        mock_chat_anthropic.return_value = mock_model
        
        result = custom_summary(sample_template, sample_transcription_dict)
        
        assert result == expected_result
    
    def test_custom_summary_missing_api_key(self, sample_transcription_dict, sample_template):
        """Test custom_summary raises error when API key is missing"""
        with patch.dict(os.environ, {}, clear=True):
            with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"):
                custom_summary(sample_template, sample_transcription_dict)
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_custom_summary_invalid_json(
        self,
        mock_prompt_template,
        mock_chat_anthropic,
        sample_transcription_dict,
        sample_template
    ):
        """Test custom_summary handles invalid JSON gracefully"""
        mock_model = MagicMock()
        mock_chain = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "This is not valid JSON"
        
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        mock_prompt_instance.__or__ = lambda self, other: mock_chain
        mock_chain.invoke.return_value = mock_response
        mock_chat_anthropic.return_value = mock_model
        
        with pytest.raises(ValueError, match="Could not parse response as JSON"):
            custom_summary(sample_template, sample_transcription_dict)


class TestSchemaSwitching:
    """Test schema switching based on plan tier"""
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_schema_switching_freemium(
        self,
        mock_prompt_template,
        mock_chat_anthropic,
        sample_transcription_dict,
        sample_basic_summary_response
    ):
        """Test that freemium plan uses BasicSummary schema"""
        mock_model = MagicMock()
        mock_chain = MagicMock()
        from scripts.generate_summary import BasicSummary
        mock_result = BasicSummary(**sample_basic_summary_response)
        
        mock_structured_model = MagicMock()
        mock_structured_model.invoke.return_value = mock_result
        mock_model.with_structured_output = MagicMock(return_value=mock_structured_model)
        
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        mock_prompt_instance.__or__ = lambda self, other: mock_chain
        mock_chain.invoke.return_value = mock_result
        mock_chat_anthropic.return_value = mock_model
        
        general_summary(sample_transcription_dict, plan_tier="freemium")
        
        # Verify BasicSummary schema was used
        mock_model.with_structured_output.assert_called_once()
        call_args = mock_model.with_structured_output.call_args
        from scripts.generate_summary import BasicSummary
        assert call_args[0][0] == BasicSummary
    
    @patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
    @patch('scripts.generate_summary.ChatAnthropic')
    @patch('scripts.generate_summary.ChatPromptTemplate')
    def test_schema_switching_pro(
        self,
        mock_prompt_template,
        mock_chat_anthropic,
        sample_transcription_dict,
        sample_advanced_summary_response
    ):
        """Test that pro plan uses AdvancedSummary schema"""
        mock_model = MagicMock()
        mock_chain = MagicMock()
        from scripts.generate_summary import AdvancedSummary
        mock_result = AdvancedSummary(**sample_advanced_summary_response)
        
        mock_structured_model = MagicMock()
        mock_structured_model.invoke.return_value = mock_result
        mock_model.with_structured_output = MagicMock(return_value=mock_structured_model)
        
        mock_prompt_instance = MagicMock()
        mock_prompt_template.from_messages.return_value = mock_prompt_instance
        mock_prompt_instance.__or__ = lambda self, other: mock_chain
        mock_chain.invoke.return_value = mock_result
        mock_chat_anthropic.return_value = mock_model
        
        general_summary(sample_transcription_dict, plan_tier="pro")
        
        # Verify AdvancedSummary schema was used
        mock_model.with_structured_output.assert_called_once()
        call_args = mock_model.with_structured_output.call_args
        from scripts.generate_summary import AdvancedSummary
        assert call_args[0][0] == AdvancedSummary