import { useState, useCallback, useRef, useEffect } from 'react'
import {
  Section,
  Block,
  BlockType,
  ColumnLayout,
  EditorSelection,
  createBlock,
  createRow,
  createSection,
  getColumnCount,
} from '../types'
import {
  UndoRedoState,
  createUndoRedoState,
  pushState,
  undo as undoState,
  redo as redoState,
  canUndo as canUndoFn,
  canRedo as canRedoFn,
} from '../utils/undo-redo'

export interface UseEmailEditorStateOptions {
  initialSections: Section[]
  onChange?: (sections: Section[]) => void
}

export interface UseEmailEditorStateReturn {
  sections: Section[]
  selection: EditorSelection | null
  setSelection: (s: EditorSelection | null) => void
  updateBlock: (
    sectionIdx: number,
    rowIdx: number,
    colIdx: number,
    blockIdx: number,
    updates: Partial<Block>
  ) => void
  addBlock: (
    sectionIdx: number,
    rowIdx: number,
    colIdx: number,
    type: BlockType,
    afterBlockIdx?: number
  ) => void
  removeBlock: (
    sectionIdx: number,
    rowIdx: number,
    colIdx: number,
    blockIdx: number
  ) => void
  duplicateBlock: (
    sectionIdx: number,
    rowIdx: number,
    colIdx: number,
    blockIdx: number
  ) => void
  moveBlock: (
    sectionIdx: number,
    rowIdx: number,
    colIdx: number,
    blockIdx: number,
    direction: 'up' | 'down'
  ) => void
  addSection: (afterIdx?: number) => void
  removeSection: (sectionIdx: number) => void
  addRow: (sectionIdx: number, layout?: ColumnLayout) => void
  changeRowLayout: (
    sectionIdx: number,
    rowIdx: number,
    layout: ColumnLayout
  ) => void
  removeRow: (sectionIdx: number, rowIdx: number) => void
  undo: () => void
  redo: () => void
  canUndo: boolean
  canRedo: boolean
  /** Low-level commit for consumers that need direct section mutation (e.g. drag-and-drop). */
  commitChange: (newSections: Section[]) => void
}

