"""
generate_supplier_set.py — Parameterized Signal CSV generator.
Produces N files for a given state/supplier/challenge-level combination.

Usage:
    python generate_supplier_set.py --state PA --supplier "Gaboro DME" --patients 30 --challenge-level high --output pa-set
    python generate_supplier_set.py --state NJ --supplier "Gaboro DME" --patients 25 --challenge-level high --output nj-set

Challenge levels:
    low:    canonical headers, clean dates, valid payers
    medium: header variants, mixed date formats, some nulls
    high:   all variants + nulls + malformed dates + transfer patients + state-change patients
"""

import argparse
import csv
import json
import os
import random
from datetime import date, timedelta

random.seed(42)

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
MARKET_DATA_PATH = os.path.join(SCRIPT_DIR, "market_data.json")

DEVICE_TYPES = ["dexcom_g7", "dexcom_g6", "freestyle_libre_3", "freestyle_libre_2"]
COMPONENTS = {
    "dexcom_g7": "sensor",
    "dexcom_g6": "sensor",
    "freestyle_libre_3": "sensor",
    "freestyle_libre_2": "sensor",
}

TODAY = date.today()

DATE_BUCKETS = {
    "active":           (TODAY - timedelta(days=10),  TODAY - timedelta(days=1)),
    "renewal_soon":     (TODAY - timedelta(days=270), TODAY - timedelta(days=260)),
    "renewal_elevated": (TODAY - timedelta(days=300), TODAY - timedelta(days=290)),
    "renewal_critical": (TODAY - timedelta(days=315), TODAY - timedelta(days=310)),
    "visit_required":   (TODAY - timedelta(days=400), TODAY - timedelta(days=380)),
    "supply_lapsed":    (TODAY - timedelta(days=600), TODAY - timedelta(days=500)),
}

