import pytest

from src.investment.preprocessing_funcs import (clean_law_symbols,
                                                filter_quick_review,
                                                remove_page_indicators)


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 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