#!/usr/bin/env python3
"""
Demo Data Loader for Query Orchestration API

This script imports demo data from a CSV file into the Query Orchestration API.
It creates a demo client, uploads unique PDF documents for each file in the CSV,
and creates field extraction records with both single-value and array fields.

The CSV is expected to have:
- One header row with column names
- Multiple data rows where rows with the same FILE_NAME belong to the same document
- Single-value fields (columns 1-18) that are identical for rows with the same FILE_NAME
- Array fields (columns 19+) that vary per row, forming multiple array items per document

Usage:
    python3 demo_data_loader.py <csv_path> <base_url> [jwt_token]

Arguments:
    csv_path   - Path to the CSV file containing demo data
    base_url   - Base URL of the API (e.g., http://localhost:8080)
    jwt_token  - Optional JWT token for authentication (placeholder for future use)

Example:
    python3 demo_data_loader.py SaaS_Demo_Data_Sam.csv http://localhost:8080
"""

import csv
import json
import sys
import time
import urllib.request
import urllib.error
import urllib.parse
from collections import defaultdict
from datetime import datetime
from typing import Dict, List, Optional, Any, Tuple


# Column indices for single-value fields (0-17, corresponding to columns 1-18 in CSV)
SINGLE_FIELD_COLUMNS = {
    0: "fileName",                   # FILE_NAME
    1: "contractTitle",              # CONTRACT_TITLE
    2: "aareteDerivedAmendmentNum",  # AARETE_DERIVED_AMENDMENT_NUM
    3: "clientName",                 # CLIENT_NAME
    4: "payerName",                  # PAYER_NAME
    5: "payerState",                 # PAYER_STATE
    6: "providerState",              # PROVIDER_STATE
    7: "filenameTin",                # FILENAME_TIN
    8: "provGroupTin",               # PROV_GROUP_TIN
    9: "provGroupNpi",               # PROV_GROUP_NPI
    10: "provGroupNameFull",         # PROV_GROUP_NAME_FULL
    11: "provOtherTin",              # PROV_OTHER_TIN
    12: "provOtherNpi",              # PROV_OTHER_NPI
    13: "provOtherNameFull",         # PROV_OTHER_NAME_FULL
    14: "aareteDerivedEffectiveDt",  # AARETE_DERIVED_EFFECTIVE_DT
    15: "aareteDerivedTerminationDt",# AARETE_DERIVED_TERMINATION_DT
    16: "autoRenewalInd",            # AUTO_RENEWAL_IND
    17: "autoRenewalTerm",           # AUTO_RENEWAL_TERM
}

