#!/usr/bin/env python3
"""CMS ACCESS API — Clinical Demo Script

Idempotent, one-command walkthrough of the full patient access lifecycle:
  1. Seed synthetic patient record
  2. Check CMS eligibility ($check-eligibility)
  3. Grant consent (simulated)
  4. Align beneficiary ($align)
  5. Show results summary

Designed for clinical and non-technical audiences. All data is synthetic.

Usage:
    python3 scripts/demo_access_lifecycle.py [--base URL] [--track eCKM|CKM|MSK]

Requirements:
    - Django project configured (DJANGO_SETTINGS_MODULE set)
    - Database exists with migrations applied
    - Mock server reachable at --base URL (default: http://l.nso.me:3001)
"""
import os, sys, uuid, datetime, json, time, requests

# Add project root to sys.path so Django can find access_core (run from repo root)
_project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _project_root not in sys.path:
    sys.path.insert(0, _project_root)

# ── CLI Arguments ───────────────────────────────────────────────────────────

MOCK_BASE = "http://l.nso.me:3001"
TRACK = "eCKM"

if "--base" in sys.argv:
    idx = sys.argv.index("--base") + 1
    if idx < len(sys.argv):
        MOCK_BASE = sys.argv[idx]

if "--track" in sys.argv:
    idx = sys.argv.index("--track") + 1
    if idx < len(sys.argv):
        # Must match TRACK_CHOICES in access_core/constants.py (no BH track exists)
        VALID_TRACKS = {"eckm": "eCKM", "ckm": "CKM", "msk": "MSK"}
        raw = sys.argv[idx].lower()
        TRACK = VALID_TRACKS.get(raw)
        if not TRACK:
            print(f"\nError: Invalid track '{sys.argv[idx]}'. Valid options: eCKM, CKM, MSK")
            print("  Note: 'eCKM' is the default (chronic kidney disease / hypertension)\n")
            sys.exit(1)

# ── Configure Django ────────────────────────────────────────────────────────

os.environ["CMS_API_BASE_URL"] = MOCK_BASE
os.environ["CMS_FHIR_VERSION_PATH"] = ""
os.environ.pop("CMS_API_CLIENT_ID", None)
os.environ.pop("CMS_API_CLIENT_SECRET", None)
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "access_core.settings.local")

import django
django.setup()

from django.conf import settings as _settings
_settings.CMS_API_BASE_URL = MOCK_BASE
_settings.CMS_FHIR_VERSION_PATH = ""

from django.core.management import call_command
from access_core.intake.models import IntakeRecord, IntakeIdentifier
from access_core.monitoring.models import (
    EligibilityCheck, EligibilityRequest, EligibilityRequestCondition,
    Alignment, CarePeriod, ClinicalTrack, Unalignment
)
from access_core.identity.models import Beneficiary, ExternalSystemIdentifier
from access_core.cms_integration.models import (
    CMSApiRequest, CMSApiResponse, CMSApiSubmissionStatus
)


def wait_for_cms_results(timeout=30.0):
    """Drive poll_cms_submissions until no open submissions remain.

    The services are submit-only; this emulates cms_polling_worker inline so the
    demo can show final model state (the mock resolves after ~ASYNC_DELAY_MS).
    """
    from io import StringIO

    deadline = time.monotonic() + timeout
    while True:
        call_command("poll_cms_submissions", force=True, stdout=StringIO())
        open_rows = CMSApiSubmissionStatus.objects.filter(
            status__in=[
                CMSApiSubmissionStatus.Status.PENDING,
                CMSApiSubmissionStatus.Status.IN_PROGRESS,
            ]
        )
        if not open_rows.exists():
            return
        if time.monotonic() >= deadline:
            raise TimeoutError(
                f"{open_rows.count()} submission(s) still processing after {timeout:.0f}s"
            )
        time.sleep(0.6)

# ── Display Helpers ──────────────────────────────────────────────────────────

BOLD = "\033[1m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
RED = "\033[91m"
CYAN = "\033[96m"
DIM = "\033[2m"
RESET = "\033[0m"

def banner(title):
    w = 65
    print(f"\n{BOLD}{'=' * w}{RESET}")
    print(f"{BOLD}  {title}{RESET}")
    print(f"{BOLD}{'=' * w}{RESET}\n")

def step(num, title):
    print(f"\n{CYAN}[Step {num}] {title}{RESET}")
    sys.stdout.flush()