# Header variants — each entry is a tuple:
# (patient_id_col, device_col, date_col, qty_col, payer_col, component_col, date_fmt, extras_factory)
# extras_factory(state_data) -> dict of extra static columns to include
BASE_HEADER_VARIANTS = [
    # 1 — canonical
    ("patient_id", "device_type", "shipment_date", "quantity", "payer", "component",
     "%Y-%m-%d", lambda sd: {}),
    # 2 — Brightree-style
    ("Patient ID", "Item Description", "Service Date", "Qty", "Insurance Name", "Item Type",
     "%m/%d/%Y", lambda sd: {"Prescriber NPI": "1234567890", "Branch": sd["branch_id"]}),
    # 3 — all caps
    ("PT_ID", "DEVICE", "SHIP DATE", "UNITS", "CARRIER", "TYPE",
     "%Y-%m-%d", lambda sd: {}),
    # 4 — MRN + text date
    ("MRN", "Product Name", "Dispense Date", "Qty Dispensed", "Plan Name", "Supply Type",
     "%d-%b-%Y", lambda sd: {"Supplier": sd["supplier_name"], "State": sd["state"]}),
    # 5 — account_number
    ("Account Number", "Product", "Fill Date", "Count", "Primary Payer", "component",
     "%m/%d/%y", lambda sd: {}),
    # 6 — external ref + ISO datetime
    ("External Patient Ref", "Item", "Date of Service", "Qty Shipped", "Insurance", "item_type",
     "%Y-%m-%dT%H:%M:%S", lambda sd: {"Notes": "batch export", "Region": sd["geographies"][0]}),
    # 7 — Acct # + YYYYMMDD
    ("Acct #", "DME", "Order Date", "quantity", "plan", "component_type",
     "%Y%m%d", lambda sd: {}),
    # 8 — patientid no space + dispense date
    ("patientid", "devicetype", "dispensedate", "qty", "payername", "supplytype",
     "%m/%d/%Y", lambda sd: {}),
    # 9 — pt_id + dos
    ("pt_id", "product_type", "dos", "qty", "ins_name", "component",
     "%Y-%m-%d", lambda sd: {}),
    # 10 — account_no + hcpcs
    ("account_no", "hcpcs_description", "service_date", "units", "primary_payer", "supply_type",
     "%m-%d-%Y", lambda sd: {"HCPCS Code": "A9277"}),
    # 11 — patient_account + commercial
    ("patient_account", "description", "ship_date", "quantity_dispensed", "carrier", "component",
     "%Y-%m-%d", lambda sd: {"Account Manager": "Staff"}),
    # 12 — id + product
    ("id", "product", "fill_date", "qty_shipped", "payer", "item_type",
     "%m/%d/%Y", lambda sd: {}),
    # 13 — PT ID spaces
    ("PT ID", "Device Type", "Shipment Date", "Quantity", "Insurance", "Component",
     "%Y-%m-%d", lambda sd: {"Region": sd["geographies"][1] if len(sd["geographies"]) > 1 else ""}),
    # 14 — lowercase no space (same as 8 — patientid/devicetype variants)
    ("patientid", "devicetype", "dispensedate", "qty", "payername", "supplytype",
     "%m/%d/%Y", lambda sd: {}),
    # 15 — account_number + d/m/y
    ("account_number", "item_description", "order_date", "units_dispensed", "plan_name", "component",
     "%d/%m/%Y", lambda sd: {"Facility": f"{sd['supplier_name']} Main"}),
    # 16 — member_id
    ("member_id", "Product Type", "Service Date", "Qty", "Insurance Name", "Component Type",
     "%Y-%m-%d", lambda sd: {}),
    # 17 — acct no
    ("acct no", "dme description", "last ship date", "qty dispensed", "payer name", "type",
     "%m/%d/%Y", lambda sd: {"Billing Cycle": "Monthly"}),
    # 18 — external_patient_ref
    ("external_patient_ref", "equipment description", "service_date", "count", "primary_payer", "supply_type",
     "%Y-%m-%dT%H:%M:%S", lambda sd: {}),
    # 19 — with doc columns: swo_status, pa_status
    ("patient_id", "device_type", "shipment_date", "quantity", "payer", "component",
     "%Y-%m-%d", lambda sd: {"swo_status": "On File", "pa_status": "Not Required", "diagnosis_on_file": "Yes"}),
    # 20 — with visit_date column
    ("patient_id", "device_type", "shipment_date", "quantity", "payer", "component",
     "%Y-%m-%d", lambda sd: {"visit_date": "", "pecos_verified": "Yes", "swo_status": "On File"}),
    # 21 — PA status column variants (Brightree-style + auth_status)
    ("Patient ID", "Item Description", "Service Date", "Qty", "Insurance Name", "Item Type",
     "%m/%d/%Y", lambda sd: {"auth_status": "Approved", "swo_status": "On File"}),
    # 22 — Prior Auth column
    ("account_no", "device_type", "shipment_date", "qty", "insurance", "component",
     "%Y-%m-%d", lambda sd: {"Prior Auth": "Pending", "diagnosis_on_file": "Yes"}),
    # 23 — with notes + auth numbers (noise columns)
    ("patient_id", "device_type", "shipment_date", "quantity", "payer", "component",
     "%Y-%m-%d", lambda sd: {"Notes": "see chart", "Auth Number": "A12345", "Staff": "K.F."}),
    # 24 — missing optional columns (no component, no quantity)
    ("patient_id", "device_type", "shipment_date", None, "payer", None,
     "%Y-%m-%d", lambda sd: {}),
    # 25 — transfer patient column
    ("patient_id", "device_type", "shipment_date", "quantity", "payer", "component",
     "%Y-%m-%d", lambda sd: {"Transfer_From": "Prior Supplier LLC", "swo_status": "Pending"}),
]


def _pick_payer(state_data: dict, payer_type: str) -> str:
    payers = state_data["payers"].get(payer_type, [])
    return random.choice(payers) if payers else "Medicare Part B"


def _format_date(d: date, fmt: str) -> str:
    return d.strftime(fmt)


def _random_date_in_bucket(bucket: tuple) -> date:
    start, end = bucket
    days = (end - start).days
    return start + timedelta(days=random.randint(0, max(days, 0)))


def _device_display(device: str, header_style: str) -> str:
    """Map device type to display string matching the header style."""
    display_map = {
        "dexcom_g7": {
            "canonical": "dexcom_g7",
            "display": "Dexcom G7",
            "short": "G7",
        },
        "dexcom_g6": {
            "canonical": "dexcom_g6",
            "display": "Dexcom G6",
            "short": "G6",
        },
        "freestyle_libre_3": {
            "canonical": "freestyle_libre_3",
            "display": "FreeStyle Libre 3",
            "short": "FSL3",
        },
        "freestyle_libre_2": {
            "canonical": "freestyle_libre_2",
            "display": "FreeStyle Libre 2",
            "short": "FSL2",
        },
    }
    dm = display_map.get(device, {"canonical": device, "display": device, "short": device})
    if header_style in (
        "Item Description", "Product Name", "product", "description",
        "dme description", "equipment description", "item_description",
        "hcpcs_description", "Product Type",
    ):
        return dm["display"]
    if header_style in ("DEVICE", "DME"):
        return dm["short"]
    return dm["canonical"]


