doczyai-pipelines/src/tests/test_code_funcs.py

import json
import re
import unittest
from unittest.mock import MagicMock, patch

import pandas as pd

import src.codes.code_funcs as code_funcs
import src.prompts.prompt_templates as prompt_templates
from src.constants.constants import Constants
from src.constants.delimiters import Delimiter


class TestCodeFuncs(unittest.TestCase):

    def setUp(self):
        """Set up test fixtures before each test method."""
        # Clear the module-level valid code set caches between tests.
        code_funcs._VALID_PROC_CACHE.clear()
        code_funcs._VALID_REV_CACHE.clear()

        self.constants = MagicMock(spec=Constants)

        # Configure mock constants
        self.constants.SYNONYM_MAP = {
            "DME": "DURABLE MEDICAL EQUIPMENT",
            "SNF": "SKILLED NURSING FACILITY",
        }
        self.constants.REMOVAL_LIST = ["INPATIENT", "OUTPATIENT", "SERVICES"]
        self.constants.STOP_WORD_LIST = ["THE", "AND", "OF"]
        self.constants.DO_NOT_RUN = ["BY REPORT", "STOP-LOSS"]

        # Setup mappings
        self.constants.CPT_LEVEL1_MAPPING = {"12345": "Test Procedure"}
        self.constants.HCPCS_LEVEL1_MAPPING = {"J0001": "Test Drug"}
        self.constants.REV_LEVEL1_MAPPING = {"123": "Test Revenue"}

        self.constants.CPT_LEVEL2_MAPPING = {"67890": "Advanced Procedure"}
        self.constants.HCPCS_LEVEL2_MAPPING = {"J0002": "Advanced Drug"}
        self.constants.REV_MAPPING = {"0456": "Advanced Revenue"}
        # Full mappings for validation and description fill (Fix 1 / Fix 2)
        self.constants.CPT_FULL_MAPPING = {
            "12345": "Test Procedure",
            "99213": "Office visit, established patient",
        }
        self.constants.HCPCS_FULL_MAPPING = {
            "J0001": "Test Drug",
            "G0525": "Care management",
        }

        # Bill type mappings matching the actual JSON structure
        self.constants.BILL_TYPE_MAPPING = {
            "011X": "Inpatient Hospital",
            "012X": "Inpatient Hospital",
            "018X": "Inpatient Hospital",
        }
        # Reverse mapping returns a list of codes for each description
        self.constants.BILL_TYPE_REVERSE_MAPPING = {
            "Inpatient Hospital": ["011X", "012X", "018X"]
        }

        # Mock the embedding model and function
        mock_embedding = MagicMock()
        mock_embedding.index.search.return_value = ([0.9], [[0]])
        mock_embedding.choices = ["Test Description"]
        self.constants.get_embedding.return_value = mock_embedding

        self.constants.EMBEDDING_MODEL = MagicMock()
        mock_encoded = MagicMock()
        mock_encoded.astype.return_value = "mock_vector"
        self.constants.EMBEDDING_MODEL.encode.return_value = mock_encoded

    def tearDown(self):
        """Clear prompt_templates constants cache so a MagicMock is not reused by other tests."""
        prompt_templates.set_constants_for_instruction_cache(None)

    def test_clean_service(self):
        """Test the clean_service function with various inputs."""
        # Test empty input
        self.assertEqual(code_funcs.clean_service("", self.constants), "")
        self.assertEqual(code_funcs.clean_service(None, self.constants), "")

        # Test synonym mapping
        self.assertEqual(
            code_funcs.clean_service("DME", self.constants), "DURABLE MEDICAL EQUIPMENT"
        )

        # Test removal list
        self.assertEqual(
            code_funcs.clean_service("INPATIENT SERVICES", self.constants), ""
        )

        # Test stop words only
        self.assertEqual(code_funcs.clean_service("THE AND OF", self.constants), "")

        # Test unlisted service
        self.assertEqual(
            code_funcs.clean_service("UNLISTED PROCEDURE", self.constants), "UNLISTED"
        )

        # Test normal case
        self.assertEqual(
            code_funcs.clean_service("CT SCAN SNF", self.constants),
            "CT SCAN SKILLED NURSING FACILITY",
        )

    @patch("src.utils.llm_utils.invoke_claude")
    @patch("src.utils.string_utils.universal_json_load")
    def test_code_explicit(self, mock_json_load, mock_invoke_claude):
        """Test the code_explicit function for extracting explicit codes."""
        # Setup mocks
        mock_invoke_claude.return_value = '{"PROCEDURE_CD": "12345"}'
        mock_json_load.return_value = {"PROCEDURE_CD": "12345"}

        # Test successful extraction
        result = code_funcs.code_explicit(
            "TEST SERVICE", "TEST_METHODOLOGY", "test.pdf"
        )
        self.assertEqual(result, {"PROCEDURE_CD": "12345"})
        mock_invoke_claude.assert_called_once()

    @patch("src.utils.llm_utils.invoke_claude")
    @patch("src.utils.string_utils.universal_json_load")
    def test_code_category(self, mock_json_load, mock_invoke_claude):
        """Test the code_category function for detecting code categories."""
        # Setup mocks
        mock_invoke_claude.return_value = '["Drug A", "Drug B"]'
        mock_json_load.return_value = ["Drug A", "Drug B"]

        # Setup test data
        proc_category = ["Category: J"]
        hcpcs_mapping = {"J0001": "Drug A", "J0002": "Drug B", "K0001": "Other Drug"}

        # Test category matching
        result = code_funcs.code_category(
            "J CODES", proc_category, hcpcs_mapping, "test.pdf"
        )

        # Verify the results - code_category returns lists (JSON conversion happens in extract_codes_from_service)
        self.assertEqual(result["PROCEDURE_CD"], ["J0001", "J0002"])
        self.assertEqual(result["PROCEDURE_CD_DESC"], ["Drug A", "Drug B"])
        self.assertIn("J0001", result["PROCEDURE_CD"])
        self.assertIn("J0002", result["PROCEDURE_CD"])
        self.assertIn("Drug A", result["PROCEDURE_CD_DESC"])
        self.assertIn("Drug B", result["PROCEDURE_CD_DESC"])

    @patch("src.utils.llm_utils.invoke_claude")
    def test_code_implicit_special(self, mock_invoke_claude):
        """Test the code_implicit_special function for special case handling."""
        # Setup mocks - parser returns a list
        mock_invoke_claude.return_value = '["Drugs"]'

        # Test drug category - returns a list
        result = code_funcs.code_implicit_special("DRUG SERVICE", "test.pdf")
        self.assertEqual(result["PROCEDURE_CD"], ["J0000-J9999"])
        self.assertEqual(result["PROCEDURE_CD_DESC"], ["Drugs"])

        # Test no match - parser returns empty list
        mock_invoke_claude.return_value = "[]"
        result = code_funcs.code_implicit_special("OTHER SERVICE", "test.pdf")
        self.assertEqual(result, {})

    def test_get_embedding_levels(self):
        """Test the get_embedding_levels function for determining embedding levels."""
        # Test level 1 with both codes enabled
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True}
        embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = (
            code_funcs.get_embedding_levels(1, implicit_run_dict, self.constants)
        )

        self.assertIn("cpt_level1", embedding_levels)
        self.assertIn("hcpcs_level1", embedding_levels)
        self.assertIn("rev_level1", embedding_levels)
        self.assertEqual(cpt_mapping, self.constants.CPT_LEVEL1_MAPPING)

        # Test level 2 with only procedure code enabled
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": False}
        embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = (
            code_funcs.get_embedding_levels(2, implicit_run_dict, self.constants)
        )

        self.assertIn("cpt_level2", embedding_levels)
        self.assertIn("hcpcs_level2", embedding_levels)
        self.assertNotIn("rev", embedding_levels)
        self.assertEqual(cpt_mapping, self.constants.CPT_LEVEL2_MAPPING)

    def test_get_match_list(self):
        """Test the get_match_list function for retrieving matches from embeddings."""
        # Configure mock for proper array-like return values
        mock_embedding = MagicMock()
        # Search returns a tuple of (scores, indices) where each is a 2D array
        mock_embedding.index.search.return_value = (
            # Scores array - 2D array with shape [1, top_k]
            [[0.9]],
            # Indices array - 2D array with shape [1, top_k]
            [[0]],
        )
        mock_embedding.choices = ["Test Description"]

        # Replace the mock embedding in constants for this test
        self.constants.get_embedding.return_value = mock_embedding

        # Test with constants mock configured in setUp
        match_list, highest_similarity = code_funcs.get_match_list(
            "TEST SERVICE", ["cpt_level1"], self.constants, 1
        )

        # Verify results
        self.assertEqual(match_list, ["Test Description"])
        self.assertEqual(highest_similarity, 0.9)

        # Verify the appropriate method calls
        self.constants.EMBEDDING_MODEL.encode.assert_called_once()
        self.constants.get_embedding.assert_called_once_with("cpt_level1")

    @patch("src.utils.llm_utils.invoke_claude")
    @patch("src.utils.string_utils.universal_json_load")
    @patch("src.codes.code_funcs.get_match_list")
    def test_code_implicit_rag(
        self, mock_get_match_list, mock_json_load, mock_invoke_claude
    ):
        """Test the code_implicit_rag function for RAG-based code extraction."""
        # Mock get_match_list to return properly formatted results
        mock_get_match_list.return_value = (["Test Procedure"], 0.9)

        # Setup mocks for Claude responses
        mock_invoke_claude.return_value = '["Test Procedure"]'
        mock_json_load.return_value = ["Test Procedure"]

        # Test with successful match
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True}

        # Patch the get_embedding_levels function to return predictable values
        with patch(
            "src.codes.code_funcs.get_embedding_levels"
        ) as mock_get_embedding_levels:
            # Configure mock to return appropriate values for embedding levels
            mock_get_embedding_levels.return_value = (
                ["cpt_level1"],
                {"12345": "Test Procedure"},
                {},
                {},
            )

            # Call the function
            result = code_funcs.code_implicit_rag(
                "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants
            )

            # Verify results - code_implicit_rag returns list of candidates with source Level 1/Level 2
            self.assertIsInstance(result, list)
            self.assertGreaterEqual(len(result), 1)
            codes_found = False
            for item in result:
                self.assertIn("source", item)
                self.assertIn("code_answer_dict", item)
                d = item["code_answer_dict"]
                if d.get("PROCEDURE_CD") == ["12345"]:
                    codes_found = True
                    self.assertEqual(d["PROCEDURE_CD_DESC"], ["Test Procedure"])
                    break
            self.assertTrue(
                codes_found, "Expected PROCEDURE_CD ['12345'] in a candidate"
            )

        # Test exception handling - returns list (possibly empty)
        mock_json_load.side_effect = Exception("Test error")
        result = code_funcs.code_implicit_rag(
            "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants
        )
        self.assertIsInstance(result, list)

    @patch("src.utils.llm_utils.invoke_claude")
    @patch("src.utils.string_utils.universal_json_load")
    @patch("src.codes.code_funcs.get_match_list")
    def test_code_implicit_rag_invalid_service_short_circuits(
        self, mock_get_match_list, mock_json_load, mock_invoke_claude
    ):
        """INVALID_SERVICE in RAG response is a hard stop: no codes added."""
        mock_get_match_list.return_value = (["Test Procedure"], 0.9)
        mock_invoke_claude.return_value = '["INVALID_SERVICE", "Test Procedure"]'
        mock_json_load.return_value = ["INVALID_SERVICE", "Test Procedure"]
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True}

        with patch(
            "src.codes.code_funcs.get_embedding_levels"
        ) as mock_get_embedding_levels:
            mock_get_embedding_levels.return_value = (
                ["cpt_level1"],
                {"12345": "Test Procedure"},
                {},
                {},
            )
            result = code_funcs.code_implicit_rag(
                "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants
            )

        # INVALID_SERVICE yields no candidates with codes
        self.assertIsInstance(result, list)
        self.assertEqual(len(result), 0)

    @patch("src.utils.llm_utils.invoke_claude")
    def test_fill_bill_type_list_reverse_mapping(self, mock_invoke_claude):
        """fill_bill_type handles list values in BILL_TYPE_REVERSE_MAPPING."""
        mock_invoke_claude.return_value = '["Inpatient Hospital"]'
        reverse_mapping = {"Inpatient Hospital": ["11X", "12X"]}
        answer_dict = {}
        result = code_funcs.fill_bill_type(
            "INPATIENT SERVICE",
            answer_dict,
            self.constants.BILL_TYPE_MAPPING,
            reverse_mapping,
        )
        self.assertEqual(result["BILL_TYPE_CD"], ["11X", "12X"])

    @patch("src.utils.llm_utils.invoke_claude")
    def test_code_last_check(self, mock_invoke_claude):
        """Test the code_last_check function for categorizing non-matched services."""
        # Parser is format-aware and returns 'Specific' or 'Generic' directly.
        # Test specific case
        mock_invoke_claude.return_value = '["Specific"]'
        result = code_funcs.code_last_check("SPECIAL SERVICE", "test.pdf")
        self.assertEqual(result, "Specific")

        # Test generic case
        mock_invoke_claude.return_value = '["Generic"]'
        result = code_funcs.code_last_check("GENERIC SERVICE", "test.pdf")
        self.assertEqual(result, "Generic")

    def test_format_implicit_candidate_for_prompt(self):
        """Test _format_implicit_candidate_for_prompt formats a candidate for arbitration."""
        candidate = {
            "source": "Category",
            "code_answer_dict": {
                "PROCEDURE_CD": ["J0001", "J0002"],
                "PROCEDURE_CD_DESC": ["Drug A", "Drug B"],
            },
        }
        result = code_funcs._format_implicit_candidate_for_prompt(candidate, 0)
        self.assertIn("Candidate 0 (Source: Category)", result)
        self.assertIn("PROCEDURE_CD", result)
        self.assertIn("J0001", result)

    @patch("src.codes.code_funcs.code_implicit_rag")
    @patch("src.codes.code_funcs.code_implicit_special")
    @patch("src.codes.code_funcs.code_category")
    def test_build_implicit_candidates_empty_when_none_return(
        self, mock_category, mock_special, mock_rag
    ):
        """Test build_implicit_candidates returns empty list when no strategy returns."""
        mock_category.return_value = {}
        mock_special.return_value = {}
        mock_rag.return_value = []  # RAG returns list of candidates
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True}
        result = code_funcs.build_implicit_candidates(
            "SERVICE",
            {},
            implicit_run_dict,
            "test.pdf",
            self.constants,
        )
        self.assertEqual(result, [])

    @patch("src.codes.code_funcs.code_implicit_rag")
    @patch("src.codes.code_funcs.code_implicit_special")
    @patch("src.codes.code_funcs.code_category")
    def test_build_implicit_candidates_category_when_explicit_has_category(
        self, mock_category, mock_special, mock_rag
    ):
        """Test build_implicit_candidates runs category and adds candidate when explicit had Category."""
        mock_category.return_value = {
            "PROCEDURE_CD": ["J0001"],
            "PROCEDURE_CD_DESC": ["Drug A"],
        }
        mock_special.return_value = {}
        mock_rag.return_value = []  # RAG returns list of candidates
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True}
        code_after_explicit = {"PROCEDURE_CD": ["Category: J"]}
        result = code_funcs.build_implicit_candidates(
            "J CODES",
            code_after_explicit,
            implicit_run_dict,
            "test.pdf",
            self.constants,
        )
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["source"], "Category")
        self.assertEqual(result[0]["code_answer_dict"]["PROCEDURE_CD"], ["J0001"])
        mock_category.assert_called_once()

    @patch("src.codes.code_funcs.code_implicit_rag")
    @patch("src.codes.code_funcs.code_implicit_special")
    @patch("src.codes.code_funcs.code_category")
    def test_build_implicit_candidates_skips_category_without_category_in_explicit(
        self, mock_category, mock_special, mock_rag
    ):
        """Test build_implicit_candidates does not run category when explicit has no Category."""
        mock_special.return_value = {
            "PROCEDURE_CD": ["J0000-J9999"],
            "PROCEDURE_CD_DESC": ["Drugs"],
        }
        mock_rag.return_value = []  # RAG returns list of candidates
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True}
        result = code_funcs.build_implicit_candidates(
            "DRUG SERVICE",
            {},
            implicit_run_dict,
            "test.pdf",
            self.constants,
        )
        mock_category.assert_not_called()
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["source"], "Special")

    @patch("src.codes.code_funcs.code_implicit_rag")
    @patch("src.codes.code_funcs.code_implicit_special")
    @patch("src.codes.code_funcs.code_category")
    def test_build_implicit_candidates_collects_category_special_rag(
        self, mock_category, mock_special, mock_rag
    ):
        """Test build_implicit_candidates collects candidates from all three when all return."""
        mock_category.return_value = {
            "PROCEDURE_CD": ["A0000-A9999"],
            "PROCEDURE_CD_DESC": ["A range"],
        }
        mock_special.return_value = {
            "PROCEDURE_CD": ["J0000-J9999"],
            "PROCEDURE_CD_DESC": ["Drugs"],
        }
        mock_rag.return_value = [
            {
                "source": "Level 2",
                "code_answer_dict": {
                    "PROCEDURE_CD": ["99213"],
                    "PROCEDURE_CD_DESC": ["Office visit"],
                    "CODE_METHODOLOGY": "Implicit - Level 2",
                },
            },
        ]
        implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True}
        code_after_explicit = {"PROCEDURE_CD": ["Category: A"]}
        result = code_funcs.build_implicit_candidates(
            "SERVICE",
            code_after_explicit,
            implicit_run_dict,
            "test.pdf",
            self.constants,
        )
        self.assertEqual(len(result), 3)
        sources = [c["source"] for c in result]
        self.assertIn("Category", sources)
        self.assertIn("Special", sources)
        self.assertIn("Level 2", sources)

    def test_code_implicit_arbitration_zero_candidates_returns_none(self):
        """Test code_implicit_arbitration returns None when candidates list is empty."""
        result = code_funcs.code_implicit_arbitration("SERVICE", [], "test.pdf")
        self.assertIsNone(result)

    def test_code_implicit_arbitration_single_candidate_returns_it_no_llm(self):
        """Test code_implicit_arbitration with one candidate returns it without calling LLM."""
        candidates = [
            {
                "source": "Category",
                "code_answer_dict": {
                    "PROCEDURE_CD": ["J0001"],
                    "PROCEDURE_CD_DESC": ["Drug A"],
                },
            }
        ]
        result = code_funcs.code_implicit_arbitration("J CODES", candidates, "test.pdf")
        self.assertIsNotNone(result)
        self.assertEqual(result["PROCEDURE_CD"], ["J0001"])
        self.assertEqual(
            result["CODE_METHODOLOGY"],
            "Implicit - Arbitration (Category)",
        )

    @patch("src.utils.llm_utils.invoke_claude")
    def test_code_implicit_arbitration_multiple_candidates_uses_chosen_index(
        self, mock_invoke_claude
    ):
        """Test code_implicit_arbitration with 2+ candidates calls LLM and returns chosen candidate."""
        mock_invoke_claude.return_value = '{"chosen_index": 1}'
        candidates = [
            {
                "source": "Category",
                "code_answer_dict": {"PROCEDURE_CD": ["A0000-A9999"]},
            },
            {"source": "Level 2", "code_answer_dict": {"PROCEDURE_CD": ["99213"]}},
        ]
        result = code_funcs.code_implicit_arbitration(
            "Office visit", candidates, "test.pdf"
        )
        self.assertIsNotNone(result)
        self.assertEqual(result["PROCEDURE_CD"], ["99213"])
        self.assertEqual(
            result["CODE_METHODOLOGY"],
            "Implicit - Arbitration (Level 2)",
        )
        mock_invoke_claude.assert_called_once()

    @patch("src.utils.llm_utils.invoke_claude")
    def test_code_implicit_arbitration_no_match_returns_none(self, mock_invoke_claude):
        """Test code_implicit_arbitration returns None when LLM returns no_match."""
        mock_invoke_claude.return_value = '{"no_match": true}'
        candidates = [
            {
                "source": "Category",
                "code_answer_dict": {"PROCEDURE_CD": ["A0000-A9999"]},
            },
            {
                "source": "Special",
                "code_answer_dict": {"PROCEDURE_CD": ["10004-69990"]},
            },
        ]
        result = code_funcs.code_implicit_arbitration(
            "Unclear service", candidates, "test.pdf"
        )
        self.assertIsNone(result)

    @patch("src.utils.llm_utils.invoke_claude")
    def test_code_implicit_arbitration_invalid_index_fallback_to_zero(
        self, mock_invoke_claude
    ):
        """Test code_implicit_arbitration falls back to index 0 when chosen_index out of range."""
        mock_invoke_claude.return_value = '{"chosen_index": 99}'
        candidates = [
            {"source": "Category", "code_answer_dict": {"PROCEDURE_CD": ["J0001"]}},
        ]
        result = code_funcs.code_implicit_arbitration("SERVICE", candidates, "test.pdf")
        self.assertIsNotNone(result)
        self.assertEqual(result["PROCEDURE_CD"], ["J0001"])
        self.assertEqual(
            result["CODE_METHODOLOGY"],
            "Implicit - Arbitration (Category)",
        )

    @patch("src.utils.llm_utils.invoke_claude")
    def test_fill_bill_type(self, mock_invoke_claude):
        """Test the fill_bill_type function for populating bill type information."""
        mock_invoke_claude.return_value = '["Inpatient Hospital"]'

        # Test with valid bill type
        answer_dict = {}
        result = code_funcs.fill_bill_type(
            "INPATIENT SERVICE",
            answer_dict,
            self.constants.BILL_TYPE_MAPPING,
            self.constants.BILL_TYPE_REVERSE_MAPPING,
        )

        # Should return all codes for "Inpatient Hospital" description
        # Reverse mapping returns ["011X", "012X", "018X"] for "Inpatient Hospital"
        self.assertEqual(result["BILL_TYPE_CD"], ["011X", "012X", "018X"])
        self.assertEqual(result["BILL_TYPE_CD_DESC"], ["Inpatient Hospital"])

        # Test with empty response
        mock_invoke_claude.return_value = "[]"
        answer_dict = {"SERVICE_TERM": "Test"}
        result = code_funcs.fill_bill_type(
            "OTHER SERVICE",
            answer_dict,
            self.constants.BILL_TYPE_MAPPING,
            self.constants.BILL_TYPE_REVERSE_MAPPING,
        )

        self.assertEqual(result, {"SERVICE_TERM": "Test"})

    def test_get_implicit_runs(self):
        """Test the get_implicit_runs function for determining which code runs to execute."""
        # Test hospital claim type
        result = code_funcs.get_implicit_runs({"AARETE_DERIVED_CLAIM_TYPE_CD": "H"})
        self.assertTrue(result["REVENUE_CD"])

        # Test medical claim type with bill type
        result = code_funcs.get_implicit_runs(
            {"AARETE_DERIVED_CLAIM_TYPE_CD": "M", "BILL_TYPE_CD_DESC": "Outpatient"}
        )
        self.assertTrue(result["REVENUE_CD"])
        self.assertTrue(result["PROCEDURE_CD"])

        # Test inpatient bill type
        result = code_funcs.get_implicit_runs(
            {
                "AARETE_DERIVED_CLAIM_TYPE_CD": "M",
                "BILL_TYPE_CD_DESC": "Inpatient Hospital",
            }
        )
        self.assertTrue(result["REVENUE_CD"])
        self.assertFalse(result["PROCEDURE_CD"])

    @patch("src.codes.code_funcs.clean_service")
    @patch("src.codes.code_funcs.fill_bill_type")
    @patch("src.codes.code_funcs.get_implicit_runs")
    @patch("src.codes.code_funcs.code_explicit")
    @patch("src.codes.code_funcs.code_category")
    @patch("src.codes.code_funcs.code_implicit_special")
    @patch("src.codes.code_funcs.code_implicit_rag")
    @patch("src.codes.code_funcs.code_last_check")
    def test_extract_codes_from_service(
        self,
        mock_last_check,
        mock_implicit_rag,
        mock_implicit_special,
        mock_category,
        mock_explicit,
        mock_implicit_runs,
        mock_fill_bill_type,
        mock_clean_service,
    ):
        """Test the extract_codes_from_service function end-to-end."""
        # Setup mocks
        mock_clean_service.return_value = "CLEAN SERVICE"
        mock_fill_bill_type.return_value = {
            "SERVICE_TERM": "TEST SERVICE",
            "BILL_TYPE_CD_DESC": "Outpatient",
        }
        mock_implicit_runs.return_value = {"PROCEDURE_CD": True, "REVENUE_CD": True}

        # Test explicit code match - extract_codes_from_service normalizes to JSON list format
        mock_explicit.return_value = {"PROCEDURE_CD": "12345"}
        result = code_funcs.extract_codes_from_service(
            {"SERVICE_TERM": "TEST SERVICE"}, self.constants
        )
        self.assertEqual(result["PROCEDURE_CD"], '["12345"]')
        self.assertEqual(result["CODE_METHODOLOGY"], "Explicit")

        # Test unlisted service
        mock_clean_service.return_value = "UNLISTED"
        result = code_funcs.extract_codes_from_service(
            {"SERVICE_TERM": "UNLISTED SERVICE"}, self.constants
        )
        self.assertEqual(result["PROCEDURE_CD_DESC"], "UNLISTED")

        # Test category match (arbitration with single candidate returns it)
        mock_clean_service.return_value = "CLEAN SERVICE"
        mock_explicit.return_value = {"PROCEDURE_CD": ["Category: J"]}
        mock_category.return_value = {
            "PROCEDURE_CD": ["J0001"],
            "PROCEDURE_CD_DESC": ["Drug A"],
        }
        result = code_funcs.extract_codes_from_service(
            {"SERVICE_TERM": "J CODES"}, self.constants
        )
        self.assertEqual(
            result["CODE_METHODOLOGY"], "Implicit - Arbitration (Category)"
        )

        # Test special case match (arbitration with single candidate)
        mock_explicit.return_value = {}
        mock_implicit_special.return_value = {
            "PROCEDURE_CD": "J0000-J9999",
            "PROCEDURE_CD_DESC": "Drugs",
        }
        result = code_funcs.extract_codes_from_service(
            {"SERVICE_TERM": "DRUG SERVICE"}, self.constants
        )
        # Range J0000-J9999 not in mock mapping so methodology gets - Unmatched
        self.assertIn("Implicit - Arbitration (Special)", result["CODE_METHODOLOGY"])
        self.assertIn("Unmatched", result["CODE_METHODOLOGY"])

        # Test RAG match (arbitration with single candidate; RAG returns list with Level 1)
        mock_implicit_special.return_value = {}
        mock_implicit_rag.return_value = [
            {
                "source": "Level 1",
                "code_answer_dict": {
                    "PROCEDURE_CD": "12345",
                    "CODE_METHODOLOGY": "Implicit - Level 1",
                },
            },
        ]
        result = code_funcs.extract_codes_from_service(
            {"SERVICE_TERM": "TEST SERVICE"}, self.constants
        )
        self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Arbitration (Level 1)")

        # Test no match - specific
        mock_implicit_rag.return_value = {}
        mock_last_check.return_value = "Specific"
        result = code_funcs.extract_codes_from_service(
            {"SERVICE_TERM": "TEST SERVICE"}, self.constants
        )
        self.assertEqual(result["CODE_METHODOLOGY"], "Specific - No Match")

        # Test no match - generic
        mock_last_check.return_value = "Generic"
        result = code_funcs.extract_codes_from_service(
            {"SERVICE_TERM": "TEST SERVICE"}, self.constants
        )
        self.assertEqual(result["CODE_METHODOLOGY"], "Generic - No Match")

    def test_code_in_range(self):
        """Test _code_in_range for numeric and alphanumeric range keys."""
        self.assertTrue(code_funcs._code_in_range("00100-01999", "01234"))
        self.assertTrue(code_funcs._code_in_range("00100-01999", "00100"))
        self.assertTrue(code_funcs._code_in_range("00100-01999", "01999"))
        self.assertFalse(code_funcs._code_in_range("00100-01999", "99999"))
        self.assertFalse(code_funcs._code_in_range("00100-01999", "02000"))
        self.assertTrue(code_funcs._code_in_range("A0021-A0999", "A0500"))
        self.assertFalse(code_funcs._code_in_range("A0021-A0999", "B0001"))
        self.assertFalse(code_funcs._code_in_range("", "12345"))
        self.assertFalse(code_funcs._code_in_range("00100-01999", ""))

    def test_lookup_procedure_description(self):
        """Test _lookup_procedure_description: exact match and range match."""
        desc = code_funcs._lookup_procedure_description("12345", self.constants)
        self.assertEqual(desc, "Test Procedure")
        desc = code_funcs._lookup_procedure_description("99213", self.constants)
        self.assertEqual(desc, "Office visit, established patient")
        desc = code_funcs._lookup_procedure_description(
            "NOT_ESTABLISHED", self.constants
        )
        self.assertEqual(desc, "")
        desc = code_funcs._lookup_procedure_description("UNMAPPED99", self.constants)
        self.assertEqual(desc, "")

    def test_lookup_revenue_description(self):
        """Test _lookup_revenue_description."""
        desc = code_funcs._lookup_revenue_description("0456", self.constants)
        self.assertEqual(desc, "Advanced Revenue")
        desc = code_funcs._lookup_revenue_description("9999", self.constants)
        self.assertEqual(desc, "")

    def test_fill_code_descriptions_from_mappings(self):
        """Test fill_code_descriptions_from_mappings fills _DESC from mappings."""
        answer_dict = {
            "PROCEDURE_CD": ["12345", "99213"],
            "REVENUE_CD": ["0456"],
        }
        code_funcs.fill_code_descriptions_from_mappings(answer_dict, self.constants)
        self.assertEqual(
            answer_dict["PROCEDURE_CD_DESC"],
            ["Test Procedure", "Office visit, established patient"],
        )
        self.assertEqual(answer_dict["REVENUE_CD_DESC"], ["Advanced Revenue"])

    def test_fill_code_descriptions_leave_unmapped_unchanged(self):
        """Test fill leaves _DESC empty for unmapped codes and NOT_ESTABLISHED."""
        answer_dict = {
            "PROCEDURE_CD": ["12345", "NOT_ESTABLISHED", "UNMAPPED99"],
            "PROCEDURE_CD_DESC": ["", "", ""],
        }
        code_funcs.fill_code_descriptions_from_mappings(answer_dict, self.constants)
        self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][0], "Test Procedure")
        self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][1], "")
        self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][2], "")

    def test_validate_explicit_codes_keeps_valid_filters_invalid(self):
        """Format-invalid codes dropped; format-valid preserved; has_unmatched if not in mapping."""
        code_answer_dict = {
            "PROCEDURE_CD": ["12345", "99213", "garbage7"],
            "REVENUE_CD": ["0456", "9999"],
        }
        filtered, has_unmatched = code_funcs.validate_explicit_codes(
            code_answer_dict, self.constants
        )
        self.assertTrue(has_unmatched)
        self.assertEqual(filtered["PROCEDURE_CD"], ["12345", "99213"])
        self.assertEqual(filtered["REVENUE_CD"], ["0456", "9999"])

    def test_validate_explicit_codes_keeps_category_and_not_established(self):
        """Test validate_explicit_codes keeps Category: and NOT_ESTABLISHED; no unmatched."""
        code_answer_dict = {"PROCEDURE_CD": ["Category: J", "NOT_ESTABLISHED"]}
        filtered, has_unmatched = code_funcs.validate_explicit_codes(
            code_answer_dict, self.constants
        )
        self.assertFalse(has_unmatched)
        self.assertEqual(filtered["PROCEDURE_CD"], ["Category: J", "NOT_ESTABLISHED"])

    def test_validate_explicit_codes_all_valid(self):
        """Test validate_explicit_codes returns has_unmatched False when all in mapping."""
        code_answer_dict = {"PROCEDURE_CD": ["12345"], "REVENUE_CD": ["0456"]}
        filtered, has_unmatched = code_funcs.validate_explicit_codes(
            code_answer_dict, self.constants
        )
        self.assertFalse(has_unmatched)
        self.assertEqual(filtered["PROCEDURE_CD"], ["12345"])
        self.assertEqual(filtered["REVENUE_CD"], ["0456"])

    def test_validate_explicit_codes_multiple_and_range(self):
        """Range key that is an explicit key in mapping is kept and counts as mapped."""
        self.constants.CPT_LEVEL1_MAPPING["00100-01999"] = "Anesthesia"
        code_answer_dict = {
            "PROCEDURE_CD": ["12345", "99213", "00100-01999"],
            "REVENUE_CD": ["0456"],
        }
        filtered, has_unmatched = code_funcs.validate_explicit_codes(
            code_answer_dict, self.constants
        )
        self.assertFalse(has_unmatched)
        self.assertEqual(filtered["PROCEDURE_CD"], ["12345", "99213", "00100-01999"])
        self.assertEqual(filtered["REVENUE_CD"], ["0456"])

    def test_validate_explicit_codes_range_not_in_mapping_kept_tagged_unmatched(self):
        """Format-valid range not in mapping is kept; has_unmatched True for tag."""
        # Mock has 99213 in CPT_FULL_MAPPING but not 10021-69990; range is kept and tagged
        code_answer_dict = {
            "PROCEDURE_CD": ["99213", "10021-69990"],
            "REVENUE_CD": ["0456"],
        }
        filtered, has_unmatched = code_funcs.validate_explicit_codes(
            code_answer_dict, self.constants
        )
        self.assertTrue(has_unmatched)
        self.assertEqual(filtered["PROCEDURE_CD"], ["99213", "10021-69990"])
        self.assertEqual(filtered["REVENUE_CD"], ["0456"])

    def test_validate_explicit_codes_single_code_unmapped(self):
        """Format-valid code not in exact mapping is preserved and has_unmatched True."""
        code_answer_dict = {"PROCEDURE_CD": ["01234"], "REVENUE_CD": []}
        filtered, has_unmatched = code_funcs.validate_explicit_codes(
            code_answer_dict, self.constants
        )
        self.assertTrue(has_unmatched)
        self.assertEqual(filtered["PROCEDURE_CD"], ["01234"])

    def test_validate_explicit_codes_rejects_icd10_pcs_in_procedure(self):
        """ICD-10-PCS (7-char) filtered from PROCEDURE_CD; CPT/HCPCS format kept; S2900 in mapping."""
        self.constants.HCPCS_FULL_MAPPING["S2900"] = "Robotic surgical system"
        code_answer_dict = {
            "PROCEDURE_CD": ["0E0V0CZ", "0E0W0CZ", "S2900"],
            "REVENUE_CD": [],
        }
        filtered, has_unmatched = code_funcs.validate_explicit_codes(
            code_answer_dict, self.constants
        )
        self.assertFalse(has_unmatched)
        self.assertEqual(filtered["PROCEDURE_CD"], ["S2900"])

    def test_has_unmatched_codes_true_for_unmapped_procedure(self):
        """_has_unmatched_codes returns True when procedure code is format-valid but not in mapping."""
        d = {"PROCEDURE_CD": ["99999"], "REVENUE_CD": []}
        self.assertTrue(code_funcs._has_unmatched_codes(d, self.constants))

    def test_has_unmatched_codes_true_for_unmapped_revenue(self):
        """_has_unmatched_codes returns True when revenue code is format-valid but not in mapping."""
        d = {"PROCEDURE_CD": [], "REVENUE_CD": ["9999"]}
        self.assertTrue(code_funcs._has_unmatched_codes(d, self.constants))

    def test_has_unmatched_codes_false_when_all_mapped(self):
        """_has_unmatched_codes returns False when all procedure/revenue codes are in mapping."""
        d = {"PROCEDURE_CD": ["99213"], "REVENUE_CD": ["0456"]}
        self.assertFalse(code_funcs._has_unmatched_codes(d, self.constants))

    @patch("src.codes.code_funcs.code_explicit")
    def test_extract_codes_explicit_unmatched(self, mock_explicit):
        """Explicit path sets Unmatched when any format-valid code not in mapping; no retry."""
        mock_explicit.return_value = {
            "PROCEDURE_CD": ["12345", "99999"],
            "REVENUE_CD": [],
        }
        with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"):
            with patch("src.codes.code_funcs.fill_bill_type", return_value={}):
                with patch(
                    "src.codes.code_funcs.get_implicit_runs",
                    return_value={"PROCEDURE_CD": True, "REVENUE_CD": True},
                ):
                    result = code_funcs.extract_codes_from_service(
                        {"SERVICE_TERM": "TEST"}, self.constants
                    )
        self.assertEqual(result["CODE_METHODOLOGY"], "Explicit - Unmatched")
        self.assertIn("12345", str(result["PROCEDURE_CD"]))
        self.assertIn("99999", str(result["PROCEDURE_CD"]))
        self.assertEqual(mock_explicit.call_count, 1)

    @patch("src.codes.code_funcs.code_explicit")
    def test_extract_codes_explicit_description_filled(self, mock_explicit):
        """Test explicit path fills PROCEDURE_CD_DESC from mapping."""
        mock_explicit.return_value = {
            "PROCEDURE_CD": ["99213"],
            "REVENUE_CD": ["0456"],
        }
        with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"):
            with patch("src.codes.code_funcs.fill_bill_type", return_value={}):
                with patch(
                    "src.codes.code_funcs.get_implicit_runs",
                    return_value={"PROCEDURE_CD": True, "REVENUE_CD": True},
                ):
                    result = code_funcs.extract_codes_from_service(
                        {"SERVICE_TERM": "Office visit"}, self.constants
                    )
        self.assertEqual(result["CODE_METHODOLOGY"], "Explicit")
        self.assertIn("99213", str(result["PROCEDURE_CD"]))
        self.assertIn("Office visit", str(result["PROCEDURE_CD_DESC"]))
        self.assertIn("Advanced Revenue", str(result["REVENUE_CD_DESC"]))

    @patch("src.codes.code_funcs.code_implicit_arbitration")
    @patch("src.codes.code_funcs.build_implicit_candidates")
    @patch("src.codes.code_funcs.code_explicit")
    def test_extract_codes_implicit_arbitration_unmatched(
        self, mock_explicit, mock_build_candidates, mock_arbitration
    ):
        """Implicit path appends ' - Unmatched' when chosen candidate has unmapped code."""
        mock_explicit.return_value = {
            "PROCEDURE_CD": ["Category: J"],
            "REVENUE_CD": [],
        }
        mock_build_candidates.return_value = [
            {
                "source": "RAG",
                "code_answer_dict": {
                    "PROCEDURE_CD": ["99999"],
                    "REVENUE_CD": [],
                },
            }
        ]
        chosen = {
            "PROCEDURE_CD": ["99999"],
            "REVENUE_CD": [],
            "CODE_METHODOLOGY": "Implicit - Arbitration (RAG)",
        }
        mock_arbitration.return_value = chosen
        with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"):
            with patch("src.codes.code_funcs.fill_bill_type", return_value={}):
                with patch(
                    "src.codes.code_funcs.get_implicit_runs",
                    return_value={"PROCEDURE_CD": True, "REVENUE_CD": True},
                ):
                    result = code_funcs.extract_codes_from_service(
                        {"SERVICE_TERM": "J CODES"}, self.constants
                    )
        self.assertEqual(
            result["CODE_METHODOLOGY"], "Implicit - Arbitration (RAG) - Unmatched"
        )
        self.assertIn("99999", str(result["PROCEDURE_CD"]))

    @patch("src.codes.code_funcs.code_implicit_arbitration")
    @patch("src.codes.code_funcs.build_implicit_candidates")
    @patch("src.codes.code_funcs.code_explicit")
    def test_extract_codes_implicit_arbitration_all_mapped(
        self, mock_explicit, mock_build_candidates, mock_arbitration
    ):
        """Implicit path does not append ' - Unmatched' when all codes are in mapping."""
        mock_explicit.return_value = {
            "PROCEDURE_CD": ["Category: J"],
            "REVENUE_CD": [],
        }
        mock_build_candidates.return_value = [
            {
                "source": "RAG",
                "code_answer_dict": {
                    "PROCEDURE_CD": ["99213"],
                    "REVENUE_CD": ["0456"],
                },
            }
        ]
        chosen = {
            "PROCEDURE_CD": ["99213"],
            "REVENUE_CD": ["0456"],
            "CODE_METHODOLOGY": "Implicit - Arbitration (RAG)",
        }
        mock_arbitration.return_value = chosen
        with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"):
            with patch("src.codes.code_funcs.fill_bill_type", return_value={}):
                with patch(
                    "src.codes.code_funcs.get_implicit_runs",
                    return_value={"PROCEDURE_CD": True, "REVENUE_CD": True},
                ):
                    result = code_funcs.extract_codes_from_service(
                        {"SERVICE_TERM": "J CODES"}, self.constants
                    )
        self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Arbitration (RAG)")
        self.assertIn("99213", str(result["PROCEDURE_CD"]))
        self.assertIn("0456", str(result["REVENUE_CD"]))

    @patch("src.codes.code_funcs.extract_codes_from_service")
    def test_code_breakout(self, mock_extract_codes):
        """Test the code_breakout function for processing DataFrame records."""
        # BILL_TYPE_CD_DESC set so pre-pass skips fill_bill_type (avoids LLM on CI).
        test_df = pd.DataFrame(
            [
                {
                    "SERVICE_TERM": "SERVICE A",
                    "BILL_TYPE_CD_DESC": "Inpatient Hospital",
                },
                {
                    "SERVICE_TERM": "SERVICE B",
                    "BILL_TYPE_CD_DESC": "Inpatient Hospital",
                },
            ]
        )

        # Return value by SERVICE_TERM so order is correct under ThreadPoolExecutor.
        def return_by_service(rep, constants):
            st = rep.get("SERVICE_TERM", "")
            if st == "SERVICE A":
                return {"SERVICE_TERM": "SERVICE A", "PROCEDURE_CD": ["12345"]}
            return {"SERVICE_TERM": "SERVICE B", "PROCEDURE_CD": ["67890"]}

        mock_extract_codes.side_effect = return_by_service

        # Test function
        result_df = code_funcs.code_breakout(test_df, self.constants)

        # Verify results - code_breakout normalizes to JSON list format
        self.assertIsInstance(result_df, pd.DataFrame)
        self.assertEqual(len(result_df), 2)
        self.assertEqual(result_df.iloc[0]["PROCEDURE_CD"], '["12345"]')
        self.assertEqual(result_df.iloc[1]["PROCEDURE_CD"], '["67890"]')

        # Verify mock calls
        self.assertEqual(mock_extract_codes.call_count, 2)


if __name__ == "__main__":
    unittest.main()