import { create } from 'zustand'
import { subscribeWithSelector } from 'zustand/middleware'
import { useShallow } from 'zustand/react/shallow'
import type {
  Kernel,
  ConvolutionState,
  AppSettings,
  ComparisonState,
  KernelComparisonResult,
  PaddingMode,
} from '@/types'

type ConvolutionStoreState = ConvolutionState & {
  startConvolution: () => void
  pauseConvolution: () => void
  resumeConvolution: () => void
  stopConvolution: () => void
  updateConvolutionStep: (step: ConvolutionState['currentStep']) => void
  updateConvolutionProgress: (progress: number, processedPixels: number, elapsedTime: number) => void
  completeConvolution: (outputImageData: ImageData) => void
  setConvolutionError: (error: string) => void
  resetConvolution: () => void
}

type SettingsStoreState = AppSettings & {
  setAnimationSpeed: (speed: number) => void
  toggleOverlay: (key: keyof AppSettings['showOverlay']) => void
  setPadding: (padding: PaddingMode) => void
  setZoomLevel: (zoom: number) => void
  setPanOffset: (offset: { x: number; y: number }) => void
  setComparisonMode: (mode: AppSettings['comparisonMode']) => void
  setStepMode: (enabled: boolean) => void
}

type KernelStoreState = {
  selectedKernel: Kernel | null
  customKernelMatrix: number[][]
  customKernelSize: number
  customKernelNormalize: boolean
  setSelectedKernel: (kernel: Kernel | null) => void
  setCustomKernel: (matrix: number[][], size: number, normalize: boolean) => void
  updateCustomKernelCell: (row: number, col: number, value: number) => void
  setCustomKernelSize: (size: number) => void
  setCustomKernelNormalize: (normalize: boolean) => void
}

type ImageStoreState = {
  originalImage: HTMLImageElement | null
  originalImageData: ImageData | null
  processedImageData: ImageData | null
  setOriginalImage: (image: HTMLImageElement, imageData: ImageData) => void
}

type ComparisonStoreState = ComparisonState & {
  setComparisonViewMode: (mode: ComparisonState['mode']) => void
  setSplitPosition: (position: number) => void
  toggleSynchronizedZoom: () => void
  toggleSynchronizedPan: () => void
  addComparisonResult: (result: KernelComparisonResult) => void
  clearComparisonResults: () => void
  setActiveComparisonTab: (kernelId: string) => void
}

type UIStoreState = {
  isLoading: boolean
  error: string | null
  showHelp: boolean
  showSettings: boolean
  setLoading: (loading: boolean) => void
  setError: (error: string | null) => void
  toggleHelp: () => void
  toggleSettings: () => void
}

interface AppStore {
  // Image state
  originalImage: HTMLImageElement | null
  originalImageData: ImageData | null
  processedImageData: ImageData | null

  // Kernel state
  selectedKernel: Kernel | null
  customKernelMatrix: number[][]
  customKernelSize: number
  customKernelNormalize: boolean

  // Convolution state
  convolution: ConvolutionState

  // Settings
  settings: AppSettings

  // Comparison
  comparison: ComparisonState
  comparisonResults: KernelComparisonResult[]

  // UI state
  isLoading: boolean
  error: string | null
  showHelp: boolean
  showSettings: boolean

  // Actions
  setOriginalImage: (image: HTMLImageElement, imageData: ImageData) => void
  setSelectedKernel: (kernel: Kernel | null) => void
  setCustomKernel: (matrix: number[][], size: number, normalize: boolean) => void
  updateCustomKernelCell: (row: number, col: number, value: number) => void
  setCustomKernelSize: (size: number) => void
  setCustomKernelNormalize: (normalize: boolean) => void

  // Convolution actions
  startConvolution: () => void
  pauseConvolution: () => void
  resumeConvolution: () => void
  stopConvolution: () => void
  updateConvolutionStep: (step: ConvolutionState['currentStep']) => void
  updateConvolutionProgress: (progress: number, processedPixels: number, elapsedTime: number) => void
  completeConvolution: (outputImageData: ImageData) => void
  setConvolutionError: (error: string) => void
  resetConvolution: () => void

  // Settings actions
  updateSettings: (settings: Partial<AppSettings>) => void
  setAnimationSpeed: (speed: number) => void
  toggleOverlay: (key: keyof AppSettings['showOverlay']) => void
  setPadding: (padding: PaddingMode) => void
  setZoomLevel: (zoomLevel: number) => void
  setPanOffset: (panOffset: { x: number; y: number }) => void
  setComparisonMode: (comparisonMode: AppSettings['comparisonMode']) => void
  setStepMode: (stepMode: boolean) => void

  // Comparison actions
  setComparisonViewMode: (mode: ComparisonState['mode']) => void
  setSplitPosition: (splitPosition: number) => void
  toggleSynchronizedZoom: () => void
  toggleSynchronizedPan: () => void
  addComparisonResult: (result: KernelComparisonResult) => void
  clearComparisonResults: () => void
  setActiveComparisonTab: (activeTab: string) => void