def generate_state_change_patients(state_data: dict, file_index: int, num_files: int) -> list:
    """
    Generate 6 state-change patients whose status evolves across the file sequence.
    Returns patient records that change based on file_index.
    """
    patients = []
    progress = file_index / max(num_files - 1, 1)  # 0.0 to 1.0

    # Patient A: starts RENEWAL_SOON, confirmed in file 10, ACTIVE by file 20
    if file_index < 10:
        shipment_date_a = TODAY - timedelta(days=270)
        visit_date_a = ""
        swo_status_a = "On File"
    elif file_index < 20:
        shipment_date_a = TODAY - timedelta(days=10)
        visit_date_a = (TODAY - timedelta(days=5)).strftime("%Y-%m-%d")
        swo_status_a = "On File"
    else:
        shipment_date_a = TODAY - timedelta(days=5)
        visit_date_a = (TODAY - timedelta(days=5)).strftime("%Y-%m-%d")
        swo_status_a = "On File"

    patients.append({
        "patient_id": "SC-A-001",
        "device_type": "dexcom_g7",
        "shipment_date": shipment_date_a,
        "quantity": 1,
        "payer": _pick_payer(state_data, "medicare_ffs"),
        "component": "sensor",
        "visit_date": visit_date_a,
        "swo_status": swo_status_a,
        "pecos_verified": "Yes",
        "pa_status": "Not Required",
        "diagnosis_on_file": "Yes",
    })

    # Patient B: starts ACTIVE, slides to RENEWAL_SOON by file 15
    base_visit_offset = 170 - int(progress * 90)  # 170 to 80 days ago
    patients.append({
        "patient_id": "SC-B-002",
        "device_type": "freestyle_libre_3",
        "shipment_date": TODAY - timedelta(days=5),
        "quantity": 1,
        "payer": _pick_payer(state_data, "medicaid"),
        "component": "sensor",
        "visit_date": (TODAY - timedelta(days=base_visit_offset)).strftime("%Y-%m-%d"),
        "swo_status": "On File",
        "pecos_verified": "Yes",
        "pa_status": "Approved",
        "diagnosis_on_file": "Yes",
    })

    # Patient C: transfer patient appearing mid-sequence, clears by file 18
    if file_index < 12:
        transfer_from_c = "Prior Supplier LLC"
        swo_status_c = "Pending"
        pa_status_c = "Pending"
        diagnosis_c = "No"
    else:
        transfer_from_c = ""
        swo_status_c = "On File"
        pa_status_c = "Approved"
        diagnosis_c = "Yes"

    patients.append({
        "patient_id": "SC-C-003",
        "device_type": "dexcom_g6",
        "shipment_date": TODAY - timedelta(days=30),
        "quantity": 1,
        "payer": _pick_payer(state_data, "commercial"),
        "component": "sensor",
        "Transfer_From": transfer_from_c,
        "swo_status": swo_status_c,
        "pa_status": pa_status_c,
        "diagnosis_on_file": diagnosis_c,
    })

    # Patient D: SUPPLY_LAPSED -> new shipment after file 8 -> ACTIVE
    shipment_d = TODAY - timedelta(days=600) if file_index < 8 else TODAY - timedelta(days=5)
    patients.append({
        "patient_id": "SC-D-004",
        "device_type": "dexcom_g7",
        "shipment_date": shipment_d,
        "quantity": 1,
        "payer": _pick_payer(state_data, "medicare_advantage"),
        "component": "sensor",
        "swo_status": "On File",
        "pa_status": "Approved" if file_index >= 8 else "Pending",
        "diagnosis_on_file": "Yes",
    })

    # Patients E and F: stable control group
    patients.append({
        "patient_id": "SC-E-005",
        "device_type": "freestyle_libre_2",
        "shipment_date": TODAY - timedelta(days=3),
        "quantity": 1,
        "payer": _pick_payer(state_data, "medicare_ffs"),
        "component": "sensor",
        "visit_date": (TODAY - timedelta(days=30)).strftime("%Y-%m-%d"),
        "swo_status": "On File",
        "pecos_verified": "Yes",
        "pa_status": "Not Required",
        "diagnosis_on_file": "Yes",
    })
    patients.append({
        "patient_id": "SC-F-006",
        "device_type": "dexcom_g7",
        "shipment_date": TODAY - timedelta(days=7),
        "quantity": 1,
        "payer": _pick_payer(state_data, "commercial"),
        "component": "sensor",
        "swo_status": "On File",
        "pa_status": "Approved",
        "diagnosis_on_file": "Yes",
    })

    return patients


