#!/usr/bin/env python3
"""End-to-end tests for CMS ACCESS mock server (access-fhir-apis sandbox).

Tests the full patient lifecycle against http://l.nso.me:3001:
  eligibility → alignment → reporting → quarterly timing → unalignment → post-unalign rejection
Plus edge cases: wrong ICD-10, invalid track, missing LOINCs.

Usage: python3 test_access_mock.py [--base URL]
"""
import json, re, sys, time, urllib.error, urllib.parse, urllib.request, datetime

BASE = "http://l.nso.me:3001"
PATIENT_ID = "pat-test-99"
MBI = "TEST0000001"

TRACK_ECKM = {
    "coding": [{
        "system": "https://dsacms.github.io/cmmi-access-model/CodeSystem/ACCESSTrackCS",
        "code": "eCKM"
    }]
}

ECKM_BASELINE_LOINCS = [
    "85354-9",  # Blood pressure panel
    "13457-7",  # LDL cholesterol
    "2093-3",   # BMI
    "2571-8",   # eGFR creatinine
    "39156-5",  # Triglycerides
    "56086-2",  # Waist circumference
    "1558-6",   # Fasting glucose
    "4548-4",   # HbA1c
]

PATIENT = {
    "resourceType": "Patient",
    "id": PATIENT_ID,
    "identifier": [{"system": "http://terminology.hl7.org/NamingSystem/cmsMBI", "value": MBI}],
    "name": [{"family": "Testerson", "given": ["Alice"]}],
    "gender": "female",
    "birthDate": "1965-03-15"
}


# ── Helpers ────────────────────────────────────────────────────────────────

def post(path, data=None):
    """POST JSON to FHIR endpoint."""
    body = json.dumps(data).encode() if data else b'{}'
    req = urllib.request.Request(
        f"{BASE}{path}", body,
        {"Content-Type": "application/json"}, method="POST"
    )
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())


def get(path):
    """GET JSON from FHIR endpoint."""
    req = urllib.request.Request(f"{BASE}{path}")
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())


def get_param(resp, name):
    """Find a parameter by name in a Parameters resource."""
    for p in resp.get("parameter", []):
        if p["name"] == name:
            return p
    return None


def extract_sub_id(text):
    """Extract submissionID from operation outcome text."""
    m = re.search(r"submissionID=(sub-\d+-\d+)", text)
    return m.group(1) if m else None


def poll(submission_id, timeout=20):
    """Poll $submission-status until result available or timeout.

    Returns (result_code, submission_id).
    """
    for _ in range(timeout * 2):
        time.sleep(0.5)
        resp = get(f"/fhir/Patient/$submission-status?submissionID={urllib.parse.quote(submission_id)}")
        if resp.get("resourceType") == "OperationOutcome":
            continue  # Not found yet, keep polling
        result_p = get_param(resp, "result")
        if result_p and "resource" in result_p:
            code_p = get_param(result_p["resource"], "result")
            if code_p:
                return (code_p["valueCodeableConcept"]["coding"][0]["code"], submission_id)
    return ("timeout", submission_id)


def issue_summary(resp):
    """Human-readable summary of OperationOutcome issues."""
    if "issue" in resp:
        parts = []
        for i in resp["issue"]:
            diag = i.get("diagnostics", "")
            parts.append(f"{i['severity']}/{i['code']}: {diag}")
        return "; ".join(parts)
    return resp.get("resourceType", "?")


def make_obs(loinc_code, value=1.0, unit=""):
    """Create a minimal Observation resource for the given LOINC code."""
    return {
        "code": {"coding": [{"system": "http://loinc.org", "code": loinc_code}]},
        "valueQuantity": {"value": value, "unit": unit}
    }


def make_bp_obs(systolic=128, diastolic=78):
    """Create BP panel observation (LOINC 85354-9 with components)."""
    return {
        "code": {"coding": [{"system": "http://loinc.org", "code": "85354-9"}]},
        "component": [
            {"code": {"coding": [{"code": "8480-6"}]},
             "valueQuantity": {"value": systolic, "unit": "mmHg"}},
            {"code": {"coding": [{"code": "8462-4"}]},
             "valueQuantity": {"value": diastolic, "unit": "mmHg"}}
        ]
    }


def baseline_payload(patient_id, date, loincs=None):
    """Build $report-data payload for a baseline report."""
    if loincs is None:
        loincs = ECKM_BASELINE_LOINCS

    observations = []
    for c in loincs:
        if c == "85354-9":
            observations.append(make_bp_obs())
        elif c == "13457-7":
            observations.append(make_obs(c, 92, "mg/dL"))
        elif c == "2093-3":
            observations.append(make_obs(c, 28.1, "kg/m2"))
        elif c == "2571-8":
            observations.append(make_obs(c, 0.6, "{mg/dL}/mg"))
        elif c == "39156-5":
            observations.append(make_obs(c, 138, "mg/dL"))
        elif c == "56086-2":
            observations.append(make_obs(c, 34, "in"))
        elif c == "1558-6":
            observations.append(make_obs(c, 95, "mg/dL"))
        elif c == "4548-4":
            observations.append(make_obs(c, 6.8, "%"))

    return {
        "resourceType": "Parameters",
        "parameter": [
            {"name": "measureReport", "resource": {
                "resourceType": "MeasureReport",
                "measure": "http://example.org/fhir/Measure/access-eckm-baseline",
                "subject": {"reference": f"Patient/{patient_id}"},
                "date": date
            }},
        ] + [
            {"name": "resource", "resource": {**o, "resourceType": "Observation"}}
            for o in observations
        ]
    }