  // UI actions
  setLoading: (isLoading: boolean) => void
  setError: (error: string | null) => void
  toggleHelp: () => void
  toggleSettings: () => void
}

const defaultSettings: AppSettings = {
  animationSpeed: 1,
  showOverlay: {
    kernelBoundary: true,
    pixelValues: true,
    multiplications: true,
    summation: true,
  },
  padding: 'zero',
  zoomLevel: 1,
  panOffset: { x: 0, y: 0 },
  comparisonMode: 'split',
  stepMode: false,
}

const defaultConvolutionState: ConvolutionState = {
  status: 'idle',
  currentStep: null,
  progress: 0,
  startTime: 0,
  elapsedTime: 0,
  processedPixels: 0,
  totalPixels: 0,
  outputImageData: null,
  error: null,
}

const defaultComparisonState: ComparisonState = {
  mode: 'split',
  splitPosition: 0.5,
  synchronizedZoom: true,
  synchronizedPan: true,
  activeTab: '',
}

export const useAppStore = create<AppStore>()(
  subscribeWithSelector((set) => ({
    originalImage: null,
    originalImageData: null,
    processedImageData: null,
    selectedKernel: null,
    customKernelMatrix: Array(3)
      .fill(0)
      .map(() => Array(3).fill(0)),
    customKernelSize: 3,
    customKernelNormalize: true,
    convolution: defaultConvolutionState,
    settings: defaultSettings,
    comparison: defaultComparisonState,
    comparisonResults: [],
    isLoading: false,
    error: null,
    showHelp: false,
    showSettings: false,

    setOriginalImage: (image, imageData) =>
      set({ originalImage: image, originalImageData: imageData, processedImageData: null }),

    setSelectedKernel: (kernel) => set({ selectedKernel: kernel }),

    setCustomKernel: (matrix, size, normalize) =>
      set({ customKernelMatrix: matrix, customKernelSize: size, customKernelNormalize: normalize }),

    updateCustomKernelCell: (row, col, value) =>
      set((state) => {
        const newMatrix = state.customKernelMatrix.map((r, i) =>
          i === row ? r.map((v, j) => (j === col ? value : v)) : r
        )
        return { customKernelMatrix: newMatrix }
      }),

    setCustomKernelSize: (size) =>
      set((state) => {
        const newMatrix = Array(size)
          .fill(0)
          .map(() => Array(size).fill(0))
        return { customKernelSize: size, customKernelMatrix: newMatrix }
      }),

    setCustomKernelNormalize: (normalize) => set({ customKernelNormalize: normalize }),

    startConvolution: () =>
      set((state) => ({
        convolution: {
          ...state.convolution,
          status: 'running',
          startTime: Date.now(),
          progress: 0,
          processedPixels: 0,
          error: null,
        },
      })),

    pauseConvolution: () =>
      set((state) => ({
        convolution: {
          ...state.convolution,
          status: 'paused',
          elapsedTime: Date.now() - state.convolution.startTime,
        },
      })),

    resumeConvolution: () =>
      set((state) => ({
        convolution: {
          ...state.convolution,
          status: 'running',
        },
      })),

    stopConvolution: () =>
      set((state) => ({
        convolution: {
          ...state.convolution,
          status: 'idle',
          currentStep: null,
          progress: 0,
        },
      })),

    updateConvolutionStep: (step) =>
      set((state) => {
        const newState = {
          convolution: {
            ...state.convolution,
            currentStep: step,
          },
        }
        // If this step has a partial output image, update the processed image data
        if (step && step.isPartialOutput && step.partialOutputImageData) {
          return {
            ...newState,
            processedImageData: step.partialOutputImageData,
          }
        }
        return newState
      }),

    updateConvolutionProgress: (progress, processedPixels, elapsedTime) =>
      set((state) => ({
        convolution: {
          ...state.convolution,
          progress,
          processedPixels,
          elapsedTime,
        },
      })),

    completeConvolution: (outputImageData) =>
      set((state) => ({
        convolution: {
          ...state.convolution,
          status: 'completed',
          progress: 1,
          outputImageData,
          elapsedTime: Date.now() - state.convolution.startTime,
        },
        processedImageData: outputImageData,
      })),

    setConvolutionError: (error) =>
      set((state) => ({
        convolution: {
          ...state.convolution,
          status: 'error',
          error,
        },
      })),

    resetConvolution: () =>
      set({
        convolution: defaultConvolutionState,
        processedImageData: null,
      }),

    updateSettings: (newSettings) =>
      set((state) => ({
        settings: { ...state.settings, ...newSettings },
      })),

    setAnimationSpeed: (speed) =>
      set((state) => ({
        settings: { ...state.settings, animationSpeed: speed },
      })),

    toggleOverlay: (key) =>
      set((state) => ({
        settings: {
          ...state.settings,
          showOverlay: { ...state.settings.showOverlay, [key]: !state.settings.showOverlay[key] },
        },
      })),

    setPadding: (padding) =>
      set((state) => ({
        settings: { ...state.settings, padding },
      })),

    setZoomLevel: (zoomLevel) =>
      set((state) => ({
        settings: { ...state.settings, zoomLevel },
      })),

    setPanOffset: (panOffset) =>
      set((state) => ({
        settings: { ...state.settings, panOffset },
      })),

    setComparisonMode: (comparisonMode) =>
      set((state) => ({
        settings: { ...state.settings, comparisonMode },
      })),

    setStepMode: (stepMode) =>
      set((state) => ({
        settings: { ...state.settings, stepMode },
      })),

    // Comparison actions
    setComparisonViewMode: (mode: ComparisonState['mode']) =>
      set((state) => ({
        comparison: { ...state.comparison, mode },
      })),

    setSplitPosition: (splitPosition: number) =>
      set((state) => ({
        comparison: { ...state.comparison, splitPosition },
      })),

    toggleSynchronizedZoom: () =>
      set((state) => ({
        comparison: { ...state.comparison, synchronizedZoom: !state.comparison.synchronizedZoom },
      })),

    toggleSynchronizedPan: () =>
      set((state) => ({
        comparison: { ...state.comparison, synchronizedPan: !state.comparison.synchronizedPan },
      })),

    addComparisonResult: (result) =>
      set((state) => ({
        comparisonResults: [...state.comparisonResults, result],
      })),

    clearComparisonResults: () => set({ comparisonResults: [] }),

    setActiveComparisonTab: (activeTab: string) =>
      set((state) => ({
        comparison: { ...state.comparison, activeTab },
      })),

    setLoading: (isLoading) => set({ isLoading }),

    setError: (error) => set({ error }),

    toggleHelp: () => set((state) => ({ showHelp: !state.showHelp })),

    toggleSettings: () => set((state) => ({ showSettings: !state.showSettings })),
  }))
)

