import pytest
from src.preprocessing_funcs import split_text
from src.preprocessing_funcs import (
    remove_page_indicators,
    split_text,
    clean_law_symbols,
    filter_quick_review,
    get_exhibit_pages,
)

from src.client.smart_chunking_funcs import smart_chunk_ac

class TestPreprocessingFuncs:
    # Test cases for clean_newlines
    @pytest.mark.parametrize("input_text, expected_output", [
        # Removes page number indicators
        ("Page 1 of 10\n\nThis is a test.", " This is a test."),
        # Handles multiple newlines
        ("Line 1\n\nLine 2\n\nLine 3", "Line 1\n\nLine 2\n\nLine 3"),
        # Handles no newlines
        ("This is a single line.", "This is a single line."),
    ])
    def test_clean_newlines(self, input_text, expected_output):
        assert remove_page_indicators(input_text) == expected_output

    # Test cases for clean_law_symbols
    @pytest.mark.parametrize("input_text, expected_output", [
        # Replaces double dollar signs
        ("This is a $$ test.", "This is a $ test."),
        # Replaces U.S.C. symbol
        ("U.S.C. $1234", "U.S.C.§1234"),
        # Replaces C.F.R. symbol
        ("C.F.R. $1234", "C.F.R.§1234"),
        # Replaces $ followed by three decimal places
        ("$123.456", "§123.456"),
        # Handles no replacements
        ("This is a normal text.", "This is a normal text."),
        # Handles mixed replacements
        ("U.S.C. $1234 and C.F.R. $5678 and $123.456", "U.S.C.§1234 and C.F.R.§5678 and §123.456"),
    ])
    def test_clean_law_symbols(self, input_text, expected_output):
        assert clean_law_symbols(input_text) == expected_output

    # Test cases for smart_chunk_ac
    @pytest.mark.parametrize("methodology, mock_return, chunking_fn_name, expected_chunk, expected_pages", [
        # OR methodology
        ("or", [1, 2], "or", "Content1\nContent2", {"test": [1, 2]}),
        # Hierarchy methodology
        ("hierarchy", [1, 2], "hierarchical", "Content1\nContent2", {"test": [1, 2]}),
        # AND methodology
        ("and", [1], "and", "Content1", {"test": [1]}),
        # Regex methodology
        ("regex", [1, 2], "regex", "Content1\nContent2", {"test": [1, 2]}),
        # Include/Exclude methodology
        ("include_exclude", [1], "with_include_exclude_keywords", "Content1", {"test": [1]}),
    ])
    
    def test_smart_chunk_ac(self, mocker, methodology, mock_return, chunking_fn_name, expected_chunk, expected_pages):
        # Mock the chunking functions
        mock_chunk_func = mocker.patch(f"src.client.smart_chunking_funcs.chunk_{chunking_fn_name}", return_value=mock_return)   
        mock_keyword_search = mocker.patch("src.client.smart_chunking_funcs.keyword_search", return_value={"test": mock_return})

        text_dict = {"1": "Content1", "2": "Content2"}
        keyword_mappings = {"group1": {"methodology": methodology, "keywords": ["test"], "case_sensitive": False}}
        
        if methodology == "regex": keyword_mappings["group1"]["regex"] = "Content"
        
        if methodology == "include_exclude": 
            del keyword_mappings["group1"]["keywords"]
            keyword_mappings["group1"]["included_keywords"] = ["Content"]
            keyword_mappings["group1"]["excluded_keywords"] = []
        
        result = smart_chunk_ac(text_dict, "Content1\nContent2", keyword_mappings)
        
        assert result["group1"] == expected_chunk
        assert result["group1_pages"] == expected_pages
        mock_chunk_func.assert_called_once()
        mock_keyword_search.assert_called_once()
    
    # Test cases for filter_quick_review
    @pytest.mark.parametrize("input_dict, expected_contract, expected_quick_review", [
        # Filters quick review pages
        (
            {"1": "QUICK REVIEW", "2": "Normal content", "3": "COVER SHEET"},
            {"2": "Normal content"},
            {"1": "QUICK REVIEW", "3": "COVER SHEET"}
        ),
        # No quick review pages
        (
            {"1": "Normal content", "2": "More content"},
            {"1": "Normal content", "2": "More content"},
            {}
        ),
        # All quick review pages
        (
            {"1": "QUICK REVIEW", "2": "TOP SHEET"},
            {},
            {"1": "QUICK REVIEW", "2": "TOP SHEET"}
        ),
        # Case-insensitive matching
        (
            {"1": "quick review", "2": "Normal content"},
            {"2": "Normal content"},
            {"1": "quick review"}
        ),
    ])
    def test_filter_quick_review(self, input_dict, expected_contract, expected_quick_review):
        contract, quick_review = filter_quick_review(input_dict)
        assert contract == expected_contract
        assert quick_review == expected_quick_review

    # Test cases for get_exhibit_pages
    @pytest.mark.parametrize("mock_response, expected_pages, text_dict", [
        # First page included by default, second page not an exhibit
        (["|N|"], ["1"], {"1": "Exhibit content", "2": "Non-exhibit content"}),
        # Multiple pages with mixed exhibit status
        (["|N|", "|Y|"], ["1", "3"], {"1": "First page", "2": "Regular content", "3": "Exhibit B"}),
        # No exhibit page detected
        ([], [], {}),
        # Multiple exhibit pages
        (["|Y|","|Y|"], ["1", "2"], {"1": "Exhibit content", "2": "Attachment"}),
        # Invalid response; still include the first page
        (["Invalid"], ["1"], {"1": "Non-exhibit content"}),
    ])
    def test_get_exhibit_pages(self, mocker, mock_response, expected_pages, text_dict):
        # Mock the LLM response
        mocker_invoke_claude = mocker.patch("src.utils.llm_utils.invoke_claude", side_effect=mock_response)
        filename = "test.pdf"
        assert get_exhibit_pages(text_dict, filename) == expected_pages
        
        # Verify calls made - should be number of pages minus 1 (first page), but never negative
        expected_calls = max(0, len(text_dict) - 1)
        assert mocker_invoke_claude.call_count == expected_calls