# Column indices for array fields (18+, corresponding to columns 19+ in CSV)
ARRAY_FIELD_COLUMNS = {
    18: "exhibitTitle",              # EXHIBIT_TITLE
    19: "exhibitPage",               # EXHIBIT_PAGE
    20: "reimbProvTin",              # REIMB_PROV_TIN
    21: "reimbProvNpi",              # REIMB_PROV_NPI
    22: "reimbProvName",             # REIMB_PROV_NAME
    23: "reimbEffectiveDt",          # REIMB_EFFECTIVE_DT
    24: "reimbTerminationDt",        # REIMB_TERMINATION_DT
    25: "aareteDerivedClaimTypeCd",  # AARETE_DERIVED_CLAIM_TYPE_CD
    26: "aareteDerivedProduct",      # AARETE_DERIVED_PRODUCT
    27: "aareteDerivedLob",          # AARETE_DERIVED_LOB
    28: "aareteDerivedProgram",      # AARETE_DERIVED_PROGRAM
    29: "aareteDerivedNetwork",      # AARETE_DERIVED_NETWORK
    30: "lobProgramRelationship",    # LOB_PROGRAM_RELATIONSHIP
    31: "lobProductRelationship",    # LOB_PRODUCT_RELATIONSHIP
    32: "aareteDerivedProvType",     # AARETE_DERIVED_PROV_TYPE
    33: "provTaxonomyCd",            # PROV_TAXONOMY_CD
    34: "provTaxonomyCdDesc",        # PROV_TAXONOMY_CD_DESC
    35: "provSpecialtyCd",           # PROV_SPECIALTY_CD
    36: "provSpecialtyCdDesc",       # PROV_SPECIALTY_CD_DESC
    37: "placeOfServiceCd",          # PLACE_OF_SERVICE_CD
    38: "placeOfServiceCdDesc",      # PLACE_OF_SERVICE_CD_DESC
    39: "billTypeCd",                # BILL_TYPE_CD
    40: "billTypeCdDesc",            # BILL_TYPE_CD_DESC
    41: "patientAgeMin",             # PATIENT_AGE_MIN
    42: "patientAgeMax",             # PATIENT_AGE_MAX
    43: "reimbTerm",                 # REIMB_TERM
    44: "carveoutInd",               # CARVEOUT_IND
    45: "carveoutCd",                # CARVEOUT_CD
    46: "lesserOfInd",               # LESSER_OF_IND
    47: "greaterOfInd",              # GREATER_OF_IND
    48: "aareteDerivedReimbMethod",  # AARETE_DERIVED_REIMB_METHOD
    49: "unitOfMeasure",             # UNIT_OF_MEASURE
    50: "reimbPctRate",              # REIMB_PCT_RATE
    51: "reimbFeeRate",              # REIMB_FEE_RATE
    52: "reimbConversionFactor",     # REIMB_CONVERSION_FACTOR
    53: "triggerCapThresholdAmt",    # TRIGGER_CAP_THRESHOLD_AMT
    54: "triggerBaseThreshold",      # TRIGGER_BASE_THRESHOLD
    55: "defaultInd",                # DEFAULT_IND
    56: "additionDesc",              # ADDITION_DESC
    57: "additionMaxPctRateInc",     # ADDITION_MAX_PCT_RATE_INC
    58: "additionMaxFeeRateInc",     # ADDITION_MAX_FEE_RATE_INC
    59: "aareteDerivedAdditionRateChangeTimeline",  # AARETE_DERIVED_ADDITION_RATE_CHANGE_TIMELINE
    60: "aareteDerivedFeeSchedule",  # AARETE_DERIVED_FEE_SCHEDULE
    61: "aareteDerivedFeeScheduleVersion",  # AARETE_DERIVED_FEE_SCHEDULE_VERSION
    62: "serviceTerm",               # SERVICE_TERM
    63: "cpt4ProcCd",                # CPT4_PROC_CD
    64: "cpt4ProcCdDesc",            # CPT4_PROC_CD_DESC
    65: "cpt4ProcMod",               # CPT4_PROC_MOD
    66: "cpt4ProcModDesc",           # CPT4_PROC_MOD_DESC
    67: "revenueCd",                 # REVENUE_CD
    68: "revenueCdDesc",             # REVENUE_CD_DESC
    69: "diagCd",                    # DIAG_CD
    70: "diagCdDesc",                # DIAG_CD_DESC
    71: "ndcCd",                     # NDC_CD
    72: "ndcCdDesc",                 # NDC_CD_DESC
    73: "claimAdmitTypeCd",          # CLAIM_ADMIT_TYPE_CD
    74: "authAdmitTypeDesc",         # AUTH_ADMIT_TYPE_DESC
    75: "claimStatusCd",             # CLAIM_STATUS_CD
    76: "claimStatusCdDesc",         # CLAIM_STATUS_CD_DESC
    77: "grouperType",               # GROUPER_TYPE
    78: "grouperCd",                 # GROUPER_CD
    79: "grouperCdDesc",             # GROUPER_CD_DESC
    80: "grouperPctRate",            # GROUPER_PCT_RATE
    81: "grouperBaseRate",           # GROUPER_BASE_RATE
    82: "aareteDerivedGrouperVersion",  # AARETE_DERIVED_GROUPER_VERSION
    83: "grouperAlternativeLevelOfCare",  # GROUPER_ALTERNATIVE_LEVEL_OF_CARE
    84: "grouperSeverityInd",        # GROUPER_SEVERITY_IND
    85: "grouperSeverity",           # GROUPER_SEVERITY
    86: "grouperRiskOfMortalitySubclass",  # GROUPER_RISK_OF_MORTALITY_SUBCLASS
    87: "grouperTransferInd",        # GROUPER_TRANSFER_IND
    88: "grouperReadmissionsInd",    # GROUPER_READMISSIONS_IND
    89: "grouperHacInd",             # GROUPER_HAC_IND
    90: "outlierTerm",               # OUTLIER_TERM
    91: "outlierFirstDollarInd",     # OUTLIER_FIRST_DOLLAR_IND
    92: "rangeNbrDays",              # RANGE_NBR_DAYS
    93: "outlierFixedLossNbrDaysThreshold",  # OUTLIER_FIXED_LOSS_NBR_DAYS_THRESHOLD
    94: "outlierFixedLossThreshold", # OUTLIER_FIXED_LOSS_THRESHOLD
    95: "outlierMaximum",            # OUTLIER_MAXIMUM
    96: "outlierMaximumFrequency",   # OUTLIER_MAXIMUM_FREQUENCY
    97: "outlierPctRate",            # OUTLIER_PCT_RATE
    98: "outlierExclusionCd",        # OUTLIER_EXCLUSION_CD
    99: "outlierExclusionCdDesc",    # OUTLIER_EXCLUSION_CD_DESC
    100: "facilityAdjustmentTerm",   # FACILITY_ADJUSTMENT_TERM
    101: "dshInd",                   # DSH_IND
    102: "dshPctRate",               # DSH_PCT_RATE
    103: "dshFeeRate",               # DSH_FEE_RATE
    104: "imeInd",                   # IME_IND
    105: "imePctRate",               # IME_PCT_RATE
    106: "imeFeeRate",               # IME_FEE_RATE
    107: "ntapInd",                  # NTAP_IND
    108: "ntapPctRate",              # NTAP_PCT_RATE
    109: "ntapFeeRate",              # NTAP_FEE_RATE
    110: "ucInd",                    # UC_IND
    111: "ucPctRate",                # UC_PCT_RATE
    112: "ucFeeRate",                # UC_FEE_RATE
    113: "gmeInd",                   # GME_IND
    114: "gmePctRate",               # GME_PCT_RATE
    115: "gmeFeeRate",               # GME_FEE_RATE
    116: "rateEscalatorInd",         # RATE_ESCALATOR_IND
    117: "rateEscalatorDesc",        # RATE_ESCALATOR_DESC
    118: "rateEscalatorMaxRateIncPct",  # RATE_ESCALATOR_MAX_RATE_INC_PCT
    119: "rateEscalatorRateChangeTimeline",  # RATE_ESCALATOR_RATE_CHANGE_TIMELINE
    120: "stopLossTerm",             # STOP_LOSS_TERM
    121: "stopLossFirstDollarInd",   # STOP_LOSS_FIRST_DOLLAR_IND
    122: "stopLossRangeNbrDays",     # STOP_LOSS_RANGE_NBR_DAYS
    123: "stopLossFixedLossThreshold",  # STOP_LOSS_FIXED_LOSS_THRESHOLD
    124: "stopLossMaximum",          # STOP_LOSS_MAXIMUM
    125: "stopLossMaximumFrequency", # STOP_LOSS_MAXIMUM_FREQUENCY
    126: "stopLossDailyMaxRate",     # STOP_LOSS_DAILY_MAX_RATE
    127: "stopLossPctRateOnExcessCharges",  # STOP_LOSS_PCT_RATE_ON_EXCESS_CHARGES
    128: "stopLossExclusionCd",      # STOP_LOSS_EXCLUSION_CD
    129: "stopLossExclusionDesc",    # STOP_LOSS_EXCLUSION_DESC
}

