stack/tests/bcda/test_express.py

"""Tests for bcda.express — FHIR R4 flattening and CCLF transforms.

Tests cover two modules:

  1. ``bcda.express.flatten`` — private helpers and public flatteners
     (flatten_patient, flatten_eob, flatten_coverage).

  2. ``bcda.express.cclf`` — narwhals @nw.narwhalify functions that
     transform flattened BCDA resources into input_layer tables
     (medical_claim, pharmacy_claim, eligibility).

Fixtures use polars DataFrames with minimal columns matching
function signatures.
"""

from __future__ import annotations

from datetime import date

import polars as pl
import pytest

from bcda.express.cclf import stg_beneficiary_xref
from bcda.express.flatten import (
    _coding_code,
    _get_ext,
    _get_ext_money,
    _get_ext_quantity,
    _get_identifier,
    _get_supporting_info,
    _parse_date,
    _safe_get,
    flatten_coverage,
    flatten_eob,
    flatten_patient,
)

# ════════════════════════════════════════════════════════════════
# Part 1: bcda.express.flatten — helpers
# ════════════════════════════════════════════════════════════════


class TestSafeGet:
    """_safe_get: nested dict traversal."""

    def test_single_key(self) -> None:
        assert _safe_get({"a": 1}, "a") == 1

    def test_nested_keys(self) -> None:
        d = {"a": {"b": {"c": 42}}}
        assert _safe_get(d, "a", "b", "c") == 42

    def test_missing_key_returns_none(self) -> None:
        assert _safe_get({"a": 1}, "b") is None

    def test_none_input_returns_none(self) -> None:
        assert _safe_get(None, "a") is None

    def test_non_dict_intermediate(self) -> None:
        assert _safe_get({"a": "string"}, "a", "b") is None


class TestGetExt:
    """_get_ext: extract value from BB extensions."""

    def test_value_coding(self) -> None:
        exts = [
            {
                "url": "https://bb.cms.gov/vars/race",
                "valueCoding": {"code": "1"},
            }
        ]
        assert _get_ext(exts, "race") == "1"

    def test_value_code(self) -> None:
        exts = [
            {
                "url": "https://bb.cms.gov/vars/sex",
                "valueCode": "M",
            }
        ]
        assert _get_ext(exts, "sex") == "M"

    def test_value_date(self) -> None:
        exts = [
            {
                "url": "https://bb.cms.gov/vars/rfrnc_yr",
                "valueDate": "2023",
            }
        ]
        assert _get_ext(exts, "rfrnc_yr") == "2023"

    def test_value_identifier(self) -> None:
        exts = [
            {
                "url": "https://bb.cms.gov/vars/id",
                "valueIdentifier": {"value": "ABC"},
            }
        ]
        assert _get_ext(exts, "id") == "ABC"

    def test_value_money(self) -> None:
        exts = [
            {
                "url": "https://bb/vars/amt",
                "valueMoney": {"value": 99.50},
            }
        ]
        assert _get_ext(exts, "amt") == "99.5"

    def test_value_quantity(self) -> None:
        exts = [
            {
                "url": "https://bb/vars/qty",
                "valueQuantity": {"value": 3},
            }
        ]
        assert _get_ext(exts, "qty") == "3"

    def test_none_extensions(self) -> None:
        assert _get_ext(None, "x") is None

    def test_empty_extensions(self) -> None:
        assert _get_ext([], "x") is None

    def test_no_match(self) -> None:
        exts = [
            {
                "url": "https://bb/vars/other",
                "valueCode": "Y",
            }
        ]
        assert _get_ext(exts, "nope") is None

    def test_exact_url_match(self) -> None:
        exts = [
            {
                "url": "dual_01",
                "valueCoding": {"code": "00"},
            }
        ]
        assert _get_ext(exts, "dual_01") == "00"


class TestGetExtMoney:
    """_get_ext_money: extract valueMoney.value."""

    def test_returns_float(self) -> None:
        exts = [
            {
                "url": "https://bb/vars/amt",
                "valueMoney": {"value": 123.45},
            }
        ]
        assert _get_ext_money(exts, "amt") == 123.45

    def test_none_when_missing(self) -> None:
        assert _get_ext_money(None, "x") is None

    def test_no_match(self) -> None:
        exts = [
            {
                "url": "https://bb/vars/other",
                "valueMoney": {"value": 1.0},
            }
        ]
        assert _get_ext_money(exts, "nope") is None


class TestGetExtQuantity:
    """_get_ext_quantity: extract valueQuantity.value."""

    def test_returns_float(self) -> None:
        exts = [
            {
                "url": "https://bb/vars/cnt",
                "valueQuantity": {"value": 7.0},
            }
        ]
        assert _get_ext_quantity(exts, "cnt") == 7.0

    def test_none_when_missing(self) -> None:
        assert _get_ext_quantity(None, "x") is None


class TestGetIdentifier:
    """_get_identifier: find identifier by system suffix."""

    def test_finds_by_suffix(self) -> None:
        ids = [
            {
                "system": "http://hl7.org/fhir/sid/bene_id",
                "value": "B123",
            }
        ]
        assert _get_identifier(ids, "bene_id") == "B123"

    def test_none_when_empty(self) -> None:
        assert _get_identifier(None, "x") is None

    def test_no_match(self) -> None:
        ids = [
            {
                "system": "http://example.com/other",
                "value": "V1",
            }
        ]
        assert _get_identifier(ids, "bene_id") is None