def info(msg):
    print(f"  {DIM}> {msg}{RESET}")

def success(msg):
    print(f"  {GREEN}v {msg}{RESET}")

def warn(msg):
    print(f"  {YELLOW}! {msg}{RESET}")

def error(msg):
    print(f"  {RED}x {msg}{RESET}")

# ── Data Helpers ─────────────────────────────────────────────────────────────

DEMO_PREFIX = "demo"

TRACK_ICD10 = {
    "eCKM": ("I10", "Essential hypertension"),
    "CKM":   ("N18", "Chronic kidney disease"),
    "MSK":   ("M47", "Spondylosis"),
    "BH":    ("F32", "Major depressive disorder"),
}

# Deterministic synthetic patient data (varies by index to avoid control group)
SYNTHETIC_PATIENTS = [
    {"first": "Maria",   "last": "Gonzalez",  "dob": datetime.date(1958, 6, 22), "sex": "female"},
    {"first": "Carlos",  "last": "Rodriguez", "dob": datetime.date(1962, 3, 14), "sex": "male"},
    {"first": "Elena",   "last": "Chen",      "dob": datetime.date(1970, 11, 8),  "sex": "female"},
    {"first": "James",   "last": "Williams",  "dob": datetime.date(1955, 1, 30),  "sex": "male"},
    {"first": "Aisha",   "last": "Patel",     "dob": datetime.date(1968, 9, 5),   "sex": "female"},
    {"first": "Robert",  "last": "Johnson",   "dob": datetime.date(1960, 4, 17),  "sex": "male"},
    {"first": "Sarah",   "last": "Kim",       "dob": datetime.date(1953, 8, 29),  "sex": "female"},
    {"first": "David",   "last": "Martinez",  "dob": datetime.date(1975, 12, 3),  "sex": "male"},
    {"first": "Linda",   "last": "Thompson",  "dob": datetime.date(1964, 7, 21),  "sex": "female"},
    {"first": "Michael", "last": "Brown",     "dob": datetime.date(1950, 2, 14),  "sex": "male"},
]

MAX_RETRIES = len(SYNTHETIC_PATIENTS)

def reset_mock_server(url):
    """Reset mock server state before demo."""
    try:
        requests.post(f"{url}/$mock-reset", timeout=10)
        return True
    except Exception:
        return False


def create_demo_patient(index=0):
    """Create a synthetic patient record for the demo."""
    int_id = uuid.uuid4()
    short = int_id.hex[:10]
    pdata = SYNTHETIC_PATIENTS[index % len(SYNTHETIC_PATIENTS)]

    patient = IntakeRecord.objects.create(
        external_form_id=f"{DEMO_PREFIX}-{short}",
        first_name=pdata["first"],
        last_name=pdata["last"],
        dob=pdata["dob"],
        sex=pdata["sex"],
        is_synthetic=True,
        status="pending_cms_eligibility",
    )

    IntakeIdentifier.objects.create(
        intake_record=patient,
        identifier_type="email",
        identifier_value=f"{DEMO_PREFIX}.{short}@demo.local",
        is_primary=True,
    )
    IntakeIdentifier.objects.create(
        intake_record=patient,
        identifier_type="phone",
        identifier_value="+15559876543",
    )
    IntakeIdentifier.objects.create(
        intake_record=patient,
        identifier_type="address",
        identifier_value=f"{42 + index} Wellness Blvd, Suite 100",
        identifier_metadata={
            "line1": f"{42 + index} Wellness Blvd, Ste 100",
            "city": "Riverside",
            "state": "CA",
            "zip": "92501"
        },
    )

    return patient


def create_demo_eligibility_request(patient):
    """Create the reviewer's EligibilityRequest (+ ICD-10 condition) for a demo patient."""
    icd10_code, icd10_desc = TRACK_ICD10.get(TRACK, ("I10", "Essential hypertension"))
    track = ClinicalTrack.objects.get_or_create(
        track_name=TRACK, defaults={"is_active": True}
    )[0]
    request = EligibilityRequest.objects.create(
        intake_record=patient, clinical_track=track
    )
    EligibilityRequestCondition.objects.create(
        request=request, icd10_code=icd10_code, icd10_description=icd10_desc
    )
    return request