# Fields that should be converted to boolean
BOOLEAN_FIELDS = {
    "autoRenewalInd", "carveoutInd", "lesserOfInd", "greaterOfInd",
    "defaultInd", "grouperSeverityInd", "grouperTransferInd",
    "grouperReadmissionsInd", "grouperHacInd", "outlierFirstDollarInd",
    "dshInd", "imeInd", "ntapInd", "ucInd", "gmeInd", "rateEscalatorInd",
    "stopLossFirstDollarInd"
}

# Fields that should be converted to integer
INTEGER_FIELDS = {"aareteDerivedAmendmentNum"}

# Fields that should be converted to float/number
NUMERIC_FIELDS = {
    "reimbPctRate", "reimbFeeRate", "reimbConversionFactor",
    "triggerCapThresholdAmt", "triggerBaseThreshold",
    "additionMaxPctRateInc", "additionMaxFeeRateInc",
    "grouperPctRate", "grouperBaseRate",
    "outlierFixedLossNbrDaysThreshold", "outlierFixedLossThreshold",
    "outlierMaximum", "outlierMaximumFrequency", "outlierPctRate",
    "dshPctRate", "dshFeeRate", "imePctRate", "imeFeeRate",
    "ntapPctRate", "ntapFeeRate", "ucPctRate", "ucFeeRate",
    "gmePctRate", "gmeFeeRate", "rateEscalatorMaxRateIncPct",
    "rateEscalatorRateChangeTimeline", "stopLossRangeNbrDays",
    "stopLossFixedLossThreshold", "stopLossMaximum",
    "stopLossMaximumFrequency", "stopLossDailyMaxRate",
    "stopLossPctRateOnExcessCharges"
}

# Fields that should be formatted as dates (YYYY-MM-DD)
DATE_FIELDS = {
    "aareteDerivedEffectiveDt", "aareteDerivedTerminationDt",
    "reimbEffectiveDt", "reimbTerminationDt"
}