class TestGetSupportingInfo:
    """_get_supporting_info: find supportingInfo by category."""

    def test_finds_by_code(self) -> None:
        info = [
            {
                "category": {"coding": [{"code": "admtype"}]},
                "code": {"coding": [{"code": "3"}]},
            }
        ]
        result = _get_supporting_info(info, "admtype")
        assert result is not None
        assert result["code"]["coding"][0]["code"] == "3"

    def test_none_when_empty(self) -> None:
        assert _get_supporting_info(None, "x") is None

    def test_no_match(self) -> None:
        info = [{"category": {"coding": [{"code": "other"}]}}]
        assert _get_supporting_info(info, "admtype") is None


class TestParseDate:
    """_parse_date: YYYY-MM-DD and datetime parsing."""

    def test_date_string(self) -> None:
        assert _parse_date("2023-06-15") == date(2023, 6, 15)

    def test_datetime_string(self) -> None:
        assert _parse_date("2023-06-15T10:30:00Z") == date(2023, 6, 15)

    def test_none_returns_none(self) -> None:
        assert _parse_date(None) is None

    def test_empty_returns_none(self) -> None:
        assert _parse_date("") is None

    def test_invalid_returns_none(self) -> None:
        assert _parse_date("not-a-date") is None


class TestCodingCode:
    """_coding_code: CodeableConcept code extraction."""

    def test_no_system_filter(self) -> None:
        cc = {"coding": [{"code": "A01"}]}
        assert _coding_code(cc) == "A01"

    def test_with_system_filter(self) -> None:
        cc = {
            "coding": [
                {
                    "system": "http://hl7.org/fhir/nch",
                    "code": "10",
                },
                {
                    "system": "http://hl7.org/fhir/eob",
                    "code": "20",
                },
            ]
        }
        assert _coding_code(cc, "nch") == "10"
        assert _coding_code(cc, "eob") == "20"

    def test_none_input(self) -> None:
        assert _coding_code(None) is None

    def test_empty_coding(self) -> None:
        assert _coding_code({"coding": []}) is None


# ════════════════════════════════════════════════════════════════
# Part 1b: bcda.express.flatten — public flatteners
# ════════════════════════════════════════════════════════════════


class TestFlattenPatient:
    """flatten_patient: FHIR Patient -> Patient model."""

    @pytest.fixture
    def minimal_patient(self) -> dict:
        return {
            "id": "PAT-001",
            "meta": {"lastUpdated": "2024-01-15T12:00:00Z"},
            "identifier": [
                {
                    "system": "http://hl7.org/fhir/sid/us-mbi",
                    "value": "1AN0TE5MK72",
                    "type": {
                        "coding": [
                            {"extension": [{"valueCoding": {"code": "current"}}]}
                        ]
                    },
                },
                {
                    "system": "http://ccw/bene_id",
                    "value": "BENE-99",
                },
            ],
            "name": [
                {
                    "given": ["ALICE", "M"],
                    "family": "JOHNSON",
                }
            ],
            "birthDate": "1955-03-22",
            "gender": "female",
            "address": [
                {
                    "line": ["100 MAIN ST"],
                    "city": "ANYTOWN",
                    "state": "MD",
                    "postalCode": "20601",
                }
            ],
            "extension": [
                {
                    "url": "https://bb/vars/race",
                    "valueCoding": {"code": "1"},
                },
                {
                    "url": "http://hl7.org/us-core-race",
                    "extension": [
                        {
                            "url": "ombCategory",
                            "valueCoding": {"code": "2106-3"},
                        }
                    ],
                },
            ],
        }

    def test_returns_patient_model(self, minimal_patient: dict) -> None:
        from bcda.table.patient import Patient

        p = flatten_patient(minimal_patient)
        assert isinstance(p, Patient)

    def test_mbi_extraction(self, minimal_patient: dict) -> None:
        p = flatten_patient(minimal_patient)
        assert p.bene_mbi_id == "1AN0TE5MK72"
        assert p.identifier_currency == "current"

    def test_bene_id(self, minimal_patient: dict) -> None:
        p = flatten_patient(minimal_patient)
        assert p.bene_id == "BENE-99"

    def test_demographics(self, minimal_patient: dict) -> None:
        p = flatten_patient(minimal_patient)
        assert p.bene_1st_name == "ALICE"
        assert p.bene_midl_name == "M"
        assert p.bene_last_name == "JOHNSON"
        assert p.bene_dob == date(1955, 3, 22)
        assert p.bene_sex_cd == "female"

    def test_address(self, minimal_patient: dict) -> None:
        p = flatten_patient(minimal_patient)
        assert p.bene_line_1_adr == "100 MAIN ST"
        assert p.geo_zip_plc_name == "ANYTOWN"
        assert p.geo_usps_state_cd == "MD"
        assert p.geo_zip5_cd == "20601"

    def test_race_codes(self, minimal_patient: dict) -> None:
        p = flatten_patient(minimal_patient)
        assert p.bene_race_cd == "1"
        assert p.us_core_race == "2106-3"

    def test_historic_mbi(self) -> None:
        """Historic MBI populates prvs_num."""
        resource = {
            "identifier": [
                {
                    "system": "http://hl7.org/us-mbi",
                    "value": "HIST-MBI",
                    "type": {
                        "coding": [
                            {"extension": [{"valueCoding": {"code": "historic"}}]}
                        ]
                    },
                    "period": {
                        "start": "2020-01-01",
                        "end": "2022-12-31",
                    },
                }
            ],
        }
        p = flatten_patient(resource)
        assert p.prvs_num == "HIST-MBI"
        assert p.efctv_bgn_dt == date(2020, 1, 1)
        assert p.efctv_end_dt == date(2022, 12, 31)


