import { create } from 'zustand';
import { ProjectFile, VisualElement } from '../types/blocks';
import type { CssRule } from '@shared/types';

type EditorMode = 'blocks' | 'visual' | 'code' | 'files' | 'preview' | 'css';
type ViewMode = 'design' | 'preview' | 'split';

interface EditorStore {
  activeFileId: string | null;
  activeFileContent: string;
  editorMode: EditorMode;
  viewMode: ViewMode;
  selectedElementId: string | null;
  selectedBlockId: string | null;
  zoom: number;
  snapToGrid: boolean;
  showMiniMap: boolean;
  visualElements: VisualElement[];
  elementRegistry: Record<string, VisualElement[]>;
  fileTree: ProjectFile[];
  blocksXml: Record<string, string>;
  blockCode: Record<string, string>;
  cssRules: Record<string, CssRule[]>;

  setActiveFile: (fileId: string | null) => void;
  setActiveFileContent: (content: string) => void;
  setEditorMode: (mode: EditorMode) => void;
  setViewMode: (mode: ViewMode) => void;
  setSelectedElement: (id: string | null) => void;
  setSelectedBlock: (id: string | null) => void;
  setZoom: (zoom: number) => void;
  toggleSnapToGrid: () => void;
  toggleMiniMap: () => void;
  setVisualElements: (elements: VisualElement[]) => void;
  addVisualElement: (element: VisualElement, parentId?: string) => void;
  updateVisualElement: (id: string, updates: Partial<VisualElement>) => void;
  removeVisualElement: (id: string) => void;
  setElementRegistry: (registry: Record<string, VisualElement[]>) => void;
  syncVisualElementsToRegistry: (fileId: string) => void;
  getElementsForFile: (fileId: string) => VisualElement[];
  setFileTree: (files: ProjectFile[]) => void;
  addFile: (file: ProjectFile, parentId?: string) => void;
  removeFile: (id: string) => void;
  updateFile: (id: string, updates: Partial<ProjectFile>) => void;
  getBlocksXml: (fileId: string) => string;
  setBlocksXml: (fileId: string, xml: string) => void;
  setBlockCode: (fileId: string, code: string) => void;
  getBlockCode: (fileId: string) => string;
  saveGeneratedCodeToFile: (fileId: string, code: string) => void;
  setCssRules: (fileId: string, rules: CssRule[]) => void;
  getCssRules: (fileId: string) => CssRule[];
}