def read_token_file(filepath: str) -> Optional[str]:
    """
    Read a JWT token from a file.

    Args:
        filepath: Path to the file containing the token.

    Returns:
        The token string (stripped of whitespace) or None if file cannot be read.
    """
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            token = f.read().strip()
            if token:
                return token
            return None
    except (IOError, OSError) as e:
        print(f"  [ERROR] Failed to read token file: {e}")
        return None


def print_header(msg: str) -> None:
    """Print a formatted section header."""
    print(f"\n{'=' * 60}")
    print(f"  {msg}")
    print(f"{'=' * 60}")


def print_success(msg: str) -> None:
    """Print a success message."""
    print(f"  [OK] {msg}")


def print_error(msg: str) -> None:
    """Print an error message."""
    print(f"  [ERROR] {msg}")


def print_info(msg: str) -> None:
    """Print an informational message."""
    print(f"  {msg}")


def parse_boolean(value: str) -> Optional[bool]:
    """
    Parse a string value to boolean.

    Args:
        value: String value like 'Y', 'N', 'true', 'false', etc.

    Returns:
        Boolean value or None if empty/invalid.
    """
    if not value or value.strip() == "":
        return None
    v = value.strip().upper()
    if v in ("Y", "YES", "TRUE", "1"):
        return True
    if v in ("N", "NO", "FALSE", "0"):
        return False
    return None


def parse_integer(value: str) -> Optional[int]:
    """
    Parse a string value to integer.

    Args:
        value: String representation of an integer.

    Returns:
        Integer value or None if empty/invalid.
    """
    if not value or value.strip() == "":
        return None
    try:
        return int(value.strip())
    except ValueError:
        return None


def parse_numeric(value: str) -> Optional[float]:
    """
    Parse a string value to float, handling percentages and special characters.

    Args:
        value: String representation of a number (may include % or other chars).

    Returns:
        Float value or None if empty/invalid.
    """
    if not value or value.strip() == "":
        return None
    v = value.strip()
    # Remove percentage sign and other common suffixes
    v = v.replace("%", "").replace("$", "").replace(",", "").strip()
    try:
        return float(v)
    except ValueError:
        return None


def parse_date(value: str) -> Optional[str]:
    """
    Parse a date string and convert to YYYY-MM-DD format.

    Args:
        value: Date string in various formats (M/D/YYYY, MM/DD/YYYY, etc.)

    Returns:
        Date in YYYY-MM-DD format or None if empty/invalid.
    """
    if not value or value.strip() == "":
        return None
    v = value.strip()

    # Try common date formats
    formats = [
        "%m/%d/%Y",    # 1/1/2025 or 01/01/2025
        "%Y-%m-%d",    # 2025-01-01
        "%d/%m/%Y",    # 01/01/2025 (European)
        "%Y/%m/%d",    # 2025/01/01
    ]

    for fmt in formats:
        try:
            dt = datetime.strptime(v, fmt)
            return dt.strftime("%Y-%m-%d")
        except ValueError:
            continue

    return None


def convert_field_value(field_name: str, value: str) -> Any:
    """
    Convert a field value to its appropriate type based on field name.

    Args:
        field_name: The API field name.
        value: The raw string value from CSV.

    Returns:
        Converted value in appropriate type.
    """
    if field_name in BOOLEAN_FIELDS:
        return parse_boolean(value)
    if field_name in INTEGER_FIELDS:
        return parse_integer(value)
    if field_name in NUMERIC_FIELDS:
        return parse_numeric(value)
    if field_name in DATE_FIELDS:
        return parse_date(value)
    # Default: return as string (None if empty)
    if not value or value.strip() == "":
        return None
    return value.strip()


def api_request(
    base_url: str,
    endpoint: str,
    method: str = "GET",
    data: Optional[Dict] = None,
    jwt_token: Optional[str] = None,
    max_retries: int = 3
) -> Tuple[int, Any]:
    """
    Make an HTTP request to the API with retry logic for rate limiting.

    Args:
        base_url: Base URL of the API.
        endpoint: API endpoint path.
        method: HTTP method (GET, POST, etc.)
        data: Request body data (will be JSON encoded).
        jwt_token: Optional JWT token for authorization.
        max_retries: Maximum number of retries for rate limiting.

    Returns:
        Tuple of (HTTP status code, response body as dict/list or error string).
    """
    url = f"{base_url.rstrip('/')}{endpoint}"
    headers = {"Content-Type": "application/json"}

    if jwt_token:
        headers["Authorization"] = f"Bearer {jwt_token}"

    body = None
    if data is not None:
        body = json.dumps(data).encode("utf-8")

    delay = 1.0
    for attempt in range(max_retries + 1):
        req = urllib.request.Request(url, data=body, headers=headers, method=method)

        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                status = resp.status
                content = resp.read().decode("utf-8")
                if content:
                    return status, json.loads(content)
                return status, {}
        except urllib.error.HTTPError as e:
            content = e.read().decode("utf-8") if e.fp else ""
            # Retry on rate limit
            if e.code == 429 and attempt < max_retries:
                time.sleep(delay)
                delay *= 2  # Exponential backoff
                continue
            try:
                return e.code, json.loads(content) if content else {"error": str(e)}
            except json.JSONDecodeError:
                return e.code, {"error": content or str(e)}
        except urllib.error.URLError as e:
            return 0, {"error": str(e)}

    return 429, {"error": "Rate limit exceeded after retries"}