class TestFlattenEob:
    """flatten_eob: FHIR EOB -> (header, items)."""

    @pytest.fixture
    def minimal_eob(self) -> dict:
        return {
            "id": "EOB-001",
            "meta": {"lastUpdated": "2024-01-20T09:00:00Z"},
            "patient": {"reference": "Patient/BENE-99"},
            "identifier": [
                {
                    "system": "http://ccw/clm_id",
                    "value": "CLM-100",
                },
            ],
            "type": {
                "coding": [
                    {
                        "system": "http://ccw/nch_clm_type_cd",
                        "code": "71",
                    }
                ]
            },
            "billablePeriod": {
                "start": "2023-06-01",
                "end": "2023-06-10",
            },
            "extension": [
                {
                    "url": "https://bb/vars/nch_near_line_rec_ident_cd",
                    "valueCoding": {"code": "V"},
                }
            ],
            "careTeam": [
                {
                    "role": {"coding": [{"code": "attending"}]},
                    "provider": {
                        "identifier": {
                            "value": "1234567890",
                            "type": {"coding": [{"code": "npi"}]},
                        }
                    },
                }
            ],
            "diagnosis": [
                {
                    "sequence": 1,
                    "diagnosisCodeableConcept": {
                        "coding": [
                            {
                                "system": "http://hl7.org/fhir/icd-10",
                                "code": "J44.1",
                            }
                        ]
                    },
                    "type": [{"coding": [{"code": "principal"}]}],
                }
            ],
            "item": [
                {
                    "sequence": 1,
                    "servicedPeriod": {
                        "start": "2023-06-01",
                        "end": "2023-06-05",
                    },
                    "productOrService": {"coding": [{"code": "99213"}]},
                    "quantity": {"value": 1.0},
                    "adjudication": [
                        {
                            "category": {
                                "coding": [
                                    {
                                        "system": "https://bluebutton.cms.gov/resources/codesystem/adjudication",
                                        "code": "line_nch_pmt_amt",
                                    }
                                ]
                            },
                            "amount": {"value": 75.50},
                        }
                    ],
                }
            ],
        }

    def test_returns_tuple(self, minimal_eob: dict) -> None:
        header, items = flatten_eob(minimal_eob)
        from bcda.table.eob import (
            EobItem,
            ExplanationOfBenefit,
        )

        assert isinstance(header, ExplanationOfBenefit)
        assert len(items) == 1
        assert isinstance(items[0], EobItem)

    def test_header_identifiers(self, minimal_eob: dict) -> None:
        header, _ = flatten_eob(minimal_eob)
        assert header.clm_id == "CLM-100"
        assert header.bene_id == "BENE-99"

    def test_header_dates(self, minimal_eob: dict) -> None:
        header, _ = flatten_eob(minimal_eob)
        assert header.clm_from_dt == date(2023, 6, 1)
        assert header.clm_thru_dt == date(2023, 6, 10)

    def test_header_near_line(self, minimal_eob: dict) -> None:
        header, _ = flatten_eob(minimal_eob)
        assert header.nch_near_line_rec_ident_cd == "V"

    def test_header_diagnosis(self, minimal_eob: dict) -> None:
        header, _ = flatten_eob(minimal_eob)
        assert header.prncpl_dgns_cd == "J44.1"
        assert header.clm_dgns_1_cd == "J44.1"
        assert header.dgns_prcdr_icd_ind == "0"

    def test_header_care_team(self, minimal_eob: dict) -> None:
        header, _ = flatten_eob(minimal_eob)
        assert header.atndg_prvdr_npi_num == "1234567890"

    def test_item_amounts(self, minimal_eob: dict) -> None:
        _, items = flatten_eob(minimal_eob)
        assert items[0].clm_line_cvrd_pd_amt == 75.50

    def test_item_clm_id_ref(self, minimal_eob: dict) -> None:
        _, items = flatten_eob(minimal_eob)
        assert items[0].clm_id == "CLM-100"

    def test_item_service_dates(self, minimal_eob: dict) -> None:
        _, items = flatten_eob(minimal_eob)
        assert items[0].clm_line_from_dt == date(2023, 6, 1)
        assert items[0].clm_line_thru_dt == date(2023, 6, 5)

    def test_item_hcpcs(self, minimal_eob: dict) -> None:
        _, items = flatten_eob(minimal_eob)
        assert items[0].clm_line_hcpcs_cd == "99213"

    def test_item_quantity(self, minimal_eob: dict) -> None:
        _, items = flatten_eob(minimal_eob)
        assert items[0].clm_line_srvc_unit_qty == 1.0

    def test_no_items_returns_empty_list(
        self,
    ) -> None:
        resource = {
            "identifier": [],
            "item": [],
        }
        _, items = flatten_eob(resource)
        assert items == []


