#!/usr/bin/env python3
"""
Build Healthie form from Cognito_to_Healthie_Form_Mapping.yaml

Three-phase API flow:
  Phase 1: createCustomModuleForm → creates the form container
  Phase 2: createCustomModule (×N) → adds each module (no options)
  Phase 3: updateCustomModule (×N) → sets options_array for choice/checkbox

Usage:
    # Preview mutations
    python scripts/build_healthie_form.py

    # Execute against Healthie sandbox (reads .env.local)
    python scripts/build_healthie_form.py --execute

    # Dry-run: show what would be sent
    python scripts/build_healthie_form.py --dry-run
"""

import argparse
import json
import os
import sys
import time
import urllib.request
import urllib.error

import yaml

YAML_PATH = "docs/intake/Cognito_to_Healthie_Form_Mapping.yaml"


def load_dotenv(path: str = ".env.local") -> None:
    if not os.path.exists(path):
        return
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" in line:
                key, _, value = line.partition("=")
                key = key.strip()
                value = value.strip().strip("\"'")
                if key and key not in os.environ:
                    os.environ[key] = value


load_dotenv()

HEALTHIE_API_BASE_URL = os.environ.get(
    "HEALTHIE_API_BASE_URL", "https://api.gethealthie.com/graphql"
)

# Healthie custom modules support: text, choice, checkbox, range, date, time
# file_upload is NOT a valid mod_type — mapped to text with metadata
# read_only: displays HTML content from options parameter
# agree_to_above: checkbox requiring acknowledgment before submission
# Built-in types: name, sex, dob, phone, location
TYPE_MAP = {
    "text": "text",
    "choice": "radio",  # All choice fields use radio buttons
    "checkbox": "checkbox",
    "date": "date",
    "file_upload": "text",
    "read_only": "read_only",
    "agree_to_above": "agree_to_above",
    "label": "label",
    "name": "name",
    "sex": "sex",
    "dob": "dob",
    "phone": "phone",
    "location": "location",
    "composite": None,
}

# Section labels for Cognito form sections
SECTION_LABELS = {
    "EligibilityAssessment": "Eligibility Assessment Agreement",
    "SafetyScreening2": "Safety Screening Acknowledgment",
    "PreEnrollmentACCESSDataRecordsParticipationReadinessAcknowledgment2": "Pre-Enrollment Acknowledgment",
    "BasicInfo": "Personal Information",
    "WholeSystem": "Whole System Health Assessment",
    "DigestionIntakeEnergyPattern": "Digestion, Intake & Energy Pattern",
    "BowelUrinaryFluidPattern": "Bowel, Urinary & Fluid Pattern",
    "Female": "Female Health Assessment",
    "Male": "Male Health Assessment",
    "MouthThroatNoseEyesEars": "Mouth, Throat, Nose, Eyes & Ears",
    "SkinSweatingSurfaceReadout": "Skin, Sweating & Surface Readout",
    "SleepStressEmotionalState": "Sleep, Stress & Emotional State",
    "BodyMetrics": "Body Metrics",
    "PainInjuryMusculoskeletalFunctionalBoundary": "Pain, Injury & Musculoskeletal",
    "MedicalHistoryMedicationSafetyReview": "Medical History & Medication Safety",
    "VascularPlaqueAtherosclerosisVascularNarrowing": "Vascular & Atherosclerosis Reports",
    "FattyLiver": "Fatty Liver Reports",
    "KidneyUrineRenalFunctionReports": "Kidney & Renal Function Reports",
    "InfectionImmuneEndocrineAxisReview": "Infection, Immune & Endocrine Review",
    "ReportUpload": "Baseline Report Upload",
    "FinalReviewSubmission": "Final Review & Submission",
}

COMPOSITES = {
    "NameEntryNumber": [
        {"key": "First", "label": "First Name"},
        {"key": "Middle", "label": "Middle Name"},
        {"key": "Last", "label": "Last Name"},
    ],
    "Address": [
        {"key": "Line1", "label": "Address Line 1"},
        {"key": "Line2", "label": "Address Line 2"},
        {"key": "City", "label": "City"},
        {"key": "State", "label": "State"},
        {"key": "PostalCode", "label": "ZIP / Postal Code"},
        {"key": "Country", "label": "Country"},
    ],
}


def load_yaml(path: str) -> dict:
    with open(path, "r") as f:
        return yaml.safe_load(f)