def generate_challenge_rows(state_data: dict, challenge_level: str, file_index: int) -> list:
    """Generate challenge-type rows based on challenge_level."""
    if challenge_level == "low":
        return []

    challenges = []
    if challenge_level in ("medium", "high"):
        # Null patient_id (should be skipped by normalizer)
        challenges.append({
            "patient_id": "",
            "device_type": "dexcom_g7",
            "shipment_date": TODAY - timedelta(days=10),
            "quantity": 1,
            "payer": _pick_payer(state_data, "medicare_ffs"),
            "component": "sensor",
        })

    if challenge_level == "high":
        # Future shipment date (edge case)
        challenges.append({
            "patient_id": f"FUTURE-{file_index:03d}",
            "device_type": "dexcom_g7",
            "shipment_date": TODAY + timedelta(days=30),
            "quantity": 1,
            "payer": _pick_payer(state_data, "medicare_ffs"),
            "component": "sensor",
        })
        # Zero quantity
        challenges.append({
            "patient_id": f"ZERQTY-{file_index:03d}",
            "device_type": "freestyle_libre_3",
            "shipment_date": TODAY - timedelta(days=5),
            "quantity": 0,
            "payer": _pick_payer(state_data, "commercial"),
            "component": "sensor",
        })
        # Duplicate patient_id (second entry should also be processed)
        challenges.append({
            "patient_id": f"DUP-PT-{file_index:03d}",
            "device_type": "dexcom_g7",
            "shipment_date": TODAY - timedelta(days=5),
            "quantity": 1,
            "payer": _pick_payer(state_data, "medicare_ffs"),
            "component": "sensor",
        })
        challenges.append({
            "patient_id": f"DUP-PT-{file_index:03d}",  # intentional duplicate
            "device_type": "freestyle_libre_3",
            "shipment_date": TODAY - timedelta(days=3),
            "quantity": 1,
            "payer": _pick_payer(state_data, "commercial"),
            "component": "sensor",
        })

    return challenges


def generate_patient_batch(
    state_data: dict,
    patient_count: int,
    file_index: int,
    num_files: int,
    challenge_level: str,
) -> list:
    """Generate the main patient batch for one file."""
    patients = []
    payer_types = ["medicare_ffs", "medicare_ffs", "medicare_advantage", "medicaid", "commercial"]
    status_buckets = [
        "active", "active", "active", "active",
        "renewal_soon", "renewal_elevated", "renewal_critical",
        "visit_required", "supply_lapsed",
    ]

    for i in range(patient_count):
        device = random.choice(DEVICE_TYPES)
        payer_type = random.choice(payer_types)
        payer_str = _pick_payer(state_data, payer_type)
        bucket_name = status_buckets[i % len(status_buckets)]
        shipment_date = _random_date_in_bucket(DATE_BUCKETS[bucket_name])
        qty = random.choice([1, 1, 1, 2, 3])

        patient = {
            "patient_id": f"{state_data['branch_prefix']}-{file_index:02d}-{i+1:04d}",
            "device_type": device,
            "shipment_date": shipment_date,
            "quantity": qty,
            "payer": payer_str,
            "component": COMPONENTS.get(device, "sensor"),
        }

        # Add doc fields based on payer type
        if payer_type == "medicare_ffs":
            patient.update({
                "swo_status": random.choice(["On File", "On File", "Pending"]),
                "pecos_verified": random.choice(["Yes", "Yes", "No"]),
                "pa_status": "Not Required",
                "diagnosis_on_file": random.choice(["Yes", "Yes", "No"]),
            })
        elif payer_type == "medicare_advantage":
            patient.update({
                "swo_status": random.choice(["On File", "Pending"]),
                "pecos_verified": random.choice(["Yes", "No"]),
                "pa_status": random.choice(["Approved", "Pending", "Denied"]),
                "diagnosis_on_file": random.choice(["Yes", "No"]),
            })
        else:
            patient.update({
                "swo_status": random.choice(["On File", "Pending", "Expired"]),
                "pa_status": random.choice(["Approved", "Pending", "Not Required"]),
                "diagnosis_on_file": random.choice(["Yes", "No"]),
            })

        # Add visit date for some patients
        if random.random() > 0.4:
            visit_offset = random.randint(30, 200)
            patient["visit_date"] = (TODAY - timedelta(days=visit_offset)).strftime("%Y-%m-%d")
        else:
            patient["visit_date"] = ""

        patients.append(patient)

    return patients