def upload_document(
    base_url: str,
    client_id: str,
    filename: str,
    content: bytes,
    jwt_token: Optional[str] = None
) -> Tuple[int, Any]:
    """
    Upload a document via multipart/form-data.

    Args:
        base_url: Base URL of the API.
        client_id: Client ID to upload document to.
        filename: Name of the file being uploaded.
        content: File content as bytes.
        jwt_token: Optional JWT token for authorization.

    Returns:
        Tuple of (HTTP status code, response body).
    """
    url = f"{base_url.rstrip('/')}/client/{client_id}/document"
    boundary = f"----WebKitFormBoundary{int(time.time() * 1000)}"

    # Build multipart form data - content type must be application/octet-stream
    body_parts = []
    body_parts.append(f"--{boundary}".encode())
    body_parts.append(
        f'Content-Disposition: form-data; name="file"; filename="{filename}"'.encode()
    )
    body_parts.append(b"Content-Type: application/octet-stream")
    body_parts.append(b"")
    body_parts.append(content)
    body_parts.append(f"--{boundary}--".encode())

    body = b"\r\n".join(body_parts)

    headers = {
        "Content-Type": f"multipart/form-data; boundary={boundary}",
        "Content-Length": str(len(body))
    }

    if jwt_token:
        headers["Authorization"] = f"Bearer {jwt_token}"

    req = urllib.request.Request(url, data=body, headers=headers, method="POST")

    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            status = resp.status
            resp_content = resp.read().decode("utf-8")
            if resp_content:
                return status, json.loads(resp_content)
            return status, {}
    except urllib.error.HTTPError as e:
        err_content = e.read().decode("utf-8") if e.fp else ""
        try:
            return e.code, json.loads(err_content) if err_content else {"error": str(e)}
        except json.JSONDecodeError:
            return e.code, {"error": err_content or str(e)}
    except urllib.error.URLError as e:
        return 0, {"error": str(e)}


def create_pdf(unique_text: str) -> bytes:
    """
    Create a minimal valid PDF file with unique content.

    The PDF contains the unique_text which makes its hash unique.

    Args:
        unique_text: Text to embed in the PDF to make it unique.

    Returns:
        PDF file content as bytes.
    """
    stream_content = f"""BT
/F1 12 Tf
50 700 Td
({unique_text}) Tj
ET"""
    stream_length = len(stream_content)

    pdf = f"""%PDF-1.4
%\xe2\xe3\xcf\xd3
1 0 obj
<<
  /Type /Catalog
  /Pages 2 0 R
>>
endobj
2 0 obj
<<
  /Type /Pages
  /Kids [3 0 R]
  /Count 1
>>
endobj
3 0 obj
<<
  /Type /Page
  /Parent 2 0 R
  /MediaBox [0 0 612 792]
  /Resources <<
    /Font <<
      /F1 <<
        /Type /Font
        /Subtype /Type1
        /BaseFont /Helvetica
      >>
    >>
  >>
  /Contents 4 0 R
>>
endobj
4 0 obj
<<
  /Length {stream_length}
>>
stream
{stream_content}
endstream
endobj
xref
0 5
0000000000 65535 f
0000000015 00000 n
0000000066 00000 n
0000000125 00000 n
0000000330 00000 n
trailer
<<
  /Size 5
  /Root 1 0 R
>>
startxref
430
%%EOF"""

    return pdf.encode("latin-1")


def read_csv_file(csv_path: str) -> Tuple[List[str], List[List[str]]]:
    """
    Read and parse the CSV file.

    Args:
        csv_path: Path to the CSV file.

    Returns:
        Tuple of (headers list, data rows list).
    """
    with open(csv_path, "r", encoding="utf-8-sig") as f:
        reader = csv.reader(f)
        headers = next(reader)
        rows = list(reader)
    return headers, rows