class TestFlattenCoverage:
    """flatten_coverage: FHIR Coverage -> Coverage model."""

    @pytest.fixture
    def minimal_coverage(self) -> dict:
        return {
            "id": "COV-001",
            "meta": {"lastUpdated": "2024-02-01T08:00:00Z"},
            "status": "active",
            "subscriberId": "1AN0TE5MK72",
            "beneficiary": {"reference": "Patient/BENE-99"},
            "type": {"coding": [{"code": "SUBSIDIZ"}]},
            "class": [
                {
                    "type": {"coding": [{"code": "group"}]},
                    "value": "Medicare",
                },
                {
                    "type": {"coding": [{"code": "plan"}]},
                    "value": "Part A",
                },
            ],
            "period": {
                "start": "2020-01-01",
                "end": "2025-12-31",
            },
            "extension": [
                {
                    "url": "https://bb/vars/ms_cd",
                    "valueCoding": {"code": "10"},
                },
                {
                    "url": "https://bb/vars/orec",
                    "valueCoding": {"code": "0"},
                },
            ],
        }

    def test_returns_coverage_model(self, minimal_coverage: dict) -> None:
        from bcda.table.coverage import Coverage

        c = flatten_coverage(minimal_coverage)
        assert isinstance(c, Coverage)

    def test_subscriber_id(self, minimal_coverage: dict) -> None:
        c = flatten_coverage(minimal_coverage)
        assert c.subscriber_id == "1AN0TE5MK72"

    def test_part_a_dates(self, minimal_coverage: dict) -> None:
        c = flatten_coverage(minimal_coverage)
        assert c.bene_part_a_enrlmt_bgn_dt == date(2020, 1, 1)
        assert c.mdcr_covg_a_end_dt == date(2025, 12, 31)

    def test_extension_values(self, minimal_coverage: dict) -> None:
        c = flatten_coverage(minimal_coverage)
        assert c.ms_cd == "10"
        assert c.orec == "0"

    def test_class_group_plan(self, minimal_coverage: dict) -> None:
        c = flatten_coverage(minimal_coverage)
        assert c.class_group == "Medicare"
        assert c.class_plan == "Part A"

    def test_part_b_dates(self) -> None:
        resource = {
            "class": [
                {
                    "type": {"coding": [{"code": "plan"}]},
                    "value": "Part B",
                }
            ],
            "period": {
                "start": "2021-07-01",
                "end": "2025-06-30",
            },
            "extension": [],
        }
        c = flatten_coverage(resource)
        assert c.bene_part_b_enrlmt_bgn_dt == date(2021, 7, 1)
        assert c.mdcr_covg_b_end_dt == date(2025, 6, 30)
        assert c.bene_part_a_enrlmt_bgn_dt is None

    def test_part_d_dates(self) -> None:
        resource = {
            "class": [
                {
                    "type": {"coding": [{"code": "plan"}]},
                    "value": "Part D",
                }
            ],
            "period": {
                "start": "2022-01-01",
                "end": "2024-12-31",
            },
            "extension": [],
        }
        c = flatten_coverage(resource)
        assert c.mdcr_covg_d_bgn_dt == date(2022, 1, 1)
        assert c.mdcr_covg_d_end_dt == date(2024, 12, 31)


# ════════════════════════════════════════════════════════════════
# Part 2: bcda.express.cclf — narwhals transforms
# ════════════════════════════════════════════════════════════════

# ── shared fixtures for CCLF transforms ─────────────────────


@pytest.fixture
def patient_df() -> pl.DataFrame:
    """Minimal bcda.patient with two beneficiaries."""
    return pl.DataFrame(
        {
            "bene_id": ["B001", "B002"],
            "bene_mbi_id": ["MBI-A", "MBI-B"],
        }
    )


@pytest.fixture
def xref_df(patient_df: pl.DataFrame) -> pl.DataFrame:
    """Staged beneficiary xref from patient."""
    return stg_beneficiary_xref(bcda__patient=patient_df)


@pytest.fixture
def eob_df() -> pl.DataFrame:
    """Minimal bcda.explanation_of_benefit with four
    claim types (V, O, M, P)."""
    return pl.DataFrame(
        {
            "bene_id": [
                "B001",
                "B001",
                "B002",
                "B002",
            ],
            "clm_id": [
                "CLM-1",
                "CLM-2",
                "CLM-3",
                "CLM-4",
            ],
            "pde_id": [None, None, None, "PDE-1"],
            "nch_near_line_rec_ident_cd": [
                "V",
                "O",
                "M",
                "P",
            ],
            "clm_from_dt": [
                date(2023, 1, 1),
                date(2023, 2, 1),
                date(2023, 3, 1),
                date(2023, 4, 1),
            ],
            "clm_thru_dt": [
                date(2023, 1, 10),
                date(2023, 2, 10),
                date(2023, 3, 10),
                date(2023, 4, 10),
            ],
            "dgns_prcdr_icd_ind": [
                "0",
                "0",
                None,
                None,
            ],
            "clm_dgns_1_cd": [
                "J44.1",
                "M54.5",
                None,
                None,
            ],
            "clm_dgns_2_cd": [None, None, None, None],
            "clm_dgns_3_cd": [None, None, None, None],
            "clm_dgns_4_cd": [None, None, None, None],
            "clm_dgns_5_cd": [None, None, None, None],
            "clm_dgns_6_cd": [None, None, None, None],
            "clm_dgns_7_cd": [None, None, None, None],
            "clm_dgns_8_cd": [None, None, None, None],
            "clm_dgns_9_cd": [None, None, None, None],
            "clm_dgns_10_cd": [
                None,
                None,
                None,
                None,
            ],
            "clm_dgns_11_cd": [
                None,
                None,
                None,
                None,
            ],
            "clm_dgns_12_cd": [
                None,
                None,
                None,
                None,
            ],
            "clm_poa_ind": ["Y", None, None, None],
            "clm_prcdr_cd": [
                "0BJ08ZZ",
                None,
                None,
                None,
            ],
            "clm_prcdr_prfrm_dt": [
                date(2023, 1, 2),
                None,
                None,
                None,
            ],
            "dgns_drg_cd": [
                "MS-DRG470",
                None,
                None,
                None,
            ],
            "clm_admsn_src_cd": [
                "1",
                None,
                None,
                None,
            ],
            "clm_admsn_type_cd": [
                "3",
                None,
                None,
                None,
            ],
            "bene_ptnt_stus_cd": [
                "01",
                None,
                None,
                None,
            ],
            "clm_bill_fac_type_cd": [
                "1",
                None,
                None,
                None,
            ],
            "clm_bill_clsfctn_cd": [
                "1",
                None,
                None,
                None,
            ],
            "clm_bill_freq_cd": [
                "1",
                None,
                None,
                None,
            ],
            "clm_mdcr_instnl_tot_chrg_amt": [
                5000.0,
                None,
                None,
                None,
            ],
            "oprtg_prvdr_npi_num": [
                "NPI-OP1",
                None,
                None,
                None,
            ],
            "fac_prvdr_npi_num": [
                "NPI-FAC1",
                None,
                None,
                None,
            ],
            "rndrg_prvdr_npi_num": [
                None,
                "NPI-RN2",
                None,
                None,
            ],
            "clm_rndrg_prvdr_tax_num": [
                None,
                "TAX-2",
                None,
                None,
            ],
            "carr_clm_blg_npi_num": [
                None,
                "NPI-BLG2",
                None,
                None,
            ],
            "clm_pos_cd": [None, "11", "12", None],
            "ordrg_prvdr_npi_num": [
                None,
                None,
                "NPI-ORD3",
                None,
            ],
            "payto_prvdr_npi_num": [
                None,
                None,
                "NPI-PAY3",
                None,
            ],
            "clm_prsbng_prvdr_gnrc_id_num": [
                None,
                None,
                None,
                "NPI-PRE4",
            ],
            "clm_srvc_prvdr_gnrc_id_num": [
                None,
                None,
                None,
                "NPI-DSP4",
            ],
        }
    )