def build_module_spec(field: dict) -> list[dict]:
    """Convert YAML field → createCustomModule spec."""
    ft = field.get("healthie_type", "text")
    mod_type = TYPE_MAP.get(ft)
    if mod_type is None:
        return []

    spec = {
        "mod_type": mod_type,
        "label": field["label"],
        "required": field.get("required", False),
        "external_id": field["cognito_field_id"],
    }

    helptext = field.get("helptext")
    if helptext:
        spec["sublabel"] = helptext

    # read_only: HTML content goes in options parameter (always set, even placeholder)
    # Healthie's UI crashes if options is null on this module type
    if mod_type == "read_only":
        html_content = field.get("options", "")
        if html_content:
            spec["options"] = html_content
        else:
            spec["options"] = "<p>[Agreement text to be provided]</p>"

    # Built-in sex type requires options despite being a built-in
    elif mod_type == "sex":
        spec["options"] = field.get("options", "Male\nFemale")

    # agree_to_above: no extra params needed
    # label: no extra params needed

    choices = field.get("choices")
    if choices and mod_type in ("radio", "checkbox"):
        spec["options"] = "\n".join(choices)

    max_len = field.get("max_length")
    if max_len:
        meta = {"max_length": max_len}
        if ft == "file_upload":
            meta["_file_upload_placeholder"] = True
        spec["metadata"] = json.dumps(meta)

    return [spec]


def build_composite_specs(field: dict) -> list[dict]:
    parent_id = field["cognito_field_id"]
    sub_fields = COMPOSITES.get(parent_id, [])
    required_keys = {"First", "Last", "Line1", "City", "State", "PostalCode"}

    return [
        {
            "mod_type": "text",
            "label": sub["label"],
            "required": sub["key"] in required_keys,
            "external_id": f"{parent_id}_{sub['key']}",
        }
        for sub in sub_fields
    ]


def build_all_specs(data: dict) -> list[dict]:
    specs = []
    page_order = data.get("page_order", {})
    skip_keys = {"form_name", "form_version", "page_order",
                 "conditional_visibility", "top_level_fields"}

    # Sections that use read_only as their header — don't add separate label
    agreement_sections = {"EligibilityAssessment", "SafetyScreening2",
                          "PreEnrollmentACCESSDataRecordsParticipationReadinessAcknowledgment2"}

    # Collect all sections (including agreement and BasicInfo)
    sections = []
    for key, val in data.items():
        if key in skip_keys:
            continue
        if not isinstance(val, dict) or "fields" not in val:
            continue
        page_num = page_order.get(key, 999)
        sections.append((page_num, key, val))

    sections.sort(key=lambda x: x[0])

    for _, section_key, section_data in sections:
        # Add section label header (skip agreement sections — they use read_only)
        if section_key not in agreement_sections:
            section_label = SECTION_LABELS.get(section_key, section_data.get("section_label", section_key))
            specs.append({
                "mod_type": "label",
                "label": section_label,
                "required": False,
                "external_id": f"__section_{section_key}",
                "options": "",
            })

        for field in section_data.get("fields", []):
            specs.extend(build_module_spec(field))

    return specs


# ---------------------------------------------------------------------------
# GraphQL helpers
# ---------------------------------------------------------------------------

def gql_value(v) -> str:
    if v is None:
        return "null"
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, int):
        return str(v)
    if isinstance(v, str):
        escaped = v.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")
        return f'"{escaped}"'
    return f'"{v}"'


def gql_input(d: dict) -> str:
    parts = []
    for k, v in d.items():
        if v is not None:
            parts.append(f"{k}: {gql_value(v)}")
    return "{" + ", ".join(parts) + "}"


def gql_input_mixed(d: dict) -> str:
    """Like gql_input but handles options_array as a list."""
    parts = []
    for k, v in d.items():
        if v is None:
            continue
        if k == "options_array" and isinstance(v, list):
            parts.append(f"options_array: {gql_list(v)}")
        else:
            parts.append(f"{k}: {gql_value(v)}")
    return "{" + ", ".join(parts) + "}"


# ---------------------------------------------------------------------------
# Mutation builders
# ---------------------------------------------------------------------------

def phase1_mutation(form_name: str) -> str:
    return (
        'mutation CreateIntakeForm {\n'
        '  createCustomModuleForm(input: {\n'
        f'    name: "{form_name}"\n'
        '    use_for_charting: true\n'
        '    use_for_program: true\n'
        '  }) {\n'
        '    customModuleForm {\n'
        '      id\n'
        '      name\n'
        '    }\n'
        '    messages {\n'
        '      field\n'
        '      message\n'
        '    }\n'
        '  }\n'
        '}'
    )


def phase2_mutation(form_id: str, specs: list[dict], start_index: int = 0) -> str:
    """Create modules (with inline options as newline-separated string)."""
    aliases = []
    for i, spec in enumerate(specs, start=start_index):
        alias = f"m{i}"
        create_spec = dict(spec)
        create_spec["custom_module_form_id"] = form_id
        create_spec["index"] = i
        aliases.append(
            f'  {alias}: createCustomModule(input: {gql_input(create_spec)}) {{\n'
            '    customModule {\n'
            '      id\n'
            '      label\n'
            '      options\n'
            '    }\n'
            '    messages {\n'
            '      field\n'
            '      message\n'
            '    }\n'
            '  }'
        )
    return 'mutation BuildModules {\n' + "\n".join(aliases) + '\n}'