from src.preprocessing_funcs import get_exhibit_dict

class TestGetExhibitDict:
    
    def test_case_1_short_table_first_page(self):
        """Case 1: Short Table on First Page"""
        text_dict = {
            '5': 'Exhibit 5 header text',
            '5.0': 'Short table content'
        }
        exhibit_chunk_mapping = {
            '5': '5',
            '5.0': '5'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.0': 'Exhibit 5 header text\nShort table content'
        }
        assert result == expected

    def test_case_2_continued_short_table(self):
        """Case 2: Continued Short Table on Subsequent Pages"""
        text_dict = {
            '5': 'Exhibit 5 header text',
            '5.0': 'Table start content',
            '5.1': 'Table continuation content'
        }
        exhibit_chunk_mapping = {
            '5': '5',
            '5.0': '5',
            '5.1': '5'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.0': 'Exhibit 5 header text\nTable start content',
            '5.1': 'Exhibit 5 header text\nTable continuation content'
        }
        assert result == expected

    def test_case_3_long_table_first_page(self):
        """Case 3: Long Table on First Page (split by row limit)"""
        text_dict = {
            '5': 'Exhibit 5 header text',
            '5.0': 'Long table part 1',
            '5.1': 'Long table part 2',
            '5.2': 'Long table part 3'
        }
        exhibit_chunk_mapping = {
            '5': '5',
            '5.0': '5',
            '5.1': '5',
            '5.2': '5'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.0': 'Exhibit 5 header text\nLong table part 1',
            '5.1': 'Exhibit 5 header text\nLong table part 2',
            '5.2': 'Exhibit 5 header text\nLong table part 3'
        }
        assert result == expected

    def test_case_4_continued_long_table(self):
        """Case 4: Continued Long Table on Subsequent Pages"""
        text_dict = {
            '5': 'Exhibit 5 header text',
            '6': 'Page 6 text',
            '6.0': 'Continued table part 1',
            '6.1': 'Continued table part 2',
            '7': 'Page 7 text'
        }
        exhibit_chunk_mapping = {
            '5': '5',
            '6': '5',
            '6.0': '5',
            '6.1': '5',
            '7': '5'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.0': 'Exhibit 5 header text\nPage 6 text\nContinued table part 1\nPage 7 text',
            '5.1': 'Exhibit 5 header text\nPage 6 text\nContinued table part 2\nPage 7 text',
            '5.2': 'Exhibit 5 header text\nPage 6 text\nPage 7 text'
        }
        assert result == expected

    def test_case_5_multiple_tables_first_page(self):
        """Case 5: Multiple Tables on First page only"""
        text_dict = {
            '5': 'Exhibit 5 header text',
            '5.0': 'First table content',
            '5.1': 'Second table content',
            '5.2': 'Third table content'
        }
        exhibit_chunk_mapping = {
            '5': '5',
            '5.0': '5',
            '5.1': '5',
            '5.2': '5'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.0': 'Exhibit 5 header text\nFirst table content',
            '5.1': 'Exhibit 5 header text\nSecond table content',
            '5.2': 'Exhibit 5 header text\nThird table content'
        }
        assert result == expected

    def test_case_6_multiple_tables_with_continuation(self):
        """Case 6: Multiple Tables on Subsequent Page, with first being continuation"""
        text_dict = {
            '5': 'Exhibit 5 header text',
            '6': 'Page 6 text',
            '6.0': 'Continuation of previous table',
            '6.1': 'New table on same page',
            '6.2': 'Another new table on same page'
        }
        exhibit_chunk_mapping = {
            '5': '5',
            '6': '5',
            '6.0': '5',
            '6.1': '5',
            '6.2': '5'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.0': 'Exhibit 5 header text\nPage 6 text\nContinuation of previous table',
            '5.1': 'Exhibit 5 header text\nPage 6 text\nNew table on same page',
            '5.2': 'Exhibit 5 header text\nPage 6 text\nAnother new table on same page'
        }
        assert result == expected

    def test_case_7_multiple_tables_across_pages(self):
        """Case 7: Multiple Tables on First and Subsequent Pages"""
        text_dict = {
            '5': 'Exhibit 5 header text',
            '5.0': 'First table on page 5',
            '5.1': 'Second table on page 5',
            '6': 'Page 6 text',
            '6.0': 'First table on page 6',
            '6.1': 'Second table on page 6',
            '7': 'Page 7 text'
        }
        exhibit_chunk_mapping = {
            '5': '5',
            '5.0': '5',
            '5.1': '5',
            '6': '5',
            '6.0': '5',
            '6.1': '5',
            '7': '5'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.0': 'Exhibit 5 header text\nFirst table on page 5\nPage 6 text\nPage 7 text',
            '5.1': 'Exhibit 5 header text\nSecond table on page 5\nPage 6 text\nPage 7 text',
            '5.2': 'Exhibit 5 header text\nPage 6 text\nFirst table on page 6\nPage 7 text',
            '5.3': 'Exhibit 5 header text\nPage 6 text\nSecond table on page 6\nPage 7 text',
            '5.4': 'Exhibit 5 header text\nPage 6 text\nPage 7 text'
        }
        assert result == expected

    def test_edge_case_exhibit_is_decimal_page(self):
        """Edge Case: Exhibit itself is a decimal page"""
        text_dict = {
            '5.1': 'Exhibit starting at decimal page'
        }
        exhibit_chunk_mapping = {
            '5.1': '5.1'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.1': 'Exhibit starting at decimal page'
        }
        assert result == expected

    def test_edge_case_no_decimal_pages(self):
        """Edge Case: No decimal pages in exhibit (pure text)"""
        text_dict = {
            '5': 'Exhibit 5 header text',
            '6': 'Page 6 text',
            '7': 'Page 7 text'
        }
        exhibit_chunk_mapping = {
            '5': '5',
            '6': '5',
            '7': '5'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '5.0': 'Exhibit 5 header text\nPage 6 text\nPage 7 text'
        }
        assert result == expected

    def test_example_from_docstring_1(self):
        """Test first example from docstring"""
        text_dict = {
            '1.0': 'Page 1.0 text',
            '2': 'Page 2 text',
            '3.0': 'Page 3.0 text',
            '3.1': 'Page 3.1 text',
            '3.2': 'Page 3.2 text',
            '4.0': 'Page 4.0 text'
        }
        exhibit_chunk_mapping = {
            '1.0': '1.0',
            '2': '2',
            '3.0': '2',
            '3.1': '2',
            '3.2': '2',
            '4.0': '2'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '1.0': 'Page 1.0 text',
            '2.0': 'Page 2 text\nPage 3.0 text',
            '2.1': 'Page 2 text\nPage 3.1 text',
            '2.2': 'Page 2 text\nPage 3.2 text',
            '2.3': 'Page 2 text\nPage 4.0 text'
        }
        assert result == expected

    def test_example_from_docstring_2(self):
        """Test second example from docstring"""
        text_dict = {
            '1.0': 'Page 1.0 text',
            '2': 'Page 2 text',
            '3.0': 'Page 3.0 text',
            '3.1': 'Page 3.1 text',
            '3.2': 'Page 3.2 text',
            '4': 'Page 4 text',
            '5': 'Page 5 text'
        }
        exhibit_chunk_mapping = {
            '1.0': '1.0',
            '2': '2',
            '3.0': '2',
            '3.1': '2',
            '3.2': '2',
            '4': '2',
            '5': '2'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            '1.0': 'Page 1.0 text',
            '2.0': 'Page 2 text\nPage 3.0 text\nPage 4 text\nPage 5 text',
            '2.1': 'Page 2 text\nPage 3.1 text\nPage 4 text\nPage 5 text',
            '2.2': 'Page 2 text\nPage 3.2 text\nPage 4 text\nPage 5 text',
            '2.3': 'Page 2 text\nPage 4 text\nPage 5 text'
        }
        assert result == expected

    def test_exhibit_with_decimal_pages(self):
        """Test exhibit with decimal pages"""
        text_dict = {
            "23.0": "Exhibit page 23 header",
            "24": "Page 24 text"
        }
        exhibit_chunk_mapping = {
            "23.0": "23.0",
            "24": "23.0"
        }
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            "23.0.0": "Exhibit page 23 header\nPage 24 text",
            "23.0.1": "Page 24 text"
        }
        assert result == expected

    def test_multiple_exhibits_with_tables(self):
        """Test multiple different exhibits, each with their own base pages and tables"""
        text_dict = {
            '10': 'Exhibit A header',
            '10.0': 'Exhibit A table 1',
            '10.1': 'Exhibit A table 2', 
            '11': 'Exhibit A page 2',
            
            '20': 'Exhibit B header',
            '21': 'Exhibit B page 2',
            '21.0': 'Exhibit B table 1',
            '22': 'Exhibit B page 3',
            
            '30': 'Exhibit C header',
            '31': 'Exhibit C page 2'
        }
        exhibit_chunk_mapping = {
            '10': 'A',
            '10.0': 'A', 
            '10.1': 'A',
            '11': 'A',
            
            '20': 'B',
            '21': 'B',
            '21.0': 'B', 
            '22': 'B',
            
            '30': 'C',
            '31': 'C'
        }
        
        result = get_exhibit_dict(text_dict, exhibit_chunk_mapping)
        expected = {
            # Exhibit A
            'A.0': 'Exhibit A header\nExhibit A table 1\nExhibit A page 2',
            'A.1': 'Exhibit A header\nExhibit A table 2\nExhibit A page 2', 
            'A.2': 'Exhibit A header\nExhibit A page 2',
            
            # Exhibit B  
            'B.0': 'Exhibit B header\nExhibit B page 2\nExhibit B table 1\nExhibit B page 3',
            'B.1': 'Exhibit B header\nExhibit B page 2\nExhibit B page 3',
            
            # Exhibit C (no decimal pages)
            'C.0': 'Exhibit C header\nExhibit C page 2'
        }
        assert result == expected