@pytest.fixture
def eob_item_df() -> pl.DataFrame:
    """Minimal bcda.eob_item with matching clm_ids."""
    return pl.DataFrame(
        {
            "clm_id": [
                "CLM-1",
                "CLM-2",
                "CLM-3",
                "CLM-4",
            ],
            "clm_line_num": [1.0, 1.0, 1.0, 1.0],
            "clm_line_from_dt": [
                date(2023, 1, 1),
                date(2023, 2, 1),
                date(2023, 3, 1),
                date(2023, 4, 1),
            ],
            "clm_line_thru_dt": [
                date(2023, 1, 5),
                date(2023, 2, 5),
                date(2023, 3, 5),
                date(2023, 4, 5),
            ],
            "clm_line_prod_rev_ctr_cd": [
                "0301",
                None,
                None,
                None,
            ],
            "clm_line_hcpcs_cd": [
                "99213",
                "99214",
                "E1390",
                None,
            ],
            "clm_line_ndc_cd": [
                None,
                None,
                None,
                "12345678901",
            ],
            "hcpcs_1_mdfr_cd": [
                "26",
                None,
                None,
                None,
            ],
            "hcpcs_2_mdfr_cd": [
                None,
                None,
                None,
                None,
            ],
            "hcpcs_3_mdfr_cd": [
                None,
                None,
                None,
                None,
            ],
            "hcpcs_4_mdfr_cd": [
                None,
                None,
                None,
                None,
            ],
            "hcpcs_5_mdfr_cd": [
                None,
                None,
                None,
                None,
            ],
            "clm_line_srvc_unit_qty": [
                1.0,
                2.0,
                None,
                30.0,
            ],
            "clm_line_cvrd_pd_amt": [
                75.50,
                120.00,
                50.00,
                10.00,
            ],
            "clm_line_alowd_chrg_amt": [
                100.00,
                150.00,
                60.00,
                None,
            ],
            "clm_line_bene_pmt_amt": [
                None,
                None,
                None,
                5.00,
            ],
            "clm_line_days_suply_qty": [
                None,
                None,
                None,
                30.0,
            ],
            "clm_line_rx_fill_num": [
                None,
                None,
                None,
                0.0,
            ],
        }
    )


# ── stg_beneficiary_xref ────────────────────────────────────


class TestStgBeneficiaryXref:
    """BCDA Patient -> MBI crosswalk."""

    def test_output_columns(self, patient_df: pl.DataFrame) -> None:
        result = stg_beneficiary_xref(bcda__patient=patient_df)
        assert set(result.columns) == {
            "bene_id",
            "current_bene_mbi_id",
        }

    def test_dedup_on_bene_id(self) -> None:
        df = pl.DataFrame(
            {
                "bene_id": ["B1", "B1", "B2"],
                "bene_mbi_id": [
                    "MBI-1",
                    "MBI-1",
                    "MBI-2",
                ],
            }
        )
        result = stg_beneficiary_xref(bcda__patient=df)
        assert result.shape[0] == 2
        assert result["bene_id"].n_unique() == result.shape[0]

    def test_alias_current_bene_mbi_id(self, patient_df: pl.DataFrame) -> None:
        result = stg_beneficiary_xref(bcda__patient=patient_df)
        assert "current_bene_mbi_id" in result.columns
        assert set(result["current_bene_mbi_id"].to_list()) == {"MBI-A", "MBI-B"}


class TestGetExtQuantityMatch:
    """_get_ext_quantity: cover the return-value path (line 155)."""

    def test_returns_value_on_match(self) -> None:
        exts = [
            {
                "url": "https://bb/vars/rev_cntr_unit_cnt",
                "valueQuantity": {"value": 42.0},
            }
        ]
        assert _get_ext_quantity(exts, "rev_cntr_unit_cnt") == 42.0

    def test_no_match_returns_none(self) -> None:
        exts = [
            {
                "url": "https://bb/vars/other",
                "valueQuantity": {"value": 1.0},
            }
        ]
        assert _get_ext_quantity(exts, "rev_cntr_unit_cnt") is None