# ── Test Runner ────────────────────────────────────────────────────────────

class Results:
    def __init__(self):
        self.passed = 0
        self.failed = 0
        self.errors = []

    def run(self, name, fn):
        print(f"\n[{self.passed + self.failed + 1}] {name}")
        try:
            detail = fn()
            print(f"  PASS — {detail}")
            self.passed += 1
        except Exception as e:
            msg = str(e) if str(e) else "AssertionError"
            print(f"  FAIL — {msg}")
            self.failed += 1
            self.errors.append((name, msg))


# ── Tests ──────────────────────────────────────────────────────────────────

def run_tests():
    TODAY = datetime.date.today().isoformat()
    r = Results()

    # Reset mock state for clean test run
    print("=" * 60)
    print("CMS ACCESS Mock Server — End-to-End Tests")
    print(f"Target: {BASE}")
    print("=" * 60)
    print("\nResetting mock state... ", end="")
    reset_resp = post("/fhir/$mock-reset", {})
    diag = issue_summary(reset_resp)
    print(f"{diag}")

    # TEST 1: Metadata / CapabilityStatement
    def t1():
        meta = get("/fhir/metadata")
        ops = [op["name"] for rest in meta["rest"] for op in rest.get("operation", [])]
        required = {"check-eligibility", "align", "unalign", "report-data", "submission-status"}
        assert required.issubset(set(ops)), f"Missing: {required - set(ops)}"
        return f"All 5 operations present: {', '.join(sorted(ops))}"

    r.run("Metadata / CapabilityStatement", t1)

    # TEST 2: Eligibility — positive (eCKM track, I10 hypertension)
    def t2():
        resp = post("/fhir/Patient/$check-eligibility", {
            "resourceType": "Parameters", "parameter": [
                {"name": "participantID", "valueIdentifier": {"value": "ACCESS1234"}},
                {"name": "patient", "resource": PATIENT},
                {"name": "track", "valueCodeableConcept": TRACK_ECKM},
                {"name": "condition", "resource": {
                    "resourceType": "Condition",
                    "code": {"coding": [{"system": "http://hl7.org/fhir/sid/icd-10-cm", "code": "I10"}]}
                }}
            ]
        })
        sub_id = extract_sub_id(issue_summary(resp))
        assert sub_id, f"No submission ID in: {issue_summary(resp)}"
        code, _ = poll(sub_id)
        assert code == "eligible", f"Expected eligible, got {code}"
        return f"Eligible (submission {sub_id})"

    r.run("Eligibility — I10 hypertension → eligible", t2)

    # TEST 3: Eligibility — wrong ICD-10 for track
    def t3():
        resp = post("/fhir/Patient/$check-eligibility", {
            "resourceType": "Parameters", "parameter": [
                {"name": "participantID", "valueIdentifier": {"value": "ACCESS1234"}},
                {"name": "patient", "resource": PATIENT},
                {"name": "track", "valueCodeableConcept": TRACK_ECKM},
                {"name": "condition", "resource": {
                    "resourceType": "Condition",
                    "code": {"coding": [{"system": "http://hl7.org/fhir/sid/icd-10-cm", "code": "J06.9"}]}
                }}
            ]
        })
        sub_id = extract_sub_id(issue_summary(resp))
        code, _ = poll(sub_id)
        assert code == "not-eligible-diagnoses", f"Expected not-eligible-diagnoses, got {code}"
        return f"Not eligible — ICD-10 mismatch (submission {sub_id})"

    r.run("Eligibility — J06.9 URI → not-eligible-diagnoses", t3)

    # TEST 4: Align patient to eCKM track
    def t4():
        resp = post("/fhir/Patient/$align", {
            "resourceType": "Parameters", "parameter": [
                {"name": "participantID", "valueIdentifier": {"value": "ACCESS1234"}},
                {"name": "patient", "resource": PATIENT},
                {"name": "track", "valueCodeableConcept": TRACK_ECKM},
                {"name": "condition", "resource": {
                    "resourceType": "Condition",
                    "code": {"coding": [{"system": "http://hl7.org/fhir/sid/icd-10-cm", "code": "I10"}]}
                }},
                {"name": "isProviderReferral", "valueBoolean": True}
            ]
        })
        sub_id = extract_sub_id(issue_summary(resp))
        code, _ = poll(sub_id)
        assert code == "aligned", f"Expected aligned, got {code}"
        return f"Aligned (submission {sub_id})"

    r.run("Align → aligned", t4)

    # TEST 5: Baseline report — MISSING HbA1c → rejected-missing-measures
    def t5():
        missing_hba1c = [c for c in ECKM_BASELINE_LOINCS if c != "4548-4"]
        resp = post("/fhir/$report-data", baseline_payload(PATIENT_ID, TODAY, missing_hba1c))
        sub_id = extract_sub_id(issue_summary(resp))
        code, _ = poll(sub_id)
        assert code == "rejected-missing-measures", f"Expected rejected-missing-measures, got {code}"
        return f"Rejected — missing LOINC 4548-4 (submission {sub_id})"

    r.run("Baseline minus HbA1c → rejected-missing-measures", t5)

    # TEST 6: Baseline report — ALL 8 LOINCs → accepted
    def t6():
        resp = post("/fhir/$report-data", baseline_payload(PATIENT_ID, TODAY))
        sub_id = extract_sub_id(issue_summary(resp))
        code, _ = poll(sub_id)
        assert code == "accepted", f"Expected accepted, got {code}"
        return f"Accepted (all 8 LOINCs, submission {sub_id})"

    r.run("Baseline all 8 LOINCs → accepted", t6)

    # TEST 7: Quarterly report — too early (< 70 days since baseline)
    def t7():
        resp = post("/fhir/$report-data", {
            "resourceType": "Parameters", "parameter": [
                {"name": "measureReport", "resource": {
                    "resourceType": "MeasureReport",
                    "measure": "http://example.org/fhir/Measure/access-eckm-quarterly",
                    "subject": {"reference": f"Patient/{PATIENT_ID}"},
                    "date": TODAY
                }},
                {"name": "resource", "resource": {**make_bp_obs(), "resourceType": "Observation"}},
            ]
        })
        sub_id = extract_sub_id(issue_summary(resp))
        code, _ = poll(sub_id)
        assert code == "rejected-outside-window", f"Expected rejected-outside-window, got {code}"
        return f"Rejected — quarterly too early (submission {sub_id})"

    r.run("Quarterly too soon → rejected-outside-window", t7)

    # TEST 8: Unalign patient
    def t8():
        resp = post("/fhir/Patient/$unalign", {
            "resourceType": "Parameters", "parameter": [
                {"name": "patient", "resource": PATIENT},
                {"name": "track", "valueCodeableConcept": TRACK_ECKM},
                {"name": "reason", "valueCoding": {
                    "system": "https://dsacms.github.io/cmmi-access-model/CodeSystem/ACCESSUnalignmentReasonCS",
                    "code": "geographic-relocated"
                }}
            ]
        })
        sub_id = extract_sub_id(issue_summary(resp))
        code, _ = poll(sub_id)
        assert code == "unaligned", f"Expected unaligned, got {code}"
        return f"Unaligned (submission {sub_id})"

    r.run("Unalign → unaligned", t8)

    # TEST 9: Report after unalignment → rejected-not-aligned
    def t9():
        resp = post("/fhir/$report-data", baseline_payload(PATIENT_ID, TODAY))
        sub_id = extract_sub_id(issue_summary(resp))
        code, _ = poll(sub_id)
        assert code == "rejected-not-aligned", f"Expected rejected-not-aligned, got {code}"
        return f"Rejected — patient not aligned (submission {sub_id})"

    r.run("Report after unalign → rejected-not-aligned", t9)

    # TEST 10: Invalid track code — handled gracefully
    def t10():
        resp = post("/fhir/Patient/$check-eligibility", {
            "resourceType": "Parameters", "parameter": [
                {"name": "participantID", "valueIdentifier": {"value": "ACCESS1234"}},
                {"name": "patient", "resource": PATIENT},
                {"name": "track", "valueCodeableConcept": {"coding": [{"code": "INVALID_TRACK"}]}},
                {"name": "condition", "resource": {
                    "resourceType": "Condition",
                    "code": {"coding": [{"system": "http://hl7.org/fhir/sid/icd-10-cm", "code": "I10"}]}
                }}
            ]
        })
        text = issue_summary(resp)
        assert "sub-" in text, f"Expected submission accepted (async validation), got: {text}"
        return f"Accepted for async validation (submission {extract_sub_id(text)})"

    r.run("Invalid track → handled gracefully", t10)

    # ── Summary ────────────────────────────────────────────────────────
    print("\n" + "=" * 60)
    total = r.passed + r.failed
    print(f"RESULTS: {r.passed}/{total} passed, {r.failed} failed")
    if r.errors:
        print("\nFailures:")
        for name, msg in r.errors:
            print(f"  - {name}: {msg}")
    else:
        print("ALL TESTS PASSED")
    print("=" * 60)
    return len(r.errors)


if __name__ == "__main__":
    # Allow overriding base URL via CLI arg or env var
    if "--base" in sys.argv:
        idx = sys.argv.index("--base") + 1
        if idx < len(sys.argv):
            BASE = sys.argv[idx]
            print(f"Using custom base URL: {BASE}")

    import os
    env_base = os.environ.get("ACCESS_MOCK_URL")
    if env_base:
        BASE = env_base
        print(f"Using ACCESS_MOCK_URL: {BASE}")

    sys.exit(run_tests())