TalkEdit/frontend/src/store/editorStore.ts

import { create } from 'zustand';
import { temporal } from 'zundo';
import type { Word, Segment, DeletedRange, CutRange, MuteRange, TranscriptionResult, ProjectFile } from '../types/project';

interface EditorState {
  videoPath: string | null;
  videoUrl: string | null;
  exportedAudioPath: string | null;  // path to modified audio from a previous export
  words: Word[];
  segments: Segment[];
  deletedRanges: DeletedRange[];
  cutRanges: CutRange[];
  muteRanges: MuteRange[];
  language: string;

  currentTime: number;
  duration: number;
  isPlaying: boolean;

  selectedWordIndices: number[];
  hoveredWordIndex: number | null;

  isTranscribing: boolean;
  transcriptionProgress: number;
  transcriptionStatus: string;
  isExporting: boolean;
  exportProgress: number;

  backendUrl: string;
}

interface EditorActions {
  setBackendUrl: (url: string) => void;
  loadVideo: (path: string) => void;
  setExportedAudioPath: (path: string | null) => void;
  saveProject: () => ProjectFile;
  setTranscription: (result: TranscriptionResult) => void;
  setCurrentTime: (time: number) => void;
  setDuration: (duration: number) => void;
  setIsPlaying: (playing: boolean) => void;
  setSelectedWordIndices: (indices: number[]) => void;
  setHoveredWordIndex: (index: number | null) => void;
  deleteSelectedWords: () => void;
  deleteWordRange: (startIndex: number, endIndex: number) => void;
  restoreRange: (rangeId: string) => void;
  addCutRange: (start: number, end: number) => void;
  addMuteRange: (start: number, end: number) => void;
  updateCutRange: (id: string, start: number, end: number) => void;
  updateMuteRange: (id: string, start: number, end: number) => void;
  removeCutRange: (id: string) => void;
  removeMuteRange: (id: string) => void;
  setTranscribing: (active: boolean, progress?: number, status?: string) => void;
  setExporting: (active: boolean, progress?: number) => void;
  getKeepSegments: () => Array<{ start: number; end: number }>;
  getWordAtTime: (time: number) => number;
  loadProject: (projectData: any) => void;
  reset: () => void;
  pauseUndo: () => void;
  resumeUndo: () => void;
}

const initialState: EditorState = {
  videoPath: null,
  videoUrl: null,
  exportedAudioPath: null,
  words: [],
  segments: [],
  deletedRanges: [],
  cutRanges: [],
  muteRanges: [],
  language: '',
  currentTime: 0,
  duration: 0,
  isPlaying: false,
  selectedWordIndices: [],
  hoveredWordIndex: null,
  isTranscribing: false,
  transcriptionProgress: 0,
  transcriptionStatus: '',
  isExporting: false,
  exportProgress: 0,
  backendUrl: 'http://127.0.0.1:8000',
};

let nextRangeId = 1;

