Gishen-Mob/lib/stores/prescriptionStore.ts

import AsyncStorage from "@react-native-async-storage/async-storage";
import { create } from "zustand";
import { createJSONStorage, persist } from "zustand/middleware";
import {
  seedDoseLogs,
  seedPrescriptions,
  seedReminders,
} from "@/lib/data/seed";

export type RxStatus =
  | "submitted"
  | "under_review"
  | "approved"
  | "queried"
  | "rejected"
  | "filled";

export type LineStatus = "ready" | "review" | "unavailable";

export type Prescription = {
  id: string;
  patientName: string;
  memberId?: string | null;
  medicineName: string;
  medicines: { name: string; status: LineStatus }[];
  imageUris: string[];
  daysSupply: number;
  lastFilledAt: string | null;
  status: RxStatus;
  notes?: string;
  /** Message from pharmacist when status is queried / rejected */
  pharmacistMessage?: string;
  /** Customer reply to a query */
  customerReply?: string;
  createdAt: string;
};

export type Reminder = {
  id: string;
  prescriptionId: string;
  medicineName: string;
  times: string[];
  daysOfWeek: number[];
  graceMinutes: number;
  escalateOnMiss: boolean;
  memberId?: string | null;
  enabled: boolean;
};

export type DoseLog = {
  id: string;
  reminderId: string;
  scheduledAt: string;
  status: "pending" | "taken" | "snoozed" | "missed_escalated";
  confirmedAt?: string;
};

type RxState = {
  prescriptions: Prescription[];
  reminders: Reminder[];
  doseLogs: DoseLog[];
  draftImages: string[];
  setDraftImages: (uris: string[]) => void;
  addPrescription: (
    p: Omit<Prescription, "id" | "createdAt" | "status">
  ) => Prescription;
  updateRxStatus: (id: string, status: RxStatus, pharmacistMessage?: string) => void;
  replyToQuery: (id: string, message: string) => void;
  markRefilled: (id: string) => void;
  removePrescription: (id: string) => void;
  upsertReminder: (r: Omit<Reminder, "id"> & { id?: string }) => void;
  removeReminder: (id: string) => void;
  confirmDose: (logId: string) => void;
  escalateDose: (logId: string) => void;
  schedulePendingDose: (reminderId: string) => DoseLog;
};

export const usePrescriptionStore = create<RxState>()(
  persist(
    (set, get) => ({
      prescriptions: seedPrescriptions,
      reminders: seedReminders,
      doseLogs: seedDoseLogs,
      draftImages: [],
      setDraftImages: (uris) => set({ draftImages: uris }),
      addPrescription: (payload) => {
        const rx: Prescription = {
          ...payload,
          id: `rx_${Date.now()}`,
          createdAt: new Date().toISOString(),
          status: "submitted",
        };
        set({ prescriptions: [rx, ...get().prescriptions] });
        return rx;
      },
      updateRxStatus: (id, status, pharmacistMessage) =>
        set({
          prescriptions: get().prescriptions.map((p) =>
            p.id === id
              ? {
                  ...p,
                  status,
                  ...(pharmacistMessage !== undefined
                    ? { pharmacistMessage }
                    : {}),
                }
              : p
          ),
        }),
      replyToQuery: (id, message) =>
        set({
          prescriptions: get().prescriptions.map((p) =>
            p.id === id
              ? {
                  ...p,
                  customerReply: message,
                  status: "under_review",
                  pharmacistMessage: undefined,
                }
              : p
          ),
        }),
      markRefilled: (id) =>
        set({
          prescriptions: get().prescriptions.map((p) =>
            p.id === id
              ? { ...p, lastFilledAt: new Date().toISOString(), status: "filled" }
              : p
          ),
        }),
      removePrescription: (id) =>
        set({
          prescriptions: get().prescriptions.filter((p) => p.id !== id),
          reminders: get().reminders.filter((r) => r.prescriptionId !== id),
        }),
      upsertReminder: (r) => {
        const existing = get().reminders.find(
          (x) =>
            (r.id != null && x.id === r.id) ||
            x.prescriptionId === r.prescriptionId
        );
        const id = r.id ?? existing?.id ?? `rem_${Date.now()}`;
        const next = { ...existing, ...r, id } as Reminder;
        const others = get().reminders.filter(
          (x) => x.id !== id && x.prescriptionId !== r.prescriptionId
        );
        set({ reminders: [next, ...others] });
      },
      removeReminder: (id) =>
        set({ reminders: get().reminders.filter((r) => r.id !== id) }),
      confirmDose: (logId) =>
        set({
          doseLogs: get().doseLogs.map((l) =>
            l.id === logId
              ? {
                  ...l,
                  status: "taken",
                  confirmedAt: new Date().toISOString(),
                }
              : l
          ),
        }),
      escalateDose: (logId) =>
        set({
          doseLogs: get().doseLogs.map((l) =>
            l.id === logId ? { ...l, status: "missed_escalated" } : l
          ),
        }),
      schedulePendingDose: (reminderId) => {
        const log: DoseLog = {
          id: `dose_${Date.now()}`,
          reminderId,
          scheduledAt: new Date().toISOString(),
          status: "pending",
        };
        set({ doseLogs: [log, ...get().doseLogs] });
        return log;
      },
    }),
    {
      name: "gishen-prescriptions",
      storage: createJSONStorage(() => AsyncStorage),
      version: 2,
      migrate: (persisted: unknown) => {
        const state = (persisted ?? {}) as Partial<RxState>;
        const empty = !state.prescriptions?.length;
        return empty
          ? {
              ...state,
              prescriptions: seedPrescriptions,
              reminders: seedReminders,
              doseLogs: seedDoseLogs,
            }
          : state;
      },
    }
  )
);