frontend/src/context/FormContext.jsx

import { createContext, useContext, useState, useCallback } from 'react';
import { getTemplateData } from '../data/templateData';

const FormContext = createContext();

function initializeFormState() {
  const data = getTemplateData();
  const TEMPLATE_SECTIONS = data?.sections || [];
  const HEADER_FIELDS = data?.headerFields || [];

  const state = {
    patientInfo: {},
    sections: {},
    customSections: [],
    enabledSections: {},
  };

  // Initialize header fields
  HEADER_FIELDS.forEach(f => {
    state.patientInfo[f.id] = '';
  });

  // Initialize all sections
  TEMPLATE_SECTIONS.forEach(section => {
    state.enabledSections[section.id] = true;
    state.sections[section.id] = {};

    section.subsections.forEach(sub => {
      sub.fields.forEach(field => {
        if (field.type === 'fill-blank') {
          state.sections[section.id][field.id] = {
            selected: [],
            customOptions: [],
          };
        } else if (field.type === 'checkbox-group') {
          state.sections[section.id][field.id] = {
            checked: [],
            customItems: [],
          };
        } else if (field.type === 'exercise-list') {
          const exercises = {};
          field.exercises.forEach(ex => {
            exercises[ex.id] = { enabled: false, sets: '', reps: '', minutes: '', seconds: '', trials: '', steps: '' };
          });
          state.sections[section.id][field.id] = { exercises, customExercises: [] };
        } else if (field.type === 'text' || field.type === 'number') {
          state.sections[section.id][field.id] = '';
        }
        // static type needs no state
      });
    });
  });

  return state;
}