def group_rows_by_filename(rows: List[List[str]]) -> Dict[str, List[List[str]]]:
    """
    Group CSV rows by FILE_NAME (first column).

    Args:
        rows: List of CSV data rows.

    Returns:
        Dictionary mapping file names to their rows.
    """
    groups = defaultdict(list)
    for row in rows:
        if row:  # Skip empty rows
            filename = row[0] if row else ""
            if filename:
                groups[filename].append(row)
    return dict(groups)


def extract_single_fields(row: List[str]) -> Dict[str, Any]:
    """
    Extract single-value fields from a CSV row.

    Args:
        row: A CSV data row.

    Returns:
        Dictionary of single field name -> value.
    """
    fields = {}
    for col_idx, field_name in SINGLE_FIELD_COLUMNS.items():
        if col_idx < len(row):
            value = convert_field_value(field_name, row[col_idx])
            if value is not None:
                fields[field_name] = value
    return fields


def extract_array_field_item(row: List[str]) -> Dict[str, Any]:
    """
    Extract array field item from a CSV row.

    Args:
        row: A CSV data row.

    Returns:
        Dictionary of array field name -> value.
    """
    item = {}
    for col_idx, field_name in ARRAY_FIELD_COLUMNS.items():
        if col_idx < len(row):
            value = convert_field_value(field_name, row[col_idx])
            if value is not None:
                item[field_name] = value
    return item


def create_client(base_url: str, jwt_token: Optional[str] = None) -> Optional[str]:
    """
    Create a demo client via the API.

    Args:
        base_url: Base URL of the API.
        jwt_token: Optional JWT token.

    Returns:
        Client ID if successful, None otherwise.
    """
    timestamp = int(time.time())
    client_data = {
        "id": f"demo-data-loader-{timestamp}",
        "name": f"Demo Data Loader Client {timestamp}"
    }

    print_info(f"Creating client: {client_data['name']}")
    status, response = api_request(base_url, "/client", "POST", client_data, jwt_token)

    if status == 201:
        client_id = response.get("id")
        print_success(f"Client created: {client_id}")
        return client_id
    else:
        print_error(f"Failed to create client: {status} - {response}")
        return None


def upload_documents(
    base_url: str,
    client_id: str,
    filenames: List[str],
    jwt_token: Optional[str] = None
) -> Dict[str, str]:
    """
    Upload unique PDF documents for each filename.

    Args:
        base_url: Base URL of the API.
        client_id: Client ID to upload to.
        filenames: List of unique filenames from CSV.
        jwt_token: Optional JWT token.

    Returns:
        Dictionary mapping original filename to document ID.
    """
    filename_to_doc_id = {}

    for i, filename in enumerate(filenames):
        # Create unique PDF with timestamp and index
        unique_text = f"Demo Document: {filename} - Created: {datetime.now().isoformat()} - Index: {i}"
        pdf_content = create_pdf(unique_text)
        pdf_filename = f"{filename}.pdf"

        print_info(f"Uploading: {pdf_filename}")
        status, response = upload_document(
            base_url, client_id, pdf_filename, pdf_content, jwt_token
        )

        if status == 200:
            # Response may contain document info or we need to fetch it
            doc_id = response.get("id") or response.get("document_id")
            if doc_id:
                filename_to_doc_id[filename] = doc_id
                print_success(f"Uploaded: {pdf_filename} -> {doc_id}")
            else:
                print_info(f"Uploaded {pdf_filename} (processing async)")
                # Will need to retrieve document ID later
        else:
            print_error(f"Failed to upload {pdf_filename}: {status} - {response}")

        # Delay to avoid rate limiting
        time.sleep(0.5)

    return filename_to_doc_id


def wait_for_documents(
    base_url: str,
    client_id: str,
    expected_count: int,
    jwt_token: Optional[str] = None,
    max_attempts: int = 30
) -> List[Dict]:
    """
    Wait for documents to be processed and return their details.

    Args:
        base_url: Base URL of the API.
        client_id: Client ID to query.
        expected_count: Number of documents expected.
        jwt_token: Optional JWT token.
        max_attempts: Maximum polling attempts.

    Returns:
        List of document records with full details including filename.
    """
    print_info(f"Waiting for {expected_count} documents to be processed...")

    doc_count = 0
    doc_ids = []

    for attempt in range(max_attempts):
        status, response = api_request(
            base_url, f"/client/{client_id}/document", "GET", jwt_token=jwt_token
        )

        if status == 200 and isinstance(response, list):
            doc_count = len(response)
            doc_ids = [d.get("id") for d in response if d.get("id")]
            print_info(f"  Attempt {attempt + 1}: {doc_count}/{expected_count} documents")

            if doc_count >= expected_count:
                print_success(f"All {doc_count} documents are available")
                break

        time.sleep(1)

    if doc_count == 0:
        print_info("Timeout: No documents available")
        return []

    # Fetch full details for each document to get filenames
    print_info("Fetching document details...")
    full_documents = []
    for doc_id in doc_ids:
        status, doc_detail = api_request(
            base_url, f"/document/{doc_id}", "GET", jwt_token=jwt_token
        )
        if status == 200:
            full_documents.append(doc_detail)
        time.sleep(0.05)  # Small delay to avoid rate limiting

    return full_documents