def cleanup_demo_data(intake_record):
    """Clean up all demo data."""
    intake_id = str(intake_record.id)

    CMSApiResponse.objects.all().delete()
    CMSApiRequest.objects.filter(beneficiary__intake_record=intake_record).exclude(
        beneficiary=None
    ).delete()
    CMSApiRequest.objects.exclude(beneficiary=None).filter(
        beneficiary__intake_record_id=intake_id
    ).delete()

    bens = list(Beneficiary.objects.filter(intake_record_id=intake_id))
    ben_ids = [b.id for b in bens] if bens else []

    if ben_ids:
        ExternalSystemIdentifier.objects.filter(beneficiary_id__in=ben_ids).delete()
        try:
            BeneficiaryIdentifier = django.apps.apps.get_model('identity', 'BeneficiaryIdentifier')
            BeneficiaryIdentifier.objects.filter(beneficiary_id__in=ben_ids).delete()
        except Exception:
            pass
        Unalignment.objects.filter(beneficiary_id__in=ben_ids).delete()
        CarePeriod.objects.filter(beneficiary_id__in=ben_ids).delete()

    Alignment.objects.filter(intake_record=intake_record).update(beneficiary=None)
    if ben_ids:
        Alignment.objects.filter(beneficiary_id__in=ben_ids).update(beneficiary=None)
        Beneficiary.objects.filter(id__in=ben_ids).delete()

    EligibilityRequestCondition.objects.filter(
        request__intake_record=intake_record
    ).delete()
    EligibilityCheck.objects.filter(intake_record=intake_record).delete()
    EligibilityRequest.objects.filter(intake_record=intake_record).delete()
    Alignment.objects.filter(intake_record=intake_record).delete()
    IntakeIdentifier.objects.filter(intake_record=intake_record).delete()
    intake_record.delete()


# ── Demo Flow ───────────────────────────────────────────────────────────────