// Selector hooks that include both state and actions
export const useConvolutionStore = () => useAppStore(useShallow((state): ConvolutionStoreState => ({
  ...state.convolution,
  startConvolution: state.startConvolution,
  pauseConvolution: state.pauseConvolution,
  resumeConvolution: state.resumeConvolution,
  stopConvolution: state.stopConvolution,
  updateConvolutionStep: state.updateConvolutionStep,
  updateConvolutionProgress: state.updateConvolutionProgress,
  completeConvolution: state.completeConvolution,
  setConvolutionError: state.setConvolutionError,
  resetConvolution: state.resetConvolution,
})))

export const useSettingsStore = () => useAppStore(useShallow((state): SettingsStoreState => ({
  ...state.settings,
  setAnimationSpeed: state.setAnimationSpeed,
  toggleOverlay: state.toggleOverlay,
  setPadding: state.setPadding,
  setZoomLevel: state.setZoomLevel,
  setPanOffset: state.setPanOffset,
  setComparisonMode: state.setComparisonMode,
  setStepMode: state.setStepMode,
})))

export const useKernelStore = () => useAppStore(useShallow((state): KernelStoreState => ({
  selectedKernel: state.selectedKernel,
  customKernelMatrix: state.customKernelMatrix,
  customKernelSize: state.customKernelSize,
  customKernelNormalize: state.customKernelNormalize,
  setSelectedKernel: state.setSelectedKernel,
  setCustomKernel: state.setCustomKernel,
  updateCustomKernelCell: state.updateCustomKernelCell,
  setCustomKernelSize: state.setCustomKernelSize,
  setCustomKernelNormalize: state.setCustomKernelNormalize,
})))

export const useImageStore = () => useAppStore(useShallow((state): ImageStoreState => ({
  originalImage: state.originalImage,
  originalImageData: state.originalImageData,
  processedImageData: state.processedImageData,
  setOriginalImage: state.setOriginalImage,
})))

export const useComparisonStore = () => useAppStore(useShallow((state): ComparisonStoreState => ({
  ...state.comparison,
  setComparisonViewMode: state.setComparisonViewMode,
  setSplitPosition: state.setSplitPosition,
  toggleSynchronizedZoom: state.toggleSynchronizedZoom,
  toggleSynchronizedPan: state.toggleSynchronizedPan,
  addComparisonResult: state.addComparisonResult,
  clearComparisonResults: state.clearComparisonResults,
  setActiveComparisonTab: state.setActiveComparisonTab,
})))

export const useUIStore = () => useAppStore(useShallow((state): UIStoreState => ({
  isLoading: state.isLoading,
  error: state.error,
  showHelp: state.showHelp,
  showSettings: state.showSettings,
  setLoading: state.setLoading,
  setError: state.setError,
  toggleHelp: state.toggleHelp,
  toggleSettings: state.toggleSettings,
})))