logic.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field, asdict
from datetime import date, datetime, timedelta
from typing import Optional


EXPECTED_KEYS = [
    "drug_name",
    "dose",
    "frequency_text",
    "quantity",
    "refill_or_expiry_date",
]


@dataclass
class Medication:
    drug_name: Optional[str] = None
    dose: Optional[str] = None
    frequency_text: Optional[str] = None
    quantity: Optional[float] = None
    refill_or_expiry_date: Optional[str] = None

    schedule_times: list[str] = field(default_factory=list)
    doses_per_day: Optional[float] = None
    run_out_date: Optional[str] = None
    refill_warning: Optional[str] = None

    def to_row(self) -> list:
        return [
            self.drug_name or "—",
            self.dose or "—",
            ", ".join(self.schedule_times) if self.schedule_times else (self.frequency_text or "—"),
            self.quantity if self.quantity is not None else "—",
            self.run_out_date or "—",
            self.refill_warning or "",
        ]

    def to_dict(self) -> dict:
        return asdict(self)


def safe_parse(raw: str) -> Optional[dict]:
    if not raw or not raw.strip():
        return None

    text = raw.strip()

    fence = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
    if fence:
        text = fence.group(1)
    else:
        brace = re.search(r"\{.*\}", text, re.DOTALL)
        if brace:
            text = brace.group(0)

    try:
        data = json.loads(text)
    except (json.JSONDecodeError, ValueError):
        return None

    if not isinstance(data, dict):
        return None

    out: dict = {}
    for k in EXPECTED_KEYS:
        v = data.get(k)
        if isinstance(v, str) and v.strip().lower() in ("", "null", "none", "n/a", "unknown"):
            v = None
        out[k] = v

    out["quantity"] = _coerce_quantity(out.get("quantity"))
    return out


def _coerce_quantity(value) -> Optional[float]:
    if value is None:
        return None
    if isinstance(value, (int, float)):
        return float(value)
    m = re.search(r"\d+(?:\.\d+)?", str(value))
    return float(m.group(0)) if m else None


_DEFAULT_TIMES = {
    1: ["08:00"],
    2: ["08:00", "20:00"],
    3: ["08:00", "14:00", "20:00"],
    4: ["08:00", "12:00", "16:00", "20:00"],
}


def build_schedule(data: dict) -> tuple[list[str], Optional[float]]:
    text = (data.get("frequency_text") or "").lower().strip()
    if not text:
        return [], None

    if "as needed" in text or "prn" in text or "if needed" in text:
        return [], None

    m = re.search(r"every\s+(\d+)\s*(?:hours|hrs|h)\b", text)
    if m:
        interval = int(m.group(1))
        if interval > 0:
            doses = max(1, round(24 / interval))
            times = _even_spread(interval)
            return times, float(doses)

    word_counts = {
        "once": 1, "one time": 1, "1 time": 1,
        "twice": 2, "two times": 2, "2 times": 2,
        "three times": 3, "thrice": 3, "3 times": 3,
        "four times": 4, "4 times": 4,
    }
    for phrase, n in word_counts.items():
        if phrase in text:
            return list(_DEFAULT_TIMES.get(n, [])), float(n)

    m = re.search(r"(\d+)\s*(?:x|times)\s*(?:a|per)?\s*(?:day|daily)", text)
    if m:
        n = int(m.group(1))
        return list(_DEFAULT_TIMES.get(n, [])), float(n)

    if "bedtime" in text or "at night" in text or "nightly" in text:
        return ["21:00"], 1.0
    if "morning" in text and "night" not in text and "evening" not in text:
        return ["08:00"], 1.0

    if re.search(r"\b(once a day|once daily|every day|daily|qd)\b", text):
        return ["08:00"], 1.0

    return [], None


def _even_spread(interval_hours: int) -> list[str]:
    times = []
    h = 8
    while h < 24:
        times.append(f"{h:02d}:00")
        h += interval_hours
    return times or ["08:00"]


def check_refill(
    data: dict,
    doses_per_day: Optional[float],
    today: Optional[date] = None,
    warn_within_days: int = 7,
) -> tuple[Optional[str], Optional[str]]:
    today = today or date.today()
    run_out_iso: Optional[str] = None
    warnings: list[str] = []

    qty = data.get("quantity")
    if qty is not None and doses_per_day and doses_per_day > 0:
        days_left = int(qty // doses_per_day)
        run_out = today + timedelta(days=days_left)
        run_out_iso = run_out.isoformat()
        if days_left <= 0:
            warnings.append("⚠️ Out of medication — refill now")
        elif days_left <= warn_within_days:
            warnings.append(f"⚠️ Refill soon — about {days_left} day(s) left")

    label_date = _parse_date(data.get("refill_or_expiry_date"))
    if label_date:
        delta = (label_date - today).days
        if delta < 0:
            warnings.append("⚠️ Label refill/expiry date has passed")
        elif delta <= warn_within_days:
            warnings.append(f"⚠️ Label date in {delta} day(s)")

    return run_out_iso, (" · ".join(warnings) if warnings else None)


def _parse_date(value) -> Optional[date]:
    if not value:
        return None
    s = str(value).strip()
    for fmt in ("%Y-%m-%d", "%m/%d/%Y", "%m/%d/%y", "%d/%m/%Y", "%b %d %Y", "%B %d, %Y"):
        try:
            return datetime.strptime(s, fmt).date()
        except ValueError:
            continue
    return None


def build_daily_agenda(meds: list["Medication"]):
    by_time: dict[str, list[str]] = {}
    as_needed: list[str] = []

    for m in meds:
        label = m.drug_name or "Medication"
        if m.dose:
            label = f"{label} ({m.dose})"
        if m.schedule_times:
            for t in m.schedule_times:
                by_time.setdefault(t, []).append(label)
        else:
            instr = m.frequency_text or "no schedule"
            as_needed.append(f"{label} — {instr}")

    timed = sorted(by_time.items(), key=lambda kv: kv[0])
    return timed, as_needed


def format_time_12h(t: str) -> str:
    try:
        h, m = (int(x) for x in t.split(":"))
        suffix = "AM" if h < 12 else "PM"
        h12 = h % 12 or 12
        return f"{h12}:{m:02d} {suffix}"
    except (ValueError, AttributeError):
        return t


def make_medication(data: dict, today: Optional[date] = None) -> Medication:
    times, doses = build_schedule(data)
    run_out, warning = check_refill(data, doses, today=today)
    return Medication(
        drug_name=data.get("drug_name"),
        dose=data.get("dose"),
        frequency_text=data.get("frequency_text"),
        quantity=data.get("quantity"),
        refill_or_expiry_date=data.get("refill_or_expiry_date"),
        schedule_times=times,
        doses_per_day=doses,
        run_out_date=run_out,
        refill_warning=warning,
    )