# ---------------------------------------------------------------------------
# API
# ---------------------------------------------------------------------------

def gql_post(query: str, api_key: str) -> dict:
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Basic {api_key}",
        "AuthorizationSource": "API",
    }
    payload = json.dumps({"query": query}).encode("utf-8")
    req = urllib.request.Request(HEALTHIE_API_BASE_URL, data=payload,
                                 headers=headers, method="POST")
    try:
        with urllib.request.urlopen(req) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8")
        print(f"HTTP {e.code}: {body}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Request failed: {e}", file=sys.stderr)
        sys.exit(1)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Build Healthie form from Cognito mapping YAML"
    )
    parser.add_argument("--output", "-o", help="Write mutations to file")
    parser.add_argument("--execute", action="store_true",
                        help="Execute against Healthie API")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show mutations without executing")
    parser.add_argument("--yaml", default=YAML_PATH, help="Path to YAML")
    parser.add_argument("--batch-size", type=int, default=20,
                        help="Modules per API call (default: 20)")
    args = parser.parse_args()

    if not os.path.exists(args.yaml):
        print(f"YAML not found: {args.yaml}", file=sys.stderr)
        sys.exit(1)

    data = load_yaml(args.yaml)
    form_name = data.get("form_name", "MetaFlow Enrollment Assessment")
    specs = build_all_specs(data)
    conditional = data.get("conditional_visibility", {})
    with_options = sum(1 for s in specs if s.get("options"))

    print(f"Form: {form_name}")
    print(f"Modules: {len(specs)} ({with_options} with options)")
    print(f"Conditional sections: {len(conditional)}")
    print()

    p1 = phase1_mutation(form_name)

    if args.dry_run:
        print("=== PHASE 1 ===")
        print(p1)
        print()
        print(f"=== PHASE 2: Create {len(specs)} modules ===")
        print(phase2_mutation("<FORM_ID>", specs[:3]))
        if len(specs) > 3:
            print(f"... ({len(specs) - 3} more)")
        return

    if args.execute:
        api_key = os.environ.get("HEALTHIE_API_KEY")
        if not api_key:
            print("HEALTHIE_API_KEY not found in .env.local", file=sys.stderr)
            sys.exit(1)

        # Phase 1: Create form
        print("Phase 1: Creating form container...", file=sys.stderr)
        result = gql_post(p1, api_key)

        if result.get("errors"):
            print(f"Phase 1 errors: {json.dumps(result['errors'], indent=2)}",
                  file=sys.stderr)
            sys.exit(1)

        form = result.get("data", {}).get("createCustomModuleForm", {})
        form_id = form.get("customModuleForm", {}).get("id")
        if not form_id:
            print(f"Form creation failed: {form.get('messages', [])}",
                  file=sys.stderr)
            sys.exit(1)
        print(f"  Form created! ID: {form_id}", file=sys.stderr)
        print()

        # Phase 2: Create modules
        batch_size = args.batch_size
        total = len(specs)
        created_total = 0

        for start in range(0, total, batch_size):
            end = min(start + batch_size, total)
            batch = specs[start:end]
            p2 = phase2_mutation(form_id, batch, start_index=start)

            print(f"Phase 2: Modules {start + 1}-{end}/{total}...",
                  file=sys.stderr, end=" ")
            result = gql_post(p2, api_key)

            if result.get("errors"):
                print(f"ERROR", file=sys.stderr)
                print(json.dumps(result["errors"], indent=2), file=sys.stderr)
                sys.exit(1)

            batch_created = 0
            for i, local_i in enumerate(range(start, end)):
                alias = f"m{local_i}"
                mod = result.get("data", {}).get(alias, {}).get("customModule", {})
                if mod and mod.get("id"):
                    batch_created += 1
                msgs = result.get("data", {}).get(alias, {}).get("messages", [])
                if msgs:
                    print(f"\n  [{batch[i].get('label', '?')}]: {msgs}",
                          file=sys.stderr)

            created_total += batch_created
            print(f"created {batch_created}/{len(batch)}", file=sys.stderr)

            if end < total:
                time.sleep(0.5)

        print()
        print(f"Done! Form ID: {form_id}, {created_total}/{total} modules created.",
              file=sys.stderr)

    elif args.output:
        with open(args.output, "w") as f:
            f.write("# Phase 1\n\n")
            f.write(p1 + "\n\n")
            f.write("# Phase 2\n\n")
            f.write(phase2_mutation("<FORM_ID>", specs) + "\n")
        print(f"Mutations written to {args.output}")

    else:
        print(p1)
        print()
        print(phase2_mutation("<FORM_ID>", specs))


if __name__ == "__main__":
    main()