export function useEmailEditorState({
  initialSections,
  onChange,
}: UseEmailEditorStateOptions): UseEmailEditorStateReturn {
  const [selection, setSelection] = useState<EditorSelection | null>(null)
  const historyRef = useRef<UndoRedoState>(createUndoRedoState(initialSections))

  // Sync history when sections change externally
  useEffect(() => {
    historyRef.current = { ...historyRef.current, present: initialSections }
  }, [initialSections])

  const sections = initialSections

  const commitChange = useCallback(
    (newSections: Section[]) => {
      historyRef.current = pushState(historyRef.current, newSections)
      onChange?.(newSections)
    },
    [onChange]
  )

  const handleUndo = useCallback(() => {
    const newState = undoState(historyRef.current)
    historyRef.current = newState
    onChange?.(newState.present)
  }, [onChange])

  const handleRedo = useCallback(() => {
    const newState = redoState(historyRef.current)
    historyRef.current = newState
    onChange?.(newState.present)
  }, [onChange])

  // Keyboard shortcuts for undo/redo
  useEffect(() => {
    const handler = (e: KeyboardEvent) => {
      if ((e.metaKey || e.ctrlKey) && e.key === 'z' && !e.shiftKey) {
        e.preventDefault()
        handleUndo()
      }
      if (
        (e.metaKey || e.ctrlKey) &&
        (e.key === 'y' || (e.key === 'z' && e.shiftKey))
      ) {
        e.preventDefault()
        handleRedo()
      }
    }
    window.addEventListener('keydown', handler)
    return () => window.removeEventListener('keydown', handler)
  }, [handleUndo, handleRedo])

  // --- Mutation helpers ---

  const updateBlock = useCallback(
    (
      sectionIdx: number,
      rowIdx: number,
      colIdx: number,
      blockIdx: number,
      updates: Partial<Block>
    ) => {
      const newSections = sections.map((section, si) => {
        if (si !== sectionIdx) return section
        return {
          ...section,
          rows: section.rows.map((row, ri) => {
            if (ri !== rowIdx) return row
            return {
              ...row,
              columns: row.columns.map((col, ci) => {
                if (ci !== colIdx) return col
                return col.map((block, bi) =>
                  bi === blockIdx ? ({ ...block, ...updates } as Block) : block
                )
              }),
            }
          }),
        }
      })
      commitChange(newSections)
    },
    [sections, commitChange]
  )

  const addBlock = useCallback(
    (
      sectionIdx: number,
      rowIdx: number,
      colIdx: number,
      blockType: BlockType,
      afterBlockIdx?: number
    ) => {
      const newBlock = createBlock(blockType)
      const newSections = sections.map((section, si) => {
        if (si !== sectionIdx) return section
        return {
          ...section,
          rows: section.rows.map((row, ri) => {
            if (ri !== rowIdx) return row
            return {
              ...row,
              columns: row.columns.map((col, ci) => {
                if (ci !== colIdx) return col
                const insertAt =
                  afterBlockIdx !== undefined ? afterBlockIdx + 1 : col.length
                return [
                  ...col.slice(0, insertAt),
                  newBlock,
                  ...col.slice(insertAt),
                ]
              }),
            }
          }),
        }
      })
      commitChange(newSections)
      // Select the new block
      const insertAt =
        afterBlockIdx !== undefined
          ? afterBlockIdx + 1
          : sections[sectionIdx].rows[rowIdx].columns[colIdx].length
      setSelection({
        sectionIndex: sectionIdx,
        rowIndex: rowIdx,
        columnIndex: colIdx,
        blockIndex: insertAt,
      })
    },
    [sections, commitChange]
  )

  const removeBlock = useCallback(
    (
      sectionIdx: number,
      rowIdx: number,
      colIdx: number,
      blockIdx: number
    ) => {
      const newSections = sections.map((section, si) => {
        if (si !== sectionIdx) return section
        return {
          ...section,
          rows: section.rows.map((row, ri) => {
            if (ri !== rowIdx) return row
            return {
              ...row,
              columns: row.columns.map((col, ci) => {
                if (ci !== colIdx) return col
                return col.filter((_, bi) => bi !== blockIdx)
              }),
            }
          }),
        }
      })
      commitChange(newSections)
      setSelection(null)
    },
    [sections, commitChange]
  )

  const duplicateBlock = useCallback(
    (
      sectionIdx: number,
      rowIdx: number,
      colIdx: number,
      blockIdx: number
    ) => {
      const block =
        sections[sectionIdx].rows[rowIdx].columns[colIdx][blockIdx]
      if (!block) return
      const copy: Block = {
        ...JSON.parse(JSON.stringify(block)),
        id: `block-${Date.now()}-${Math.random().toString(36).substr(2, 6)}`,
      }
      const newSections = sections.map((section, si) => {
        if (si !== sectionIdx) return section
        return {
          ...section,
          rows: section.rows.map((row, ri) => {
            if (ri !== rowIdx) return row
            return {
              ...row,
              columns: row.columns.map((col, ci) => {
                if (ci !== colIdx) return col
                return [
                  ...col.slice(0, blockIdx + 1),
                  copy,
                  ...col.slice(blockIdx + 1),
                ]
              }),
            }
          }),
        }
      })
      commitChange(newSections)
      setSelection({
        sectionIndex: sectionIdx,
        rowIndex: rowIdx,
        columnIndex: colIdx,
        blockIndex: blockIdx + 1,
      })
    },
    [sections, commitChange]
  )

  const moveBlock = useCallback(
    (
      sectionIdx: number,
      rowIdx: number,
      colIdx: number,
      blockIdx: number,
      direction: 'up' | 'down'
    ) => {
      const col = sections[sectionIdx].rows[rowIdx].columns[colIdx]
      const newIndex = direction === 'up' ? blockIdx - 1 : blockIdx + 1
      if (newIndex < 0 || newIndex >= col.length) return
      const newCol = [...col]
      const [moved] = newCol.splice(blockIdx, 1)
      newCol.splice(newIndex, 0, moved)
      const newSections = sections.map((section, si) => {
        if (si !== sectionIdx) return section
        return {
          ...section,
          rows: section.rows.map((row, ri) => {
            if (ri !== rowIdx) return row
            return {
              ...row,
              columns: row.columns.map((c, ci) =>
                ci === colIdx ? newCol : c
              ),
            }
          }),
        }
      })
      commitChange(newSections)
      setSelection({
        sectionIndex: sectionIdx,
        rowIndex: rowIdx,
        columnIndex: colIdx,
        blockIndex: newIndex,
      })
    },
    [sections, commitChange]
  )

  // --- Section/Row operations ---

  const addSection = useCallback(
    (afterIdx?: number) => {
      const newSection = createSection()
      const insertAfter = afterIdx ?? sections.length - 1
      const newSections = [
        ...sections.slice(0, insertAfter + 1),
        newSection,
        ...sections.slice(insertAfter + 1),
      ]
      commitChange(newSections)
    },
    [sections, commitChange]
  )

  const removeSection = useCallback(
    (index: number) => {
      if (sections.length <= 1) return
      commitChange(sections.filter((_, i) => i !== index))
      setSelection(null)
    },
    [sections, commitChange]
  )

  const addRow = useCallback(
    (sectionIdx: number, layout: ColumnLayout = '1') => {
      const newRow = createRow(layout)
      const newSections = sections.map((section, si) => {
        if (si !== sectionIdx) return section
        return { ...section, rows: [...section.rows, newRow] }
      })
      commitChange(newSections)
    },
    [sections, commitChange]
  )

  const changeRowLayout = useCallback(
    (sectionIdx: number, rowIdx: number, layout: ColumnLayout) => {
      const newColCount = getColumnCount(layout)
      const newSections = sections.map((section, si) => {
        if (si !== sectionIdx) return section
        return {
          ...section,
          rows: section.rows.map((row, ri) => {
            if (ri !== rowIdx) return row
            const columns = [...row.columns]
            // Add or trim columns as needed
            while (columns.length < newColCount) columns.push([])
            // If reducing columns, merge extra columns into last
            if (columns.length > newColCount) {
              const extra = columns.splice(newColCount - 1)
              columns[newColCount - 1] = extra.flat()
            }
            return { ...row, layout, columns }
          }),
        }
      })
      commitChange(newSections)
    },
    [sections, commitChange]
  )

  const removeRow = useCallback(
    (sectionIdx: number, rowIdx: number) => {
      const section = sections[sectionIdx]
      if (section.rows.length <= 1) return
      const newSections = sections.map((s, si) => {
        if (si !== sectionIdx) return s
        return { ...s, rows: s.rows.filter((_, ri) => ri !== rowIdx) }
      })
      commitChange(newSections)
      setSelection(null)
    },
    [sections, commitChange]
  )

  return {
    sections,
    selection,
    setSelection,
    updateBlock,
    addBlock,
    removeBlock,
    duplicateBlock,
    moveBlock,
    addSection,
    removeSection,
    addRow,
    changeRowLayout,
    removeRow,
    undo: handleUndo,
    redo: handleRedo,
    canUndo: canUndoFn(historyRef.current),
    canRedo: canRedoFn(historyRef.current),
    commitChange,
  }
}