# epic_vitals_and_therapy.py
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Any, Dict, Optional
import requests

from token_service import get_or_create_epic_token

from therapy_engine import TherapyEngine, ConflictFinder
from fhir_to_engine import flattened_vitals_to_engine_readings

# --- keep your critical mapping + flatten function from your script ---
CRITICAL_VITAL_LOINC = {
    "8867-4": "heart_rate",
    "9279-1": "respiratory_rate",
    "8310-5": "body_temperature",
    "2708-6": "oxygen_saturation",
    "8302-2": "body_height",
    "29463-7": "body_weight",
    "8478-0": "mean_arterial_pressure",
    "8480-6": "systolic_blood_pressure",
    "8462-4": "diastolic_blood_pressure",
}

def _flatten_vital_observation(obs: Dict[str, Any]) -> list[Dict[str, Any]]:
    subject = (obs.get("subject") or {}).get("reference")
    if not subject:
        return []

    base = {
        "observation_id": obs.get("id"),
        "patient_ref": subject,
        "encounter_ref": (obs.get("encounter") or {}).get("reference"),
        "effective_dt": obs.get("effectiveDateTime"),
        "status": obs.get("status", "unknown"),
    }

    rows: list[Dict[str, Any]] = []

    def _row(coding: Dict[str, Any], value: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        code = coding.get("code")
        if not code or code not in CRITICAL_VITAL_LOINC:
            return None
        val = value.get("value")
        if val is None:
            return None
        r = dict(base)
        r.update(
            {
                "loinc_code": code,
                "value": val,
                "unit": value.get("unit"),
            }
        )
        return r

    if "component" in obs:
        for comp in obs.get("component") or []:
            coding = (((comp.get("code") or {}).get("coding")) or [{}])[0]
            value_q = comp.get("valueQuantity") or {}
            row = _row(coding, value_q)
            if row:
                rows.append(row)
    else:
        coding = (((obs.get("code") or {}).get("coding")) or [{}])[0]
        value_q = obs.get("valueQuantity") or {}
        row = _row(coding, value_q)
        if row:
            rows.append(row)

    return rows

def normalize_resource_url(fhir_base: str, resource: str) -> str:
    base = fhir_base.strip().rstrip("/")
    lowered = base.lower()
    r = resource.lower()
    if lowered.endswith(f"/{r}"):
        return base
    if lowered.endswith("/fhir/r4"):
        return base + f"/{resource}"
    return base + f"/{resource}"

def to_patient_reference(patient: str) -> str:
    p = patient.strip()
    if not p:
        return p
    if "/" in p:
        return p
    return f"Patient/{p}"

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--fhir-base", default=os.getenv("FHIR_BASE_URL", "https://fhir.epic.com/interconnect-fhir-oauth/api/FHIR/R4"))
    ap.add_argument("--app-config-id", type=int, default=int(os.getenv("EPIC_APP_CONFIG_ID", "3")))
    ap.add_argument("--patient", required=True, help="Patient id or reference")
    ap.add_argument("--json", default="observation_vitals.json", help="Path to Observation JSON payload.")
    ap.add_argument("--timeout", type=int, default=60)
    args = ap.parse_args()

    token = get_or_create_epic_token(args.app_config_id)
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/fhir+json",
        "Content-Type": "application/fhir+json",
    }

    observation_url = normalize_resource_url(args.fhir_base, "Observation")

    with open(args.json, "r", encoding="utf-8") as f:
        payload = json.load(f)

    # Ensure payload has correct subject
    payload.setdefault("subject", {"reference": to_patient_reference(args.patient)})

    # Flatten and compute therapy from the vitals inside the payload
    vitals_rows = _flatten_vital_observation(payload)
    vitals_readings = flattened_vitals_to_engine_readings(vitals_rows)

    # If you also have labs available, pass them here in the same shape:
    labs_readings: Dict[str, Dict[str, Any]] = {}  # e.g. {"K": {"value": 3.2, "id": "lab123"}}

    engine = TherapyEngine(conflict_finder=ConflictFinder())
    therapy = engine.compute_doses(labs=labs_readings, vitals=vitals_readings, lite=False)

    print("Computed therapy recommendations (engine output):")
    print(json.dumps(therapy, indent=2))

    # Gate: only POST if there is at least one critical vital row
    if not vitals_rows:
        print("No critical vitals found in payload; skipping Observation POST.")
        return 0

    try:
        resp = requests.post(observation_url, headers=headers, json=payload, timeout=args.timeout)
    except requests.RequestException as e:
        print(f"ERROR: Observation POST failed: {e}", file=sys.stderr)
        return 1

    print(f"Observation POST URL: {observation_url}")
    print("Observation POST HTTP:", resp.status_code)

    try:
        print(json.dumps(resp.json(), indent=2))
    except Exception:
        print(resp.text)

    return 0 if resp.status_code in (200, 201) else 1

if __name__ == "__main__":
    raise SystemExit(main())