class TestExtractCareTeamRoles:
    """_extract_care_team: cover all role branches (lines 348-379)."""

    def test_no_role_coding_skips(self) -> None:
        from bcda.express.flatten import _extract_care_team

        care_team = [{"provider": {"identifier": {"value": "NPI1"}}}]
        result = _extract_care_team(care_team)
        assert result["atndg_prvdr_npi_num"] is None

    def test_performing_role_with_npi_and_qualification(self) -> None:
        from bcda.express.flatten import _extract_care_team

        care_team = [
            {
                "role": {"coding": [{"code": "performing"}]},
                "provider": {
                    "identifier": {
                        "value": "NPI-PERF",
                        "type": {"coding": [{"code": "npi"}]},
                    }
                },
                "qualification": {"coding": [{"code": "01"}]},
                "extension": [
                    {
                        "url": "https://bb/vars/carr_line_prvdr_type_cd",
                        "valueCoding": {"code": "1"},
                    },
                    {
                        "url": "https://bb/vars/prtcptng_ind_cd",
                        "valueCoding": {"code": "2"},
                    },
                ],
            }
        ]
        result = _extract_care_team(care_team)
        assert result["rndrg_prvdr_npi_num"] == "NPI-PERF"
        assert result["clm_prvdr_spclty_cd"] == "01"
        assert result["carr_line_prvdr_type_cd"] == "1"
        assert result["prtcptng_ind_cd"] == "2"

    def test_prescribing_role(self) -> None:
        from bcda.express.flatten import _extract_care_team

        care_team = [
            {
                "role": {"coding": [{"code": "prescribing"}]},
                "provider": {
                    "identifier": {
                        "value": "PRES-001",
                        "type": {"coding": [{"code": "npi"}]},
                    }
                },
            }
        ]
        result = _extract_care_team(care_team)
        assert result["clm_prsbng_prvdr_gnrc_id_num"] == "PRES-001"

    def test_dispensing_role(self) -> None:
        from bcda.express.flatten import _extract_care_team

        care_team = [
            {
                "role": {"coding": [{"code": "dispensing"}]},
                "provider": {
                    "identifier": {
                        "value": "DISP-001",
                        "type": {"coding": [{"code": "npi"}]},
                    }
                },
            }
        ]
        result = _extract_care_team(care_team)
        assert result["clm_srvc_prvdr_gnrc_id_num"] == "DISP-001"

    def test_referring_role_with_npi(self) -> None:
        from bcda.express.flatten import _extract_care_team

        care_team = [
            {
                "role": {"coding": [{"code": "referring"}]},
                "provider": {
                    "identifier": {
                        "value": "REF-NPI",
                        "type": {"coding": [{"code": "npi"}]},
                    }
                },
            }
        ]
        result = _extract_care_team(care_team)
        assert result["ordrg_prvdr_npi_num"] == "REF-NPI"

    def test_operating_role(self) -> None:
        from bcda.express.flatten import _extract_care_team

        care_team = [
            {
                "role": {"coding": [{"code": "operating"}]},
                "provider": {
                    "identifier": {
                        "value": "OP-NPI",
                        "type": {"coding": [{"code": "npi"}]},
                    }
                },
            }
        ]
        result = _extract_care_team(care_team)
        assert result["oprtg_prvdr_npi_num"] == "OP-NPI"

    def test_otheroperating_role(self) -> None:
        from bcda.express.flatten import _extract_care_team

        care_team = [
            {
                "role": {"coding": [{"code": "otheroperating"}]},
                "provider": {
                    "identifier": {
                        "value": "OTHER-NPI",
                        "type": {"coding": [{"code": "npi"}]},
                    }
                },
            }
        ]
        result = _extract_care_team(care_team)
        assert result["othr_prvdr_npi_num"] == "OTHER-NPI"


class TestExtractDiagnosesGaps:
    """_extract_diagnoses: typed diagnoses and POA (lines 405-423)."""

    def test_icd9_system(self) -> None:
        from bcda.express.flatten import _extract_diagnoses

        dx_list = [
            {
                "sequence": 1,
                "diagnosisCodeableConcept": {
                    "coding": [
                        {
                            "system": "http://hl7.org/fhir/icd-9",
                            "code": "250.00",
                        }
                    ]
                },
            }
        ]
        result = _extract_diagnoses(dx_list)
        assert result["dgns_prcdr_icd_ind"] == "9"
        assert result["clm_dgns_1_cd"] == "250.00"

    def test_admitting_type(self) -> None:
        from bcda.express.flatten import _extract_diagnoses

        dx_list = [
            {
                "sequence": 2,
                "diagnosisCodeableConcept": {
                    "coding": [
                        {
                            "system": "http://hl7.org/fhir/icd-10",
                            "code": "R07.9",
                        }
                    ]
                },
                "type": [{"coding": [{"code": "admitting"}]}],
            }
        ]
        result = _extract_diagnoses(dx_list)
        assert result["admtg_dgns_cd"] == "R07.9"
        assert result["clm_dgns_2_cd"] == "R07.9"

    def test_poa_indicator(self) -> None:
        from bcda.express.flatten import _extract_diagnoses

        dx_list = [
            {
                "sequence": 1,
                "diagnosisCodeableConcept": {
                    "coding": [
                        {
                            "system": "http://hl7.org/fhir/icd-10",
                            "code": "J44.1",
                        }
                    ]
                },
                "extension": [
                    {
                        "url": "https://bb/vars/clm_poa_ind_sw1",
                        "valueCoding": {"code": "Y"},
                    }
                ],
            }
        ]
        result = _extract_diagnoses(dx_list)
        assert result["clm_poa_ind"] == "Y"

    def test_poa_value_code_fallback(self) -> None:
        from bcda.express.flatten import _extract_diagnoses

        dx_list = [
            {
                "sequence": 1,
                "diagnosisCodeableConcept": {
                    "coding": [{"system": "icd-10", "code": "A00"}]
                },
                "extension": [
                    {
                        "url": "https://bb/vars/POA",
                        "valueCode": "N",
                    }
                ],
            }
        ]
        result = _extract_diagnoses(dx_list)
        assert result["clm_poa_ind"] == "N"