def run_demo():
    icd10_code, icd10_desc = TRACK_ICD10.get(TRACK, ("I10", "Essential hypertension"))

    banner("CMS ACCESS API - Clinical Demo")
    info(f"Synthetic data only | Track: {TRACK} ({icd10_code}: {icd10_desc})")
    info(f"CMS endpoint: {MOCK_BASE}")

    # Reset mock server for clean state
    if reset_mock_server(MOCK_BASE):
        success("Mock server reset - clean state ready")
    else:
        warn("Could not reset mock server (continuing anyway)")

    patient = None
    ben = None
    cp = None
    check = None

    try:
        # ── Step 1 + 2: Create Patient & Check Eligibility (with retry) ─
        step(1, "Register Synthetic Patient")

        eligible = False
        for attempt in range(MAX_RETRIES):
            if patient:
                cleanup_demo_data(patient)
                info(f"Retrying with different synthetic data ({attempt + 1}/{MAX_RETRIES})...")

            # Reset mock server between attempts to get fresh randomization
            if attempt > 0 and reset_mock_server(MOCK_BASE):
                pass  # Silently reset

            patient = create_demo_patient(index=attempt)
            pdata = SYNTHETIC_PATIENTS[attempt % len(SYNTHETIC_PATIENTS)]
            info(f"Name: {pdata['first']} {pdata['last']} | DOB: {pdata['dob'].strftime('%m/%d/%Y')}")
            info(f"Record marked as SYNTHETIC (no real PHI)")

            intake_id = str(patient.id)

            step(2, "Check CMS Eligibility ($check-eligibility)")
            info("Sending synthetic patient to CMS ACCESS API...")

            # Suppress management command output for cleaner demo output
            from io import StringIO as _StringIO
            import sys as _sys
            old_stdout = _sys.stdout
            _sys.stdout = _StringIO()

            t0 = time.time()
            try:
                create_demo_eligibility_request(patient)  # reviewer-determined track + ICD-10
                call_command("check_cms_eligibility", intake_id=intake_id, verbosity=0)
                wait_for_cms_results()  # worker applies the outcome (submit-only services)
            finally:
                _sys.stdout = old_stdout  # Restore stdout after command

            elapsed = time.time() - t0

            patient.refresh_from_db()
            check = EligibilityCheck.objects.filter(intake_record=patient).first()

            if patient.status == "pending_consent" and check:
                eligible = True
                success(f"Patient is ELIGIBLE for {TRACK} track ({elapsed:.1f}s)")
                info(f"CMS response: {check.check_status}")
                break
            elif check and check.check_status == "control_group":
                info("CMS assigned this patient to the control group (randomized)")
                warn("Trying next synthetic patient...")
                continue
            else:
                resp_preview = str(check.cms_response)[:200] if check and check.cms_response else "no response data"
                warn(f"Status={patient.status} | CMS response preview: {resp_preview}")
                info("This may indicate a mock server issue — trying next synthetic patient...")
                continue
        else:
            error("Could not find eligible patient after trying all synthetic patients")
            info("The mock server may have assigned all test patients to the control group.")
            info("Try running again, or check that the CMS endpoint is reachable:")
            info(f"  {MOCK_BASE}")

            # Show diagnostic info from last attempt
            if check and check.cms_response:
                import json as _json
                resp = check.cms_response
                params = resp.get('parameter', []) if isinstance(resp, dict) else []
                result_param = next((p for p in params if p.get('name') == 'result'), None)
                if result_param and 'resource' in result_param:
                    nested = result_param['resource'].get('parameter', [])
                    result_code = next(
                        (np.get('valueCodeableConcept', {}).get('coding', [{}])[0].get('code')
                         for np in nested if np.get('name') == 'result'),
                        "unknown"
                    )
                    info(f"Last CMS response code: {result_code}")
                    if result_code.startswith("not-eligible"):
                        reason = result_code.replace("not-eligible-", "").title().replace("_", " ")
                        warn(f"Reason: Patient data does not satisfy '{reason}' criteria for track {TRACK}")

            return False

        # ── Step 3: Grant Consent (Simulated) ──────────────────────────
        step(3, "Grant Patient Consent (simulated)")
        info("In production, patient signs consent form via portal.")
        info(f"Status transition: pending_consent -> pending_alignment")

        old_status = patient.status
        patient.status = "pending_alignment"
        patient.save(update_fields=["status"])
        success(f"Consent granted ({old_status} -> pending_alignment)")

        # ── Step 4: Align Beneficiary ($align) ─────────────────────────
        step(4, "Align with CMS ACCESS Program ($align)")
        info("Submitting alignment request to CMS...")

        t0 = time.time()
        _sys.stdout = _StringIO()  # Suppress management command output
        try:
            call_command("run_cms_alignment", track=TRACK, intake_id=intake_id, verbosity=0)
            wait_for_cms_results()  # worker applies the outcome (submit-only services)
        finally:
            _sys.stdout = old_stdout
        elapsed = time.time() - t0

        patient.refresh_from_db()

        if patient.status == "active_access_beneficiary":
            alignment = Alignment.objects.filter(intake_record=patient).first()
            ben = alignment.beneficiary if alignment and alignment.beneficiary else None
            cp = CarePeriod.objects.filter(beneficiary=ben).first() if ben else None

            success(f"Alignment successful ({elapsed:.1f}s)")
            info(f"CMS Beneficiary ID: {ben.id}")
            info(f"Care Period status: {cp.operational_status if cp else 'N/A'}")
            info(f"Track: {TRACK} | ICD-10: {icd10_code} ({icd10_desc})")
        else:
            error(f"Alignment did not complete. Status: {patient.status}")
            return False

        # ── Step 5: Summary ─────────────────────────────────────────────
        step(5, "Demo Complete - Results Summary")

        api_requests = CMSApiRequest.objects.filter(beneficiary=ben) if ben else CMSApiRequest.objects.none()

        banner("Summary")
        info(f"Patient:       {pdata['first']} {pdata['last']} (synthetic)")
        info(f"Track:         {TRACK} ({icd10_code})")
        info(f"Eligibility:   {check.check_status if check else 'N/A'}")
        info(f"Beneficiary:   {ben.id if ben else 'N/A'}")
        info(f"Care Period:   {'Active' if cp and cp.operational_status == 'active' else 'Inactive'}")
        info(f"CMS API Calls: {api_requests.count()}")

    except Exception as exc:
        error(f"Demo failed: {exc}")
        import traceback
        print(f"\n{DIM}{traceback.format_exc()}{RESET}")
        return False

    finally:
        if patient:
            cleanup_demo_data(patient)
            info("Cleanup complete - demo data removed from database")

    success("Full ACCESS lifecycle demonstrated successfully!")
    return True


# ── Entry Point ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    ok = run_demo()
    sys.exit(0 if ok else 1)