def write_csv_file(
    output_path: str,
    patients: list,
    variant: tuple,
    state_data: dict,
    challenge_level: str,
    file_index: int,
) -> None:
    """Write one CSV file using the given header variant."""
    pid_col, dev_col, date_col, qty_col, payer_col, comp_col, date_fmt, extras_factory = variant
    extras = extras_factory(state_data)

    # Build the ordered list of headers
    core_headers = [h for h in [pid_col, dev_col, date_col, qty_col, payer_col, comp_col] if h is not None]
    extra_headers = list(extras.keys())
    all_headers = core_headers + extra_headers

    # Malformed date files: indices 7 and 14 get some intentionally broken dates
    malformed_indices = set()
    if challenge_level == "high" and file_index in (7, 14):
        malformed_indices = {0, 1, 2}

    with open(output_path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=all_headers, extrasaction="ignore")
        writer.writeheader()

        for row_idx, p in enumerate(patients):
            row = {}

            # Map patient dict fields to this variant's column names
            if pid_col:
                row[pid_col] = p.get("patient_id", "")
            if dev_col:
                row[dev_col] = _device_display(p.get("device_type", "dexcom_g7"), dev_col)
            if date_col:
                sd = p.get("shipment_date", TODAY)
                if isinstance(sd, date):
                    if row_idx in malformed_indices:
                        row[date_col] = random.choice(["not on file", "pending", "TBD"])
                    else:
                        row[date_col] = _format_date(sd, date_fmt)
                else:
                    row[date_col] = str(sd)
            if qty_col:
                row[qty_col] = p.get("quantity", 1)
            if payer_col:
                row[payer_col] = p.get("payer", "Medicare Part B")
            if comp_col:
                row[comp_col] = p.get("component", "sensor")

            # Static extras (branch, prescriber NPI, region, etc.)
            for k, v in extras.items():
                row[k] = v

            writer.writerow(row)


def generate_set(
    state: str,
    supplier: str,
    patients: int,
    challenge_level: str,
    output_dir: str,
    num_files: int = 25,
) -> None:
    """Generate num_files CSV files for the given state."""
    with open(MARKET_DATA_PATH) as f:
        market = json.load(f)

    state_data = market.get(state.upper())
    if not state_data:
        raise ValueError(
            f"State '{state}' not found in market_data.json. Available: {list(market.keys())}"
        )

    state_data = dict(state_data)  # avoid mutating the original
    state_data["supplier_name"] = supplier

    os.makedirs(output_dir, exist_ok=True)

    for i in range(num_files):
        variant = BASE_HEADER_VARIANTS[i % len(BASE_HEADER_VARIANTS)]

        # Core patient batch
        patient_rows = generate_patient_batch(state_data, patients, i, num_files, challenge_level)

        # State-change patients (appear in every file, evolve across the sequence)
        sc_patients = generate_state_change_patients(state_data, i, num_files)
        all_patients = sc_patients + patient_rows

        # Challenge rows appended at end
        if challenge_level in ("medium", "high"):
            challenge_rows = generate_challenge_rows(state_data, challenge_level, i)
            all_patients = all_patients + challenge_rows

        filename = f"batch-{i+1:02d}-{state.lower()}-{challenge_level}.csv"
        output_path = os.path.join(output_dir, filename)

        write_csv_file(output_path, all_patients, variant, state_data, challenge_level, i)
        print(f"  Written: {filename} ({len(all_patients)} rows)")

    print(f"\nGenerated {num_files} files in {output_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate Signal CSV test files")
    parser.add_argument("--state", default="PA")
    parser.add_argument("--supplier", default="Gaboro DME")
    parser.add_argument("--patients", type=int, default=30)
    parser.add_argument("--challenge-level", choices=["low", "medium", "high"], default="high")
    # Accept both --output and --output-dir for compatibility
    parser.add_argument("--output", "--output-dir", dest="output", default="pa-set")
    # Accept both --files and --count for compatibility
    parser.add_argument("--files", "--count", dest="files", type=int, default=25)
    args = parser.parse_args()

    output_path = os.path.join(SCRIPT_DIR, args.output)
    generate_set(
        state=args.state,
        supplier=args.supplier,
        patients=args.patients,
        challenge_level=args.challenge_level,
        output_dir=output_path,
        num_files=args.files,
    )