export const useEditorStore = create<EditorState & EditorActions>()(
  temporal(
    (set, get) => ({
      ...initialState,

      setBackendUrl: (url) => set({ backendUrl: url }),

      setExportedAudioPath: (path) => set({ exportedAudioPath: path }),

      saveProject: (): ProjectFile => {
        const { videoPath, words, segments, deletedRanges, cutRanges, muteRanges, language, exportedAudioPath } = get();
        if (!videoPath) throw new Error('No video loaded');
        const now = new Date().toISOString();
        // Strip globalStartIndex (runtime-only field) before persisting
        const persistSegments = segments.map(({ globalStartIndex: _drop, ...rest }) => rest);
        return {
          version: 1,
          videoPath,
          exportedAudioPath: exportedAudioPath ?? undefined,
          words,
          segments: persistSegments as unknown as Segment[],
          deletedRanges,
          cutRanges,
          muteRanges,
          language,
          createdAt: now,  // will be overwritten if we track original creation time later
          modifiedAt: now,
        };
      },

      loadVideo: (path) => {
        const backend = get().backendUrl;
        const buildMediaUrl = (filePath: string) => {
          const isWav = filePath.toLowerCase().endsWith('.wav');
          return isWav
            ? `${backend}/file?path=${encodeURIComponent(filePath)}&format=mp3`
            : `${backend}/file?path=${encodeURIComponent(filePath)}`;
        };
        const url = buildMediaUrl(path);
        set({
          ...initialState,
          backendUrl: backend,
          videoPath: path,
          videoUrl: url,
        });
      },

      setTranscription: (result) => {
        let globalIdx = 0;
        const annotatedSegments = result.segments.map((seg) => {
          const annotated = { ...seg, globalStartIndex: globalIdx };
          globalIdx += seg.words.length;
          return annotated;
        });
        set({
          words: result.words,
          segments: annotatedSegments,
          language: result.language,
          deletedRanges: [],
          selectedWordIndices: [],
        });
      },

      setCurrentTime: (time) => set({ currentTime: time }),
      setDuration: (duration) => set({ duration }),
      setIsPlaying: (playing) => set({ isPlaying: playing }),
      setSelectedWordIndices: (indices) => set({ selectedWordIndices: indices }),
      setHoveredWordIndex: (index) => set({ hoveredWordIndex: index }),

      deleteSelectedWords: () => {
        const { selectedWordIndices, words, deletedRanges } = get();
        if (selectedWordIndices.length === 0) return;

        const sorted = [...selectedWordIndices].sort((a, b) => a - b);
        const startWord = words[sorted[0]];
        const endWord = words[sorted[sorted.length - 1]];

        const newRange: DeletedRange = {
          id: `dr_${nextRangeId++}`,
          start: startWord.start,
          end: endWord.end,
          wordIndices: sorted,
        };

        set({
          deletedRanges: [...deletedRanges, newRange],
          selectedWordIndices: [],
        });
      },

      deleteWordRange: (startIndex, endIndex) => {
        const { words, deletedRanges } = get();
        const indices = [];
        for (let i = startIndex; i <= endIndex; i++) indices.push(i);

        const newRange: DeletedRange = {
          id: `dr_${nextRangeId++}`,
          start: words[startIndex].start,
          end: words[endIndex].end,
          wordIndices: indices,
        };

        set({ deletedRanges: [...deletedRanges, newRange] });
      },

      restoreRange: (rangeId) => {
        const { deletedRanges } = get();
        set({ deletedRanges: deletedRanges.filter((r) => r.id !== rangeId) });
      },

      addCutRange: (start, end) => {
        const { cutRanges } = get();
        const newRange: CutRange = {
          id: `cut_${nextRangeId++}`,
          start,
          end,
        };
        set({ cutRanges: [...cutRanges, newRange] });
      },

      addMuteRange: (start, end) => {
        const { muteRanges } = get();
        const newRange: MuteRange = {
          id: `mute_${nextRangeId++}`,
          start,
          end,
        };
        set({ muteRanges: [...muteRanges, newRange] });
      },

      updateCutRange: (id, start, end) => {
        const { cutRanges } = get();
        set({
          cutRanges: cutRanges.map((r) =>
            r.id === id ? { ...r, start, end } : r
          ),
        });
      },

      updateMuteRange: (id, start, end) => {
        const { muteRanges } = get();
        set({
          muteRanges: muteRanges.map((r) =>
            r.id === id ? { ...r, start, end } : r
          ),
        });
      },

      removeCutRange: (id) => {
        const { cutRanges } = get();
        set({ cutRanges: cutRanges.filter((r) => r.id !== id) });
      },

      removeMuteRange: (id) => {
        const { muteRanges } = get();
        set({ muteRanges: muteRanges.filter((r) => r.id !== id) });
      },

      setTranscribing: (active, progress, status) =>
        set({
          isTranscribing: active,
          transcriptionProgress: progress ?? (active ? 0 : 100),
          transcriptionStatus: status ?? (active ? 'Transcribing...' : ''),
        }),

      setExporting: (active, progress) =>
        set({
          isExporting: active,
          exportProgress: progress ?? (active ? 0 : 100),
        }),

      getKeepSegments: () => {
        const { words, deletedRanges, cutRanges, duration } = get();
        if (words.length === 0) return [{ start: 0, end: duration }];

        const deletedSet = new Set<number>();
        for (const range of deletedRanges) {
          for (const idx of range.wordIndices) deletedSet.add(idx);
        }

        // Also exclude words that fall within cut ranges
        for (const cutRange of cutRanges) {
          for (let i = 0; i < words.length; i++) {
            const word = words[i];
            if (word.start >= cutRange.start && word.end <= cutRange.end) {
              deletedSet.add(i);
            }
          }
        }

        const segments: Array<{ start: number; end: number }> = [];
        let segStart: number | null = null;

        for (let i = 0; i < words.length; i++) {
          if (!deletedSet.has(i)) {
            if (segStart === null) segStart = words[i].start;
          } else {
            if (segStart !== null) {
              segments.push({ start: segStart, end: words[i - 1].end });
              segStart = null;
            }
          }
        }

        if (segStart !== null) {
          segments.push({ start: segStart, end: words[words.length - 1].end });
        }

        return segments;
      },

      getWordAtTime: (time) => {
        const { words } = get();
        let lo = 0;
        let hi = words.length - 1;
        while (lo <= hi) {
          const mid = (lo + hi) >>> 1;
          if (words[mid].end < time) lo = mid + 1;
          else if (words[mid].start > time) hi = mid - 1;
          else return mid;
        }
        return lo < words.length ? lo : words.length - 1;
      },

      loadProject: (data) => {
        const backend = get().backendUrl;
        const isWav = data.videoPath.toLowerCase().endsWith('.wav');
        const url = isWav
          ? `${backend}/file?path=${encodeURIComponent(data.videoPath)}&format=mp3`
          : `${backend}/file?path=${encodeURIComponent(data.videoPath)}`;

        let globalIdx = 0;
        const annotatedSegments = (data.segments || []).map((seg: Segment) => {
          const annotated = { ...seg, globalStartIndex: globalIdx };
          globalIdx += seg.words.length;
          return annotated;
        });

        set({
          ...initialState,
          backendUrl: backend,
          videoPath: data.videoPath,
          videoUrl: url,
          words: data.words || [],
          segments: annotatedSegments,
          deletedRanges: data.deletedRanges || [],
          cutRanges: data.cutRanges || [],
          muteRanges: data.muteRanges || [],
          language: data.language || '',
          exportedAudioPath: data.exportedAudioPath ?? null,
        });
      },

      reset: () => set(initialState),

      pauseUndo: () => {
        // Access the temporal store through the useEditorStore
        const temporalStore = (useEditorStore as any).temporal;
        if (temporalStore) {
          temporalStore.getState().pause();
        }
      },

      resumeUndo: () => {
        // Access the temporal store through the useEditorStore
        const temporalStore = (useEditorStore as any).temporal;
        if (temporalStore) {
          temporalStore.getState().resume();
        }
      },
    }),
    { limit: 100 },
  ),
);