export const useEditorStore = create<EditorStore>((set, get) => ({
  activeFileId: null,
  activeFileContent: '',
  editorMode: 'blocks',
  viewMode: 'design',
  selectedElementId: null,
  selectedBlockId: null,
  zoom: 1,
  snapToGrid: true,
  showMiniMap: true,
  visualElements: [],
  elementRegistry: {},
  fileTree: [],
  blocksXml: {},
  blockCode: {},
  cssRules: {},

  setActiveFile: (fileId) => {
    const state = get();
    if (fileId && fileId !== state.activeFileId) {
      // Cache current visual elements in registry before switching
      const registry = { ...state.elementRegistry };
      if (state.activeFileId) {
        registry[state.activeFileId] = state.visualElements || [];
      }
      const file = findFileById(state.fileTree, fileId);
      set({
        activeFileId: fileId,
        activeFileContent: file?.content || '',
        visualElements: registry[fileId] || [],
        elementRegistry: registry,
      });
    } else {
      set({ activeFileId: fileId });
    }
  },
  setActiveFileContent: (content) => set({ activeFileContent: content }),
  setEditorMode: (mode) => set({ editorMode: mode }),
  setViewMode: (mode) => set({ viewMode: mode }),
  setSelectedElement: (id) => set({ selectedElementId: id }),
  setSelectedBlock: (id) => set({ selectedBlockId: id }),
  setZoom: (zoom) => set({ zoom: Math.max(0.25, Math.min(2, zoom)) }),
  toggleSnapToGrid: () => set((s) => ({ snapToGrid: !s.snapToGrid })),
  toggleMiniMap: () => set((s) => ({ showMiniMap: !s.showMiniMap })),
  setVisualElements: (elements) => set({ visualElements: elements || [] }),
  addVisualElement: (element, parentId) =>
    set((state) => {
      const elements = state.visualElements || [];
      let newElements: VisualElement[];
      if (!parentId) {
        newElements = [...elements, element];
      } else {
        const addToParent = (els: VisualElement[]): VisualElement[] =>
          (els || []).map((el) => {
            if (el.id === parentId) {
              return { ...el, children: [...(el.children || []), element] };
            }
            if (el.children) {
              return { ...el, children: addToParent(el.children) };
            }
            return el;
          });
        newElements = addToParent(elements);
      }
      // Sync to registry if active file is set
      const registry = { ...state.elementRegistry };
      if (state.activeFileId) {
        registry[state.activeFileId] = newElements;
      }
      return { visualElements: newElements, elementRegistry: registry };
    }),
  updateVisualElement: (id, updates) =>
    set((state) => {
      const elements = state.visualElements || [];
      const updateRecursive = (els: VisualElement[]): VisualElement[] =>
        (els || []).map((el) => {
          if (el.id === id) return { ...el, ...updates };
          if (el.children) return { ...el, children: updateRecursive(el.children) };
          return el;
        });
      const newElements = updateRecursive(elements);
      const registry = { ...state.elementRegistry };
      if (state.activeFileId) {
        registry[state.activeFileId] = newElements;
      }
      return { visualElements: newElements, elementRegistry: registry };
    }),
  removeVisualElement: (id) =>
    set((state) => {
      const elements = state.visualElements || [];
      const removeRecursive = (els: VisualElement[]): VisualElement[] =>
        (els || []).filter((el) => {
          if (el.id === id) return false;
          if (el.children) el.children = removeRecursive(el.children || []);
          return true;
        });
      const newElements = removeRecursive(elements);
      const registry = { ...state.elementRegistry };
      if (state.activeFileId) {
        registry[state.activeFileId] = newElements;
      }
      return {
        visualElements: newElements,
        elementRegistry: registry,
        selectedElementId: state.selectedElementId === id ? null : state.selectedElementId,
      };
    }),
  setElementRegistry: (registry) => set({ elementRegistry: registry || {} }),
  syncVisualElementsToRegistry: (fileId) =>
    set((state) => ({
      elementRegistry: { ...state.elementRegistry, [fileId]: state.visualElements || [] },
    })),
  getElementsForFile: (fileId) => {
    return get().elementRegistry[fileId] || [];
  },
  setFileTree: (files) => set({ fileTree: files || [] }),
  addFile: (file, parentId) =>
    set((state) => {
      const files = state.fileTree || [];
      if (!parentId) return { fileTree: [...files, file] };
      const addToFolder = (files: ProjectFile[]): ProjectFile[] =>
        (files || []).map((f) => {
          if (f.id === parentId && f.type === 'folder') {
            return { ...f, children: [...(f.children || []), file] };
          }
          if (f.children) return { ...f, children: addToFolder(f.children) };
          return f;
        });
      return { fileTree: addToFolder(files) };
    }),
  removeFile: (id) =>
    set((state) => {
      const files = state.fileTree || [];
      const removeRecursive = (files: ProjectFile[]): ProjectFile[] =>
        (files || []).filter((f) => {
          if (f.id === id) return false;
          if (f.children) f.children = removeRecursive(f.children || []);
          return true;
        });
      return {
        fileTree: removeRecursive(files),
        activeFileId: state.activeFileId === id ? null : state.activeFileId,
      };
    }),
  getBlocksXml: (fileId: string) => {
    return get().blocksXml[fileId] || '';
  },
  setBlocksXml: (fileId: string, xml: string) =>
    set((state) => ({
      blocksXml: { ...state.blocksXml, [fileId]: xml },
    })),
  setBlockCode: (fileId: string, code: string) =>
    set((state) => ({
      blockCode: { ...state.blockCode, [fileId]: code },
    })),
  getBlockCode: (fileId: string) => {
    return get().blockCode[fileId] || '';
  },
  saveGeneratedCodeToFile: (fileId: string, code: string) =>
    set((state) => {
      const files = state.fileTree || [];
      const updateRecursive = (entries: ProjectFile[]): ProjectFile[] =>
        (entries || []).map((f) => {
          if (f.id === fileId) return { ...f, content: code };
          if (f.children) return { ...f, children: updateRecursive(f.children) };
          return f;
        });
      return {
        fileTree: updateRecursive(files),
        activeFileContent: fileId === state.activeFileId ? code : state.activeFileContent,
      };
    }),
  updateFile: (id, updates) =>
    set((state) => {
      const files = state.fileTree || [];
      const updateRecursive = (files: ProjectFile[]): ProjectFile[] =>
        (files || []).map((f) => {
          if (f.id === id) return { ...f, ...updates };
          if (f.children) return { ...f, children: updateRecursive(f.children) };
          return f;
        });
      return { fileTree: updateRecursive(files) };
    }),
  setCssRules: (fileId, rules) =>
    set((state) => ({
      cssRules: { ...state.cssRules, [fileId]: rules },
    })),
  getCssRules: (fileId) => {
    return get().cssRules[fileId] || [];
  },
}));

function findFileById(files: ProjectFile[], id: string): ProjectFile | null {
  for (const f of files) {
    if (f.id === id) return f;
    if (f.children) {
      const found = findFileById(f.children, id);
      if (found) return found;
    }
  }
  return null;
}