class TestExtractItemNoncovered:
    """_extract_item: noncovered adjudication code (lines 471-472)."""

    def test_noncovered_adjudication(self) -> None:
        from bcda.express.flatten import _extract_item

        item = {
            "sequence": 1,
            "adjudication": [
                {
                    "category": {"coding": [{"code": "noncovered"}]},
                    "amount": {"value": 25.00},
                }
            ],
        }
        result = _extract_item(item, "CLM-1")
        assert result.clm_line_ncvrd_chrg_amt == 25.00


class TestFlattenEobContainedAndProvider:
    """flatten_eob: contained Organization, billing NPI (lines 588-620)."""

    def test_contained_organization_prn_and_npi(self) -> None:
        resource = {
            "identifier": [{"system": "http://ccw/clm_id", "value": "CLM-C1"}],
            "contained": [
                {
                    "resourceType": "Organization",
                    "identifier": [
                        {
                            "type": {"coding": [{"code": "PRN"}]},
                            "value": "OSCAR-123",
                        },
                        {
                            "type": {"coding": [{"code": "npi"}]},
                            "value": "NPI-FAC-456",
                        },
                    ],
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.prvdr_oscar_num == "OSCAR-123"
        assert header.fac_prvdr_npi_num == "NPI-FAC-456"

    def test_contained_non_organization_skipped(self) -> None:
        resource = {
            "identifier": [],
            "contained": [
                {
                    "resourceType": "Practitioner",
                    "identifier": [
                        {
                            "type": {"coding": [{"code": "PRN"}]},
                            "value": "SHOULD-SKIP",
                        }
                    ],
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.prvdr_oscar_num is None

    def test_billing_npi_from_provider(self) -> None:
        resource = {
            "identifier": [],
            "provider": {
                "identifier": {
                    "system": "http://hl7.org/fhir/blg_npi",
                    "value": "BLG-NPI-999",
                }
            },
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.carr_clm_blg_npi_num == "BLG-NPI-999"

    def test_billing_npi_not_set_without_npi_system(self) -> None:
        resource = {
            "identifier": [],
            "provider": {
                "identifier": {
                    "system": "http://hl7.org/fhir/tax_id",
                    "value": "TAX-123",
                }
            },
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.carr_clm_blg_npi_num is None


class TestFlattenEobSupportingInfoHelpers:
    """flatten_eob: inner _si_code, _si_date, _si_quantity, _si_period."""

    def test_si_code_returns_code(self) -> None:
        resource = {
            "identifier": [],
            "supportingInfo": [
                {
                    "category": {"coding": [{"code": "discharge-status"}]},
                    "code": {"coding": [{"code": "01"}]},
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.bene_ptnt_stus_cd == "01"

    def test_si_date_returns_date(self) -> None:
        resource = {
            "identifier": [],
            "supportingInfo": [
                {
                    "category": {"coding": [{"code": "clmrecvddate"}]},
                    "timingDate": "2023-08-15",
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.clm_efctv_dt == date(2023, 8, 15)

    def test_si_quantity_returns_value(self) -> None:
        resource = {
            "identifier": [],
            "supportingInfo": [
                {
                    "category": {"coding": [{"code": "nch_blood_pnts_frnshd_qty"}]},
                    "valueQuantity": {"value": 3.0},
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.nch_blood_pnts_frnshd_qty == 3.0

    def test_si_period_returns_dates(self) -> None:
        resource = {
            "identifier": [],
            "supportingInfo": [
                {
                    "category": {"coding": [{"code": "admissionperiod"}]},
                    "timingPeriod": {
                        "start": "2023-07-01",
                        "end": "2023-07-10",
                    },
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        # admissionperiod is not directly on header but
        # the EOB processing doesn't have admsn_start on the model
        # — test that the code runs without error
        assert header is not None


class TestFlattenEobTotalCharge:
    """flatten_eob: total charge from 'submitted' code (lines 671-673)."""

    def test_total_charge_extraction(self) -> None:
        resource = {
            "identifier": [],
            "total": [
                {
                    "category": {"coding": [{"code": "submitted"}]},
                    "amount": {"value": 12500.00},
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.clm_mdcr_instnl_tot_chrg_amt == 12500.00

    def test_total_charge_not_set_for_other_codes(self) -> None:
        resource = {
            "identifier": [],
            "total": [
                {
                    "category": {"coding": [{"code": "benefit"}]},
                    "amount": {"value": 999.00},
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.clm_mdcr_instnl_tot_chrg_amt is None


class TestFlattenEobProcedure:
    """flatten_eob: procedure extraction (lines 617-620)."""

    def test_procedure_extraction(self) -> None:
        resource = {
            "identifier": [],
            "procedure": [
                {
                    "sequence": 1,
                    "procedureCodeableConcept": {"coding": [{"code": "0BJ08ZZ"}]},
                    "date": "2023-06-02",
                }
            ],
            "item": [],
        }
        header, _ = flatten_eob(resource)
        assert header.clm_val_sqnc_num == "1"
        assert header.clm_prcdr_cd == "0BJ08ZZ"
        assert header.clm_prcdr_prfrm_dt == date(2023, 6, 2)


class TestBatchFlatteners:
    """flatten_patients, flatten_eobs, flatten_coverages (lines 911-931)."""

    def test_flatten_patients_batch(self) -> None:
        from bcda.express.flatten import flatten_patients

        resources = [
            {
                "id": "P1",
                "identifier": [],
            },
            {
                "id": "P2",
                "identifier": [],
            },
        ]
        result = flatten_patients(resources)
        assert len(result) == 2
        assert result[0].id == "P1"
        assert result[1].id == "P2"

    def test_flatten_eobs_batch(self) -> None:
        from bcda.express.flatten import flatten_eobs

        resources = [
            {
                "id": "E1",
                "identifier": [{"system": "http://ccw/clm_id", "value": "C1"}],
                "item": [
                    {
                        "sequence": 1,
                        "productOrService": {"coding": [{"code": "99213"}]},
                    }
                ],
            },
            {
                "id": "E2",
                "identifier": [],
                "item": [],
            },
        ]
        headers, items = flatten_eobs(resources)
        assert len(headers) == 2
        assert len(items) == 1
        assert headers[0].id == "E1"
        assert headers[1].id == "E2"
        assert items[0].clm_id == "C1"

    def test_flatten_coverages_batch(self) -> None:
        from bcda.express.flatten import flatten_coverages

        resources = [
            {
                "id": "CV1",
                "extension": [],
            },
            {
                "id": "CV2",
                "extension": [],
            },
        ]
        result = flatten_coverages(resources)
        assert len(result) == 2
        assert result[0].id == "CV1"
        assert result[1].id == "CV2"


class TestReadNdjsonAndWriteParquet:
    """_read_ndjson and _write_parquet helpers (lines 939-956)."""

    def test_read_ndjson(self, tmp_path) -> None:
        from bcda.express.flatten import _read_ndjson

        ndjson_file = tmp_path / "test.ndjson"
        ndjson_file.write_text(
            '{"id": "1", "name": "alice"}\n{"id": "2", "name": "bob"}\n\n'
        )
        result = _read_ndjson(ndjson_file)
        assert len(result) == 2
        assert result[0]["id"] == "1"
        assert result[1]["name"] == "bob"

    def test_write_parquet_empty_rows(self, tmp_path) -> None:
        import fsspec

        from bcda.express.flatten import _write_parquet

        fs, _ = fsspec.core.url_to_fs(str(tmp_path))
        path = str(tmp_path / "empty.parquet")
        n = _write_parquet(fs, path, [])
        assert n == 0

    def test_write_parquet_with_rows(self, tmp_path) -> None:
        import fsspec

        from bcda.express.flatten import _write_parquet

        fs, _ = fsspec.core.url_to_fs(str(tmp_path))
        path = str(tmp_path / "sub" / "data.parquet")
        rows = [
            {"id": "1", "val": 10},
            {"id": "2", "val": 20},
        ]
        n = _write_parquet(fs, path, rows)
        assert n == 2
        # Verify file was created and is readable
        import pyarrow.parquet as pq

        table = pq.read_table(path)
        assert table.num_rows == 2


class TestFlattenExport:
    """flatten_export orchestrator (lines 992-1078)."""

    def test_flatten_export_all_files(self, tmp_path) -> None:
        import json
        from unittest.mock import patch

        from bcda.express.flatten import flatten_export

        ndjson_dir = tmp_path / "ndjson"
        ndjson_dir.mkdir()
        store_dir = tmp_path / "store"

        # Patient.ndjson
        patient = {
            "id": "PAT-1",
            "identifier": [],
        }
        (ndjson_dir / "Patient.ndjson").write_text(json.dumps(patient) + "\n")

        # ExplanationOfBenefit.ndjson
        eob = {
            "id": "EOB-1",
            "identifier": [{"system": "http://ccw/clm_id", "value": "C1"}],
            "item": [
                {
                    "sequence": 1,
                    "productOrService": {"coding": [{"code": "99213"}]},
                }
            ],
        }
        (ndjson_dir / "ExplanationOfBenefit.ndjson").write_text(json.dumps(eob) + "\n")

        # Coverage.ndjson
        coverage = {
            "id": "COV-1",
            "extension": [],
        }
        (ndjson_dir / "Coverage.ndjson").write_text(json.dumps(coverage) + "\n")

        with patch("bcda.log.setup"):
            counts = flatten_export(ndjson_dir, store_dir)

        assert counts["bcda.patient"] == 1
        assert counts["bcda.explanation_of_benefit"] == 1
        assert counts["bcda.eob_item"] == 1
        assert counts["bcda.coverage"] == 1

        # Verify parquet files exist
        assert (store_dir / "flat" / "patient.parquet").exists()
        assert (store_dir / "flat" / "explanation_of_benefit.parquet").exists()
        assert (store_dir / "flat" / "eob_item.parquet").exists()
        assert (store_dir / "flat" / "coverage.parquet").exists()

    def test_flatten_export_missing_files(self, tmp_path) -> None:
        from unittest.mock import patch

        from bcda.express.flatten import flatten_export

        ndjson_dir = tmp_path / "empty_ndjson"
        ndjson_dir.mkdir()
        store_dir = tmp_path / "store2"

        with patch("bcda.log.setup"):
            counts = flatten_export(ndjson_dir, store_dir)

        assert counts == {}

    def test_flatten_export_patient_only(self, tmp_path) -> None:
        import json
        from unittest.mock import patch

        from bcda.express.flatten import flatten_export

        ndjson_dir = tmp_path / "ndjson_pat"
        ndjson_dir.mkdir()
        store_dir = tmp_path / "store3"

        patient = {"id": "P1", "identifier": []}
        (ndjson_dir / "Patient.ndjson").write_text(json.dumps(patient) + "\n")

        with patch("bcda.log.setup"):
            counts = flatten_export(ndjson_dir, store_dir)

        assert "bcda.patient" in counts
        assert "bcda.explanation_of_benefit" not in counts
        assert "bcda.coverage" not in counts


# ── v3 coverage shape through the CCLF demographics stage ───────