import { create } from 'zustand';
import type { AuthUser } from '../../services/auth';
import type { APIDiagnosis, APIHospitalVisit, APIPrescription, APITest } from '../../services/medical';
import { fetchDiagnoses, fetchHospitalVisits, fetchPrescriptions, fetchTests } from '../../services/medical';
import { fetchUsersByIds, fetchPractitionerVisits } from '../services/practitionerApi';

export interface PatientData {
    user: AuthUser;
    diagnoses: APIDiagnosis[];
    visits: APIHospitalVisit[];
    prescriptions: APIPrescription[];
    tests: APITest[];
}

export interface ConsultationDraft {
    patientId: string;
    chiefComplaint: string;
    diagnosis: { icdCode: string; diseaseName: string; severity: string; notes: string } | null;
    prescription: {
        medicationId: string;
        medicationName: string;
        dosage: string;
        frequency: string;
        startDate: string;
        endDate: string;
    } | null;
    notes: string;
}

interface PractitionerState {
    patients: AuthUser[];
    patientsLoading: boolean;

    viewingPatientId: string | null;

    currentPatientId: string | null;
    currentPatientData: PatientData | null;
    currentPatientLoading: boolean;

    appointments: APIHospitalVisit[];
    appointmentsLoading: boolean;
    appointmentsCacheKey: string | null;

    consultationOpen: boolean;
    consultationDraft: ConsultationDraft | null;

    loadPatients: (patientIds: string[]) => Promise<void>;
    loadAppointments: (startDate?: string, endDate?: string) => Promise<void>;
    loadPatientDetail: (patientId: string) => Promise<void>;
    setCurrentPatient: (id: string | null) => void;
    enterPatientView: (patientId: string) => void;
    exitPatientView: () => void;
    openConsultation: (patientId: string) => void;
    closeConsultation: () => void;
    updateConsultationDraft: (update: Partial<ConsultationDraft>) => void;
    clearCurrentPatient: () => void;
}

export const usePractitionerStore = create<PractitionerState>((set, get) => ({
    patients: [],
    patientsLoading: false,

    viewingPatientId: null,

    currentPatientId: null,
    currentPatientData: null,
    currentPatientLoading: false,

    appointments: [],
    appointmentsLoading: false,
    appointmentsCacheKey: null,

    consultationOpen: false,
    consultationDraft: null,

    loadAppointments: async (startDate?: string, endDate?: string) => {
        const key = `${startDate || ''}_${endDate || ''}`;
        if (get().appointmentsCacheKey === key && get().appointments.length > 0) return;

        set({ appointmentsLoading: true });
        try {
            const visits = await fetchPractitionerVisits(startDate, endDate);
            set({ appointments: visits, appointmentsLoading: false, appointmentsCacheKey: key });
        } catch {
            set({ appointmentsLoading: false });
        }
    },

    loadPatients: async (patientIds: string[]) => {
        if (!patientIds.length) {
            set({ patients: [], patientsLoading: false });
            return;
        }
        set({ patientsLoading: true });
        try {
            const users = await fetchUsersByIds(patientIds);
            set({ patients: users, patientsLoading: false });
        } catch {
            set({ patientsLoading: false });
        }
    },

    loadPatientDetail: async (patientId: string) => {
        if (get().currentPatientId === patientId && get().currentPatientData) return;

        set({ currentPatientId: patientId, currentPatientLoading: true });
        try {
            const patients = get().patients;
            let user = patients.find((p) => p._id === patientId);

            if (!user) {
                const fetched = await fetchUsersByIds([patientId]);
                user = fetched[0];
            }

            if (!user) {
                set({ currentPatientLoading: false });
                return;
            }

            const [diagnoses, visits, prescriptions, tests] = await Promise.all([
                fetchDiagnoses(patientId),
                fetchHospitalVisits(patientId),
                fetchPrescriptions(patientId),
                fetchTests(patientId),
            ]);

            set({
                currentPatientData: { user, diagnoses, visits, prescriptions, tests },
                currentPatientLoading: false,
            });
        } catch {
            set({ currentPatientLoading: false });
        }
    },

    setCurrentPatient: (id) => set({ currentPatientId: id }),

    enterPatientView: (patientId: string) =>
        set({ viewingPatientId: patientId, currentPatientId: patientId }),

    exitPatientView: () =>
        set({ viewingPatientId: null }),

    clearCurrentPatient: () =>
        set({ currentPatientId: null, currentPatientData: null }),

    openConsultation: (patientId: string) =>
        set({
            consultationOpen: true,
            consultationDraft: {
                patientId,
                chiefComplaint: '',
                diagnosis: null,
                prescription: null,
                notes: '',
            },
        }),

    closeConsultation: () =>
        set({ consultationOpen: false, consultationDraft: null }),

    updateConsultationDraft: (update) =>
        set((state) => ({
            consultationDraft: state.consultationDraft
                ? { ...state.consultationDraft, ...update }
                : null,
        })),
}));
