src/store.py

# src/store.py
from __future__ import annotations
"""Case persistence helpers (Schema v2) with legacy API compatibility.

Public API preserved from V1:
- new_case_id(), case_path(), read_case(), save_case(), create_case(), update_case(), list_cases()

Additions in V2:
- upgrade_to_v2(obj), set_status(case_id, status), touch(case_id)
- seed_cases(reset: bool = False)
"""

import json, time, string
from pathlib import Path
from typing import List, Dict, Optional, Any

from .paths import cases_dir
from .config import SCHEMA_VERSION

# --------------------------- utils ---------------------------

def _now_iso() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%S")

def _rand(n=4) -> str:
    import secrets
    alphabet = string.ascii_uppercase + string.digits
    return ''.join(secrets.choice(alphabet) for _ in range(n))

def _deep_merge(a: Dict[str, Any], b: Dict[str, Any]) -> Dict[str, Any]:
    out = dict(a)
    for k, v in b.items():
        if isinstance(v, dict) and isinstance(out.get(k), dict):
            out[k] = _deep_merge(out[k], v)  # type: ignore
        else:
            out[k] = v
    return out

# --------------------------- paths ---------------------------

def case_path(case_id: str) -> Path:
    d = cases_dir()
    d.mkdir(parents=True, exist_ok=True)
    return d / f"{case_id}.json"

def new_case_id() -> str:
    return f"EC{_rand(6)}"

# ------------------------ migrations -------------------------

def _migrate_soup_to_soap(obj: Dict[str, Any]) -> Dict[str, Any]:
    """Ensure SOAP shape; migrate legacy 'soup' key to 'soap_draft'."""
    if isinstance(obj, dict) and "soup" in obj and "soap_draft" not in obj:
        soup = obj.pop("soup", {})
        if isinstance(soup, dict):
            obj["soap_draft"] = {
                "subjective": soup.get("subjective", ""),
                "objective": soup.get("objective", ""),
                "assessment": soup.get("assessment", ""),
                "plan": soup.get("plan", ""),
            }
    return obj

def _ensure_v2_fields(obj: Dict[str, Any]) -> Dict[str, Any]:
    """Add Schema v2 fields if missing."""
    obj.setdefault("schema_version", SCHEMA_VERSION)
    obj.setdefault("created_at", _now_iso())
    obj.setdefault("updated_at", _now_iso())
    obj.setdefault("status", "draft")
    obj.setdefault("review", {"state": None, "notes": ""})
    obj.setdefault("billing", {"minutes": 5, "spoke": False, "cpt_code": None, "attested": False})
    obj.setdefault("explainability", {"manually_modified": False})
    obj.setdefault("patient", {"name": "", "age": None, "sex": ""})
    obj.setdefault("consult", {"question": "", "history": "", "medications": "", "labs": ""})
    obj.setdefault("soap_draft", {"subjective": "", "objective": "", "assessment": "", "plan": ""})
    return obj

def upgrade_to_v2(obj: Dict[str, Any]) -> Dict[str, Any]:
    """Idempotent upgrader to Schema v2."""
    if obj is None:
        return obj
    obj = _migrate_soup_to_soap(obj)
    obj = _ensure_v2_fields(obj)
    obj["schema_version"] = SCHEMA_VERSION
    return obj

# ------------------------ core ops ---------------------------

def read_case(case_id: str) -> Optional[Dict[str, Any]]:
    p = case_path(case_id)
    if not p.exists():
        return None
    try:
        data = json.loads(p.read_text(encoding='utf-8'))
    except Exception:
        return None
    return upgrade_to_v2(data)

def save_case(obj: Dict[str, Any]) -> None:
    obj = upgrade_to_v2(obj or {})
    obj["updated_at"] = _now_iso()
    p = case_path(obj["case_id"])
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(obj, ensure_ascii=False, indent=2), encoding='utf-8')