def create_field_extraction(
    base_url: str,
    document_id: str,
    single_fields: Dict[str, Any],
    array_fields: List[Dict[str, Any]],
    jwt_token: Optional[str] = None
) -> bool:
    """
    Create a field extraction record for a document.

    Args:
        base_url: Base URL of the API.
        document_id: Document ID to create extraction for.
        single_fields: Single-value fields dictionary.
        array_fields: List of array field items.
        jwt_token: Optional JWT token.

    Returns:
        True if successful, False otherwise.
    """
    extraction_data = {
        "documentId": document_id,
        "singleFields": single_fields,
        "arrayFields": array_fields,
        "createdBy": "demo-data-loader@example.com"
    }

    status, response = api_request(
        base_url, "/field-extractions", "POST", extraction_data, jwt_token
    )

    if status == 201:
        version = response.get("version", "?")
        print_success(f"Field extraction created: version={version}")
        return True
    else:
        print_error(f"Failed to create extraction: {status} - {response}")
        return False


def verify_field_extractions(
    base_url: str,
    document_ids: List[str],
    csv_data: Dict[str, Dict],
    jwt_token: Optional[str] = None
) -> Tuple[int, int]:
    """
    Verify that field extractions match the CSV input data.

    Args:
        base_url: Base URL of the API.
        document_ids: List of document IDs to verify.
        csv_data: Dictionary mapping filename to expected extraction data.
        jwt_token: Optional JWT token.

    Returns:
        Tuple of (success_count, failure_count).
    """
    success_count = 0
    failure_count = 0

    for doc_id in document_ids:
        status, response = api_request(
            base_url, f"/field-extractions?documentId={doc_id}", "GET", jwt_token=jwt_token
        )

        if status == 200:
            single_fields = response.get("singleFields", {})
            array_fields = response.get("arrayFields", [])
            filename = single_fields.get("fileName", "")

            # Basic validation: check that fields exist
            if single_fields and array_fields:
                print_success(
                    f"Verified {doc_id}: {len(single_fields)} single fields, "
                    f"{len(array_fields)} array items"
                )
                success_count += 1
            else:
                print_error(f"Incomplete extraction for {doc_id}")
                failure_count += 1
        elif status == 404:
            print_error(f"No extraction found for {doc_id}")
            failure_count += 1
        else:
            print_error(f"Failed to verify {doc_id}: {status}")
            failure_count += 1

    return success_count, failure_count


def parse_arguments() -> Tuple[str, str, Optional[str]]:
    """
    Parse command line arguments.

    Supports:
      - positional: csv_path, base_url, [jwt_token]
      - named: tokenfile=<path> to read token from file

    Returns:
        Tuple of (csv_path, base_url, jwt_token or None).
    """
    if len(sys.argv) < 3:
        print("Usage: python3 demo_data_loader.py <csv_path> <base_url> [jwt_token | tokenfile=<path>]")
        print("")
        print("Arguments:")
        print("  csv_path   - Path to the CSV file containing demo data")
        print("  base_url   - Base URL of the API (e.g., http://localhost:8080)")
        print("  jwt_token  - Optional JWT token for authentication")
        print("  tokenfile= - Optional path to file containing JWT token")
        print("")
        print("Examples:")
        print("  python3 demo_data_loader.py SaaS_Demo_Data_Sam.csv http://localhost:8080")
        print("  python3 demo_data_loader.py SaaS_Demo_Data_Sam.csv http://localhost:8080 tokenfile=./token.txt")
        sys.exit(1)

    csv_path = sys.argv[1]
    base_url = sys.argv[2]
    jwt_token = None

    # Check for optional third argument
    if len(sys.argv) > 3:
        arg3 = sys.argv[3]
        if arg3.startswith("tokenfile="):
            # Read token from file
            token_path = arg3[len("tokenfile="):]
            jwt_token = read_token_file(token_path)
            if jwt_token is None:
                print(f"  [ERROR] Could not read token from: {token_path}")
                sys.exit(1)
        else:
            # Treat as direct token
            jwt_token = arg3

    return csv_path, base_url, jwt_token