export function FormProvider({ children }) {
  const [formState, setFormState] = useState(initializeFormState);

  // ── Patient info ──
  const updatePatientInfo = useCallback((fieldId, value) => {
    setFormState(prev => ({
      ...prev,
      patientInfo: { ...prev.patientInfo, [fieldId]: value },
    }));
  }, []);

  // ── Toggle section ──
  const toggleSection = useCallback((sectionId) => {
    setFormState(prev => ({
      ...prev,
      enabledSections: {
        ...prev.enabledSections,
        [sectionId]: !prev.enabledSections[sectionId],
      },
    }));
  }, []);

  // ── Fill-blank selection ──
  const toggleOption = useCallback((sectionId, fieldId, option, multiSelect) => {
    setFormState(prev => {
      const current = prev.sections[sectionId]?.[fieldId]?.selected || [];
      let updated;
      if (multiSelect) {
        updated = current.includes(option)
          ? current.filter(o => o !== option)
          : [...current, option];
      } else {
        updated = current.includes(option) ? [] : [option];
      }
      return {
        ...prev,
        sections: {
          ...prev.sections,
          [sectionId]: {
            ...prev.sections[sectionId],
            [fieldId]: {
              ...prev.sections[sectionId][fieldId],
              selected: updated,
            },
          },
        },
      };
    });
  }, []);

  // ── Add custom option to a fill-blank ──
  const addCustomOption = useCallback((sectionId, fieldId, option) => {
    if (!option.trim()) return;
    setFormState(prev => {
      const field = prev.sections[sectionId]?.[fieldId] || { selected: [], customOptions: [] };
      if (field.customOptions.includes(option)) return prev;
      return {
        ...prev,
        sections: {
          ...prev.sections,
          [sectionId]: {
            ...prev.sections[sectionId],
            [fieldId]: {
              ...field,
              customOptions: [...(field.customOptions || []), option],
              selected: [...(field.selected || []), option],
            },
          },
        },
      };
    });
  }, []);

  // ── Checkbox toggle ──
  const toggleCheckbox = useCallback((sectionId, fieldId, item) => {
    setFormState(prev => {
      const current = prev.sections[sectionId]?.[fieldId]?.checked || [];
      const updated = current.includes(item)
        ? current.filter(i => i !== item)
        : [...current, item];
      return {
        ...prev,
        sections: {
          ...prev.sections,
          [sectionId]: {
            ...prev.sections[sectionId],
            [fieldId]: {
              ...prev.sections[sectionId][fieldId],
              checked: updated,
            },
          },
        },
      };
    });
  }, []);

  // ── Add custom checkbox item ──
  const addCustomCheckboxItem = useCallback((sectionId, fieldId, item) => {
    if (!item.trim()) return;
    setFormState(prev => {
      const field = prev.sections[sectionId]?.[fieldId] || { checked: [], customItems: [] };
      if (field.customItems?.includes(item)) return prev;
      return {
        ...prev,
        sections: {
          ...prev.sections,
          [sectionId]: {
            ...prev.sections[sectionId],
            [fieldId]: {
              ...field,
              customItems: [...(field.customItems || []), item],
              checked: [...(field.checked || []), item],
            },
          },
        },
      };
    });
  }, []);

  // ── Exercise toggle & update ──
  const toggleExercise = useCallback((sectionId, fieldId, exerciseId) => {
    setFormState(prev => {
      const exState = prev.sections[sectionId]?.[fieldId]?.exercises?.[exerciseId];
      if (!exState) return prev;
      return {
        ...prev,
        sections: {
          ...prev.sections,
          [sectionId]: {
            ...prev.sections[sectionId],
            [fieldId]: {
              ...prev.sections[sectionId][fieldId],
              exercises: {
                ...prev.sections[sectionId][fieldId].exercises,
                [exerciseId]: { ...exState, enabled: !exState.enabled },
              },
            },
          },
        },
      };
    });
  }, []);

  const updateExerciseValue = useCallback((sectionId, fieldId, exerciseId, key, value) => {
    setFormState(prev => {
      const exState = prev.sections[sectionId]?.[fieldId]?.exercises?.[exerciseId];
      if (!exState) return prev;
      return {
        ...prev,
        sections: {
          ...prev.sections,
          [sectionId]: {
            ...prev.sections[sectionId],
            [fieldId]: {
              ...prev.sections[sectionId][fieldId],
              exercises: {
                ...prev.sections[sectionId][fieldId].exercises,
                [exerciseId]: { ...exState, [key]: value },
              },
            },
          },
        },
      };
    });
  }, []);

  // ── Text / number update ──
  const updateTextField = useCallback((sectionId, fieldId, value) => {
    setFormState(prev => ({
      ...prev,
      sections: {
        ...prev.sections,
        [sectionId]: {
          ...prev.sections[sectionId],
          [fieldId]: value,
        },
      },
    }));
  }, []);

  // ── Custom sections ──
  const addCustomSection = useCallback((title) => {
    const id = `custom_${Date.now()}`;
    setFormState(prev => ({
      ...prev,
      customSections: [
        ...prev.customSections,
        {
          id,
          title,
          fields: [],
        },
      ],
      enabledSections: { ...prev.enabledSections, [id]: true },
    }));
    return id;
  }, []);

  const addFieldToCustomSection = useCallback((sectionId, field) => {
    setFormState(prev => ({
      ...prev,
      customSections: prev.customSections.map(s =>
        s.id === sectionId
          ? { ...s, fields: [...s.fields, field] }
          : s
      ),
    }));
  }, []);

  const removeCustomSection = useCallback((sectionId) => {
    setFormState(prev => ({
      ...prev,
      customSections: prev.customSections.filter(s => s.id !== sectionId),
    }));
  }, []);

  // ── Build payload for API ──
  const buildPayload = useCallback(() => {
    const data = getTemplateData();
    const TEMPLATE_SECTIONS = data?.sections || [];

    const payload = {
      patientInfo: formState.patientInfo,
      sections: [],
      customSections: [],
    };

    TEMPLATE_SECTIONS.forEach(section => {
      if (!formState.enabledSections[section.id]) return;

      const sectionData = {
        sectionId: section.id,
        sectionTitle: section.title,
        enabled: true,
        fields: [],
      };

      section.subsections.forEach(sub => {
        sub.fields.forEach(field => {
          const val = formState.sections[section.id]?.[field.id];
          if (!val) return;

          if (field.type === 'fill-blank') {
            if (val.selected?.length > 0) {
              sectionData.fields.push({
                fieldId: field.id,
                fieldLabel: field.label,
                selectedOptions: val.selected,
              });
            }
          } else if (field.type === 'checkbox-group') {
            if (val.checked?.length > 0) {
              sectionData.fields.push({
                fieldId: field.id,
                fieldLabel: field.label,
                selectedOptions: val.checked,
              });
            }
          } else if (field.type === 'exercise-list') {
            const enabledExercises = [];
            Object.entries(val.exercises || {}).forEach(([exId, exState]) => {
              if (exState.enabled) {
                const ex = field.exercises.find(e => e.id === exId);
                if (ex) {
                  let desc = ex.name;
                  if (exState.sets && exState.reps) desc += ` \u2014 ${exState.sets} \u00d7 ${exState.reps} reps`;
                  else if (exState.minutes) desc += ` \u2014 ${exState.minutes} min`;
                  else if (exState.seconds) desc += ` \u2014 ${exState.seconds} sec`;
                  else if (exState.trials) desc += ` \u2014 ${exState.trials} trials`;
                  else if (exState.steps) desc += ` \u2014 ${exState.steps} steps`;
                  enabledExercises.push(desc);
                }
              }
            });
            if (enabledExercises.length > 0) {
              sectionData.fields.push({
                fieldId: field.id,
                fieldLabel: field.label,
                selectedOptions: enabledExercises,
              });
            }
          } else if ((field.type === 'text' || field.type === 'number') && val) {
            sectionData.fields.push({
              fieldId: field.id,
              fieldLabel: field.label,
              selectedOptions: [],
              customText: String(val),
            });
          }
        });
      });

      if (sectionData.fields.length > 0) {
        payload.sections.push(sectionData);
      }
    });

    return payload;
  }, [formState]);

  // ── Count filled fields per section ──
  const getFilledCount = useCallback((sectionId) => {
    const sectionState = formState.sections[sectionId];
    if (!sectionState) return 0;
    let count = 0;
    Object.values(sectionState).forEach(val => {
      if (typeof val === 'string' && val) count++;
      else if (val?.selected?.length > 0) count++;
      else if (val?.checked?.length > 0) count++;
      else if (val?.exercises) {
        const hasEnabled = Object.values(val.exercises).some(e => e.enabled);
        if (hasEnabled) count++;
      }
    });
    return count;
  }, [formState]);

  const value = {
    formState,
    updatePatientInfo,
    toggleSection,
    toggleOption,
    addCustomOption,
    toggleCheckbox,
    addCustomCheckboxItem,
    toggleExercise,
    updateExerciseValue,
    updateTextField,
    addCustomSection,
    addFieldToCustomSection,
    removeCustomSection,
    buildPayload,
    getFilledCount,
  };

  return (
    <FormContext.Provider value={value}>
      {children}
    </FormContext.Provider>
  );
}

export function useForm() {
  const ctx = useContext(FormContext);
  if (!ctx) throw new Error('useForm must be used within FormProvider');
  return ctx;
}