def create_case(initial: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    case_id = (initial or {}).get("case_id") or new_case_id()
    obj: Dict[str, Any] = {"case_id": case_id}
    obj = _deep_merge(obj, initial or {})
    obj = upgrade_to_v2(obj)
    save_case(obj)
    return obj

def update_case(case_id: str, patch: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    cur = read_case(case_id)
    if cur is None:
        return None
    updated = _deep_merge(cur, patch)
    save_case(updated)
    return updated

def list_cases(status: Optional[str] = None) -> List[Dict[str, Any]]:
    items: List[Dict[str, Any]] = []
    d = cases_dir()
    d.mkdir(parents=True, exist_ok=True)
    for p in d.glob("*.json"):
        try:
            obj = json.loads(p.read_text(encoding='utf-8'))
            obj = upgrade_to_v2(obj)
            if (status is None) or (obj.get("status") == status):
                items.append({
                    "case_id": obj.get("case_id"),
                    "status": obj.get("status"),
                    "created_at": obj.get("created_at"),
                    "updated_at": obj.get("updated_at"),
                    "patient_name": (obj.get("patient") or {}).get("name"),
                    "specialty": (obj.get("consult") or {}).get("specialty"),
                    "question": (obj.get("consult") or {}).get("question"),
                })
        except Exception:
            continue
    items.sort(key=lambda x: x.get("updated_at") or "", reverse=True)
    return items

def set_status(case_id: str, status: str) -> Optional[Dict[str, Any]]:
    """Set status with simple tri-state guard: draft -> submitted -> completed."""
    allowed = ["draft", "submitted", "completed"]
    if status not in allowed:
        raise ValueError(f"Invalid status: {status}")
    cur = read_case(case_id)
    if cur is None:
        return None
    current = cur.get("status", "draft")
    order = {s:i for i, s in enumerate(allowed)}
    if order[status] < order.get(current, 0):
        return cur  # disallow backward moves
    cur["status"] = status
    save_case(cur)
    return cur

def touch(case_id: str) -> None:
    obj = read_case(case_id)
    if obj:
        obj["updated_at"] = _now_iso()
        save_case(obj)

# ------------------------ seeding ----------------------------

def _seed_payloads() -> Dict[str, Dict[str, Any]]:
    """Construct seed cases as per spec: 3 Draft, 1 In Review (submitted)."""
    seeds: Dict[str, Dict[str, Any]] = {}

    # Cardiology — LDL/statin intolerance (In Review + mock SOAP)
    cid = "CARD-LDL"
    seeds[cid] = {
        "case_id": cid,
        "status": "submitted",  # UI can render as "In Review"
        "review": {"state": "in_review", "notes": "Specialist reviewing; mock SOAP present."},
        "patient": {"name": "Pat Demo", "age": 58, "sex": "F"},
        "consult": {
            "specialty": "Cardiology",
            "question": "LDL management with statin intolerance?",
            "history": "Primary hyperlipidemia; myalgias with atorvastatin and simvastatin.",
            "medications": "Ezetimibe 10 mg daily",
            "labs": "LDL 162 mg/dL; HDL 48 mg/dL; TG 190 mg/dL",
            "consent_obtained": True,
        },
        "soap_draft": {
            "subjective": "Reports muscle aches with prior statins; adherent to ezetimibe.",
            "objective": "LDL 162 mg/dL; ASCVD risk elevated; BP 132/78; weight stable.",
            "assessment": "Statin intolerance vs nocebo; persistent LDL elevation despite ezetimibe.",
            "plan": "Consider PCSK9 inhibitor or inclisiran; reinforce lifestyle; check CK/Vit D; baseline LFTs; shared decision-making.",
        },
        "billing": {"minutes": 12, "spoke": True, "cpt_code": None, "attested": False},
        "explainability": {"manually_modified": False},
    }

    # Endocrinology — T2DM + CKD (Draft)
    cid = "ENDO-T2DM-CKD"
    seeds[cid] = {
        "case_id": cid,
        "status": "draft",
        "patient": {"name": "Alex Demo", "age": 67, "sex": "M"},
        "consult": {
            "specialty": "Endocrinology",
            "question": "Optimization of T2DM regimen in CKD stage 3?",
            "history": "A1c drifting up; eGFR 48; metformin on board; no SGLT2 yet.",
            "medications": "Metformin 1000 mg BID",
            "labs": "A1c 8.2%; eGFR 48 mL/min/1.73m²; K 4.6",
        },
        "soap_draft": {"subjective": "", "objective": "", "assessment": "", "plan": ""},
        "billing": {"minutes": 5, "spoke": False, "cpt_code": None, "attested": False},
        "explainability": {"manually_modified": False},
    }

    # Dermatology — chronic rash (Draft)
    cid = "DERM-RASH-CHRONIC"
    seeds[cid] = {
        "case_id": cid,
        "status": "draft",
        "patient": {"name": "Sam Demo", "age": 44, "sex": "M"},
        "consult": {
            "specialty": "Dermatology",
            "question": "Chronic pruritic rash on forearms; failed topical steroid.",
            "history": "6 months intermittent erythematous plaques; worse with sun exposure.",
            "medications": "Topical triamcinolone 0.1%",
            "labs": "CBC normal; no eosinophilia",
        },
        "soap_draft": {"subjective": "", "objective": "", "assessment": "", "plan": ""},
        "billing": {"minutes": 5, "spoke": False, "cpt_code": None, "attested": False},
        "explainability": {"manually_modified": False},
    }

    # Infectious Disease — diabetic foot infection (Draft)
    cid = "ID-DFI"
    seeds[cid] = {
        "case_id": cid,
        "status": "draft",
        "patient": {"name": "Jamie Demo", "age": 72, "sex": "F"},
        "consult": {
            "specialty": "Infectious Disease",
            "question": "Antibiotic strategy for diabetic foot infection with poor perfusion?",
            "history": "Ulcer at plantar surface; prior MRSA colonization; neuropathy.",
            "medications": "No abx started",
            "labs": "ESR 44; CRP 2.1; WBC 9.8",
        },
        "soap_draft": {"subjective": "", "objective": "", "assessment": "", "plan": ""},
        "billing": {"minutes": 5, "spoke": False, "cpt_code": None, "attested": False},
        "explainability": {"manually_modified": False},
    }

    return seeds

def seed_cases(reset: bool = False) -> List[str]:
    """Create four seed cases; returns created/updated case_ids."""
    d = cases_dir()
    d.mkdir(parents=True, exist_ok=True)
    if reset:
        for p in d.glob("*.json"):
            p.unlink()
    ids: List[str] = []
    for cid, payload in _seed_payloads().items():
        save_case(payload)
        ids.append(cid)
    return ids