def main():
    """Main entry point for the demo data loader script."""
    # Parse command line arguments
    csv_path, base_url, jwt_token = parse_arguments()

    print_header("Demo Data Loader")
    print_info(f"CSV File: {csv_path}")
    print_info(f"Base URL: {base_url}")
    if jwt_token:
        # Show truncated token for security
        token_preview = jwt_token[:20] + "..." if len(jwt_token) > 20 else jwt_token
        print_info(f"JWT Token: {token_preview}")
    else:
        print_info("JWT Token: not provided")

    # Step 1: Read and parse CSV
    print_header("Step 1: Reading CSV File")
    try:
        headers, rows = read_csv_file(csv_path)
        print_success(f"Read {len(rows)} data rows with {len(headers)} columns")
    except Exception as e:
        print_error(f"Failed to read CSV: {e}")
        sys.exit(1)

    # Step 2: Group rows by filename
    print_header("Step 2: Analyzing Data")
    file_groups = group_rows_by_filename(rows)
    unique_files = list(file_groups.keys())
    print_success(f"Found {len(unique_files)} unique documents to import")

    for filename in unique_files[:5]:  # Show first 5
        row_count = len(file_groups[filename])
        print_info(f"  - {filename}: {row_count} array item(s)")
    if len(unique_files) > 5:
        print_info(f"  ... and {len(unique_files) - 5} more")

    # Step 3: Create demo client
    print_header("Step 3: Creating Demo Client")
    client_id = create_client(base_url, jwt_token)
    if not client_id:
        print_error("Cannot proceed without a client")
        sys.exit(1)

    # Step 4: Upload documents
    print_header("Step 4: Uploading Documents")
    filename_to_doc_id = upload_documents(base_url, client_id, unique_files, jwt_token)

    # Step 5: Wait for documents to be processed
    print_header("Step 5: Waiting for Document Processing")
    documents = wait_for_documents(base_url, client_id, len(unique_files), jwt_token)

    if not documents:
        print_error("No documents available after upload")
        sys.exit(1)

    # Build mapping from filename to document ID
    for doc in documents:
        doc_filename = doc.get("filename", "")
        doc_id = doc.get("id", "")
        if not doc_id or not doc_filename:
            continue
        # Remove .pdf extension for matching
        base_name = doc_filename
        if base_name.lower().endswith(".pdf"):
            base_name = base_name[:-4]
        if base_name in file_groups:
            filename_to_doc_id[base_name] = doc_id

    print_success(f"Mapped {len(filename_to_doc_id)} filenames to document IDs")

    # Debug: show any unmatched files
    unmatched_csv = [f for f in file_groups.keys() if f not in filename_to_doc_id]
    if unmatched_csv:
        print_info(f"Unmatched CSV files ({len(unmatched_csv)}):")
        for f in unmatched_csv[:5]:
            print_info(f"  - {f}")

    # Step 6: Create field extractions
    print_header("Step 6: Creating Field Extractions")
    extraction_success = 0
    extraction_failed = 0

    for filename, doc_rows in file_groups.items():
        doc_id = filename_to_doc_id.get(filename)
        if not doc_id:
            print_error(f"No document ID for {filename}")
            extraction_failed += 1
            continue

        print_info(f"Processing: {filename}")

        # Extract single fields from first row (they should be same for all rows)
        single_fields = extract_single_fields(doc_rows[0])

        # Extract array fields from all rows
        array_fields = []
        for row in doc_rows:
            array_item = extract_array_field_item(row)
            if array_item:  # Only add non-empty items
                array_fields.append(array_item)

        # Create field extraction
        if create_field_extraction(base_url, doc_id, single_fields, array_fields, jwt_token):
            extraction_success += 1
        else:
            extraction_failed += 1

        # Delay to avoid rate limiting
        time.sleep(0.5)

    print_success(f"Created {extraction_success} extractions, {extraction_failed} failed")

    # Step 7: Verify extractions
    print_header("Step 7: Verifying Field Extractions")
    doc_ids_to_verify = list(filename_to_doc_id.values())
    verified, failed = verify_field_extractions(base_url, doc_ids_to_verify, {}, jwt_token)

    # Final summary
    print_header("Summary")
    print_info(f"Client ID: {client_id}")
    print_info(f"Documents uploaded: {len(documents)}")
    print_info(f"Field extractions created: {extraction_success}")
    print_info(f"Verifications passed: {verified}")
    print_info(f"Verifications failed: {failed}")

    if extraction_failed == 0 and failed == 0:
        print_success("All operations completed successfully!")
        sys.exit(0)
    else:
        print_error(f"Some operations failed: {extraction_failed} extractions, {failed} verifications")
        sys.exit(1)


if __name__ == "__main__":
    main()