import { ComfyButton } from '../components/button'
import { prop, getStorageValue, setStorageValue } from '../../utils'
import { $el } from '../../ui'
import { api } from '../../api'
import { ComfyPopup } from '../components/popup'
import { createSpinner } from '../spinner'
import { ComfyWorkflow, trimJsonExt } from '../../workflows'
import { ComfyAsyncDialog } from '../components/asyncDialog'
import type { ComfyApp } from '@/scripts/app'
import type { ComfyComponent } from '../components'

export class ComfyWorkflowsMenu {
  #first = true
  element = $el('div.comfyui-workflows')
  popup: ComfyPopup
  app: ComfyApp
  buttonProgress: HTMLElement
  workflowLabel: HTMLElement
  button: ComfyButton
  content: ComfyWorkflowsContent
  unsaved: boolean

  get open() {
    return this.popup.open
  }

  set open(open) {
    this.popup.open = open
  }

  constructor(app: ComfyApp) {
    this.app = app
    this.#bindEvents()

    const classList = {
      'comfyui-workflows-button': true,
      'comfyui-button': true,
      unsaved: getStorageValue('Comfy.PreviousWorkflowUnsaved') === 'true',
      running: false
    }
    this.buttonProgress = $el('div.comfyui-workflows-button-progress')
    this.workflowLabel = $el('span.comfyui-workflows-label', '')
    this.button = new ComfyButton({
      content: $el('div.comfyui-workflows-button-inner', [
        $el('i.mdi.mdi-graph'),
        this.workflowLabel,
        this.buttonProgress
      ]),
      icon: 'chevron-down',
      classList
    })

    this.element.append(this.button.element)

    this.popup = new ComfyPopup({
      target: this.element,
      classList: 'comfyui-workflows-popup'
    })
    this.content = new ComfyWorkflowsContent(app, this.popup)
    this.popup.children = [this.content.element]
    this.popup.addEventListener('change', () => {
      this.button.icon = 'chevron-' + (this.popup.open ? 'up' : 'down')
    })
    this.button.withPopup(this.popup)

    this.unsaved = prop(this, 'unsaved', classList.unsaved, (v) => {
      classList.unsaved = v
      this.button.classList = classList
      setStorageValue('Comfy.PreviousWorkflowUnsaved', v)
    })
  }

  #updateProgress = () => {
    const prompt = this.app.workflowManager.activePrompt
    let percent = 0
    if (this.app.workflowManager.activeWorkflow === prompt?.workflow) {
      const total = Object.values(prompt.nodes)
      const done = total.filter(Boolean)
      percent = (done.length / total.length) * 100
    }
    this.buttonProgress.style.width = percent + '%'
  }

  #updateActive = () => {
    const active = this.app.workflowManager.activeWorkflow
    this.button.tooltip = active.path
    this.workflowLabel.textContent = active.name
    this.unsaved = active.unsaved

    if (this.#first) {
      this.#first = false
      this.content.load()
    }

    this.#updateProgress()
  }

  #bindEvents() {
    this.app.workflowManager.addEventListener(
      'changeWorkflow',
      this.#updateActive
    )
    this.app.workflowManager.addEventListener('rename', this.#updateActive)
    this.app.workflowManager.addEventListener('delete', this.#updateActive)

    this.app.workflowManager.addEventListener('save', () => {
      this.unsaved = this.app.workflowManager.activeWorkflow.unsaved
    })

    this.app.workflowManager.addEventListener('execute', (e) => {
      this.#updateProgress()
    })

    api.addEventListener('graphChanged', () => {
      this.unsaved = true
    })
  }

  #getMenuOptions(callback) {
    const menu = []
    const directories = new Map()
    for (const workflow of this.app.workflowManager.workflows || []) {
      const path = workflow.pathParts
      if (!path) continue
      let parent = menu
      let currentPath = ''
      for (let i = 0; i < path.length - 1; i++) {
        currentPath += '/' + path[i]
        let newParent = directories.get(currentPath)
        if (!newParent) {
          newParent = {
            title: path[i],
            has_submenu: true,
            submenu: {
              options: []
            }
          }
          parent.push(newParent)
          newParent = newParent.submenu.options
          directories.set(currentPath, newParent)
        }
        parent = newParent
      }
      parent.push({
        title: trimJsonExt(path[path.length - 1]),
        callback: () => callback(workflow)
      })
    }
    return menu
  }

  #getFavoriteMenuOptions(callback) {
    const menu = []
    for (const workflow of this.app.workflowManager.workflows || []) {
      if (workflow.isFavorite) {
        menu.push({
          title: '⭐ ' + workflow.name,
          callback: () => callback(workflow)
        })
      }
    }
    return menu
  }

  registerExtension(app: ComfyApp) {
    const self = this
    app.registerExtension({
      name: 'Comfy.Workflows',
      async beforeRegisterNodeDef(nodeType) {
        function getImageWidget(node) {
          const inputs = {
            ...node.constructor?.nodeData?.input?.required,
            ...node.constructor?.nodeData?.input?.optional
          }
          for (const input in inputs) {
            if (inputs[input][0] === 'IMAGEUPLOAD') {
              const imageWidget = node.widgets.find(
                (w) => w.name === (inputs[input]?.[1]?.widget ?? 'image')
              )
              if (imageWidget) return imageWidget
            }
          }
        }

        function setWidgetImage(node, widget, img) {
          const url = new URL(img.src)
          const filename = url.searchParams.get('filename')
          const subfolder = url.searchParams.get('subfolder')
          const type = url.searchParams.get('type')
          const imageId = `${subfolder ? subfolder + '/' : ''}${filename} [${type}]`
          widget.value = imageId
          node.imgs = [img]
          app.graph.setDirtyCanvas(true, true)
        }

        async function sendToWorkflow(
          img: HTMLImageElement,
          workflow: ComfyWorkflow
        ) {
          const openWorkflow = app.workflowManager.openWorkflows.find(
            (w) => w.path === workflow.path
          )
          if (openWorkflow) {
            workflow = openWorkflow
          }

          await workflow.load()
          let options = []
          const nodes = app.graph.computeExecutionOrder(false)
          for (const node of nodes) {
            const widget = getImageWidget(node)
            if (widget == null) continue

            if (node.title?.toLowerCase().includes('input')) {
              options = [{ widget, node }]
              break
            } else {
              options.push({ widget, node })
            }
          }

          if (!options.length) {
            alert('No image nodes have been found in this workflow!')
            return
          } else if (options.length > 1) {
            const dialog = new WidgetSelectionDialog(options)
            const res = await dialog.show(app)
            if (!res) return
            options = [res]
          }

          setWidgetImage(options[0].node, options[0].widget, img)
        }

        const getExtraMenuOptions = nodeType.prototype['getExtraMenuOptions']
        nodeType.prototype['getExtraMenuOptions'] = function (
          this: { imageIndex?: number; overIndex?: number; imgs: string[] },
          _,
          options
        ) {
          const r = getExtraMenuOptions?.apply?.(this, arguments)
          const setting = app.ui.settings.getSettingValue(
            'Comfy.UseNewMenu',
            false
          )
          if (setting && setting != 'Disabled') {
            const t = this
            let img
            if (t.imageIndex != null) {
              // An image is selected so select that
              img = t.imgs?.[t.imageIndex]
            } else if (t.overIndex != null) {
              // No image is selected but one is hovered
              img = t.imgs?.[t.overIndex]
            }

            if (img) {
              let pos = options.findIndex((o) => o.content === 'Save Image')
              if (pos === -1) {
                pos = 0
              } else {
                pos++
              }

              options.splice(pos, 0, {
                content: 'Send to workflow',
                has_submenu: true,
                submenu: {
                  options: [
                    {
                      callback: () =>
                        sendToWorkflow(img, app.workflowManager.activeWorkflow),
                      title: '[Current workflow]'
                    },
                    ...self.#getFavoriteMenuOptions(
                      sendToWorkflow.bind(null, img)
                    ),
                    null,
                    ...self.#getMenuOptions(sendToWorkflow.bind(null, img))
                  ]
                }
              })
            }
          }

          return r
        }
      }
    })
  }
}

export class ComfyWorkflowsContent {
  element = $el('div.comfyui-workflows-panel')
  treeState = {}
  treeFiles: Record<string, WorkflowElement> = {}
  openFiles: Map<ComfyWorkflow, WorkflowElement<ComfyComponent>> = new Map()
  activeElement: WorkflowElement<ComfyComponent> = null
  spinner: Element
  openElement: HTMLElement
  favoritesElement: HTMLElement
  treeElement: HTMLElement
  app: ComfyApp
  popup: ComfyPopup
  actions: HTMLElement
  filterText: string | undefined
  treeRoot: HTMLElement

  constructor(app: ComfyApp, popup: ComfyPopup) {
    this.app = app
    this.popup = popup
    this.actions = $el('div.comfyui-workflows-actions', [
      new ComfyButton({
        content: 'Default',
        icon: 'file-code',
        iconSize: 18,
        classList: 'comfyui-button primary',
        tooltip: 'Load default workflow',
        action: () => {
          popup.open = false
          app.loadGraphData()
          app.resetView()
        }
      }).element,
      new ComfyButton({
        content: 'Browse',
        icon: 'folder',
        iconSize: 18,
        tooltip: 'Browse for an image or exported workflow',
        action: () => {
          popup.open = false
          app.ui.loadFile()
        }
      }).element,
      new ComfyButton({
        content: 'Blank',
        icon: 'plus-thick',
        iconSize: 18,
        tooltip: 'Create a new blank workflow',
        action: () => {
          app.workflowManager.setWorkflow(null)
          app.clean()
          app.graph.clear()
          app.workflowManager.activeWorkflow.track()
          popup.open = false
        }
      }).element
    ])

    this.spinner = createSpinner()
    this.element.replaceChildren(this.actions, this.spinner)

    this.popup.addEventListener('open', () => this.load())
    this.popup.addEventListener('close', () =>
      this.element.replaceChildren(this.actions, this.spinner)
    )

    this.app.workflowManager.addEventListener('favorite', (e) => {
      const workflow = e['detail']
      const button = this.treeFiles[workflow.path]?.primary
      if (!button) return // Can happen when a workflow is renamed
      button.icon = this.#getFavoriteIcon(workflow)
      button.overIcon = this.#getFavoriteOverIcon(workflow)
      this.updateFavorites()
    })

    for (const e of ['save', 'open', 'close', 'changeWorkflow']) {
      // TODO: dont be lazy and just update the specific element
      app.workflowManager.addEventListener(e, () => this.updateOpen())
    }
    this.app.workflowManager.addEventListener('rename', () => this.load())
    this.app.workflowManager.addEventListener('execute', (e) =>
      this.#updateActive()
    )
  }

  async load() {
    await this.app.workflowManager.loadWorkflows()
    this.updateTree()
    this.updateFavorites()
    this.updateOpen()
    this.element.replaceChildren(
      this.actions,
      this.openElement,
      this.favoritesElement,
      this.treeElement
    )
  }

  updateOpen() {
    const current = this.openElement
    this.openFiles.clear()

    this.openElement = $el('div.comfyui-workflows-open', [
      $el('h3', 'Open'),
      ...this.app.workflowManager.openWorkflows.map((w) => {
        const wrapper = new WorkflowElement(this, w, {
          primary: { element: $el('i.mdi.mdi-18px.mdi-progress-pencil') },
          buttons: [
            this.#getRenameButton(w),
            new ComfyButton({
              icon: 'close',
              iconSize: 18,
              classList: 'comfyui-button comfyui-workflows-file-action',
              tooltip: 'Close workflow',
              action: (e) => {
                e.stopImmediatePropagation()
                this.app.workflowManager.closeWorkflow(w)
              }
            })
          ]
        })
        if (w.unsaved) {
          wrapper.element.classList.add('unsaved')
        }
        if (w === this.app.workflowManager.activeWorkflow) {
          wrapper.element.classList.add('active')
        }

        this.openFiles.set(w, wrapper)
        return wrapper.element
      })
    ])

    this.#updateActive()
    current?.replaceWith(this.openElement)
  }

  updateFavorites() {
    const current = this.favoritesElement
    const favorites = [
      ...this.app.workflowManager.workflows.filter((w) => w.isFavorite)
    ]

    this.favoritesElement = $el('div.comfyui-workflows-favorites', [
      $el('h3', 'Favorites'),
      ...favorites
        .map((w) => {
          return this.#getWorkflowElement(w).element
        })
        .filter(Boolean)
    ])

    current?.replaceWith(this.favoritesElement)
  }

  filterTree() {
    if (!this.filterText) {
      this.treeRoot.classList.remove('filtered')
      // Unfilter whole tree
      for (const item of Object.values(this.treeFiles)) {
        item.element.parentElement.style.removeProperty('display')
        this.showTreeParents(item.element.parentElement)
      }
      return
    }
    this.treeRoot.classList.add('filtered')
    const searchTerms = this.filterText.toLocaleLowerCase().split(' ')
    for (const item of Object.values(this.treeFiles)) {
      const parts = item.workflow.pathParts
      let termIndex = 0
      let valid = false
      for (const part of parts) {
        let currentIndex = 0
        do {
          currentIndex = part.indexOf(searchTerms[termIndex], currentIndex)
          if (currentIndex > -1) currentIndex += searchTerms[termIndex].length
        } while (currentIndex !== -1 && ++termIndex < searchTerms.length)

        if (termIndex >= searchTerms.length) {
          valid = true
          break
        }
      }
      if (valid) {
        item.element.parentElement.style.removeProperty('display')
        this.showTreeParents(item.element.parentElement)
      } else {
        item.element.parentElement.style.display = 'none'
        this.hideTreeParents(item.element.parentElement)
      }
    }
  }

  hideTreeParents(element) {
    // Hide all parents if no children are visible
    if (
      element.parentElement?.classList.contains('comfyui-workflows-tree') ===
      false
    ) {
      for (let i = 1; i < element.parentElement.children.length; i++) {
        const c = element.parentElement.children[i]
        if (c.style.display !== 'none') {
          return
        }
      }
      element.parentElement.style.display = 'none'
      this.hideTreeParents(element.parentElement)
    }
  }

  showTreeParents(element) {
    if (
      element.parentElement?.classList.contains('comfyui-workflows-tree') ===
      false
    ) {
      element.parentElement.style.removeProperty('display')
      this.showTreeParents(element.parentElement)
    }
  }

  updateTree() {
    const current = this.treeElement
    const nodes = {}
    let typingTimeout

    this.treeFiles = {}
    this.treeRoot = $el('ul.comfyui-workflows-tree')
    this.treeElement = $el('section', [
      $el('header', [
        $el('h3', 'Browse'),
        $el('div.comfy-ui-workflows-search', [
          $el('i.mdi.mdi-18px.mdi-magnify'),
          $el('input', {
            placeholder: 'Search',
            value: this.filterText ?? '',
            oninput: (e: InputEvent) => {
              this.filterText = e.target['value']?.trim()
              clearTimeout(typingTimeout)
              typingTimeout = setTimeout(() => this.filterTree(), 250)
            }
          })
        ])
      ]),
      this.treeRoot
    ])

    for (const workflow of this.app.workflowManager.workflows) {
      if (!workflow.pathParts) continue

      let currentPath = ''
      let currentRoot = this.treeRoot

      for (let i = 0; i < workflow.pathParts.length; i++) {
        currentPath += (currentPath ? '\\' : '') + workflow.pathParts[i]
        const parentNode =
          nodes[currentPath] ??
          this.#createNode(currentPath, workflow, i, currentRoot)

        nodes[currentPath] = parentNode
        currentRoot = parentNode
      }
    }

    current?.replaceWith(this.treeElement)
    this.filterTree()
  }

  #expandNode(el, workflow, thisPath, i) {
    const expanded = !el.classList.toggle('closed')
    if (expanded) {
      let c = ''
      for (let j = 0; j <= i; j++) {
        c += (c ? '\\' : '') + workflow.pathParts[j]
        this.treeState[c] = true
      }
    } else {
      let c = thisPath
      for (let j = i + 1; j < workflow.pathParts.length; j++) {
        c += (c ? '\\' : '') + workflow.pathParts[j]
        delete this.treeState[c]
      }
      delete this.treeState[thisPath]
    }
  }

  #updateActive() {
    this.#removeActive()

    const active = this.app.workflowManager.activePrompt
    if (!active?.workflow) return

    const open = this.openFiles.get(active.workflow)
    if (!open) return

    this.activeElement = open

    const total = Object.values(active.nodes)
    const done = total.filter(Boolean)
    const percent = done.length / total.length
    open.element.classList.add('running')
    open.element.style.setProperty('--progress', percent * 100 + '%')
    open.primary.element.classList.remove('mdi-progress-pencil')
    open.primary.element.classList.add('mdi-play')
  }

  #removeActive() {
    if (!this.activeElement) return
    this.activeElement.element.classList.remove('running')
    this.activeElement.element.style.removeProperty('--progress')
    this.activeElement.primary.element.classList.add('mdi-progress-pencil')
    this.activeElement.primary.element.classList.remove('mdi-play')
  }

  #getFavoriteIcon(workflow: ComfyWorkflow) {
    return workflow.isFavorite ? 'star' : 'file-outline'
  }

  #getFavoriteOverIcon(workflow: ComfyWorkflow) {
    return workflow.isFavorite ? 'star-off' : 'star-outline'
  }

  #getFavoriteTooltip(workflow: ComfyWorkflow) {
    return workflow.isFavorite
      ? 'Remove this workflow from your favorites'
      : 'Add this workflow to your favorites'
  }

  #getFavoriteButton(workflow: ComfyWorkflow, primary: boolean) {
    return new ComfyButton({
      icon: this.#getFavoriteIcon(workflow),
      overIcon: this.#getFavoriteOverIcon(workflow),
      iconSize: 18,
      classList:
        'comfyui-button comfyui-workflows-file-action-favorite' +
        (primary ? ' comfyui-workflows-file-action-primary' : ''),
      tooltip: this.#getFavoriteTooltip(workflow),
      action: (e) => {
        e.stopImmediatePropagation()
        workflow.favorite(!workflow.isFavorite)
      }
    })
  }

  #getDeleteButton(workflow: ComfyWorkflow) {
    const deleteButton = new ComfyButton({
      icon: 'delete',
      tooltip: 'Delete this workflow',
      classList: 'comfyui-button comfyui-workflows-file-action',
      iconSize: 18,
      action: async (e, btn) => {
        e.stopImmediatePropagation()

        if (btn.icon === 'delete-empty') {
          btn.enabled = false
          await workflow.delete()
          await this.load()
        } else {
          btn.icon = 'delete-empty'
          btn.element.style.background = 'red'
        }
      }
    })
    deleteButton.element.addEventListener('mouseleave', () => {
      deleteButton.icon = 'delete'
      deleteButton.element.style.removeProperty('background')
    })
    return deleteButton
  }

  #getInsertButton(workflow: ComfyWorkflow) {
    return new ComfyButton({
      icon: 'file-move-outline',
      iconSize: 18,
      tooltip: 'Insert this workflow into the current workflow',
      classList: 'comfyui-button comfyui-workflows-file-action',
      action: (e) => {
        if (!this.app.shiftDown) {
          this.popup.open = false
        }
        e.stopImmediatePropagation()
        if (!this.app.shiftDown) {
          this.popup.open = false
        }
        workflow.insert()
      }
    })
  }

  /** @param {ComfyWorkflow} workflow */
  #getRenameButton(workflow: ComfyWorkflow) {
    return new ComfyButton({
      icon: 'pencil',
      tooltip: workflow.path
        ? 'Rename this workflow'
        : "This workflow can't be renamed as it hasn't been saved.",
      classList: 'comfyui-button comfyui-workflows-file-action',
      iconSize: 18,
      enabled: !!workflow.path,
      action: async (e) => {
        e.stopImmediatePropagation()
        const newName = prompt('Enter new name', workflow.path)
        if (newName) {
          await workflow.rename(newName)
        }
      }
    })
  }

  #getWorkflowElement(workflow: ComfyWorkflow) {
    return new WorkflowElement(this, workflow, {
      primary: this.#getFavoriteButton(workflow, true),
      buttons: [
        this.#getInsertButton(workflow),
        this.#getRenameButton(workflow),
        this.#getDeleteButton(workflow)
      ]
    })
  }

  #createLeafNode(workflow: ComfyWorkflow) {
    const fileNode = this.#getWorkflowElement(workflow)
    this.treeFiles[workflow.path] = fileNode
    return fileNode
  }

  #createNode(currentPath, workflow, i, currentRoot) {
    const part = workflow.pathParts[i]

    const parentNode = $el(
      'ul' + (this.treeState[currentPath] ? '' : '.closed'),
      {
        $: (el) => {
          el.onclick = (e) => {
            this.#expandNode(el, workflow, currentPath, i)
            e.stopImmediatePropagation()
          }
        }
      }
    )
    currentRoot.append(parentNode)

    // Create a node for the current part and an inner UL for its children if it isnt a leaf node
    const leaf = i === workflow.pathParts.length - 1
    let nodeElement
    if (leaf) {
      nodeElement = this.#createLeafNode(workflow).element
    } else {
      nodeElement = $el('li', [
        $el('i.mdi.mdi-18px.mdi-folder'),
        $el('span', part)
      ])
    }
    parentNode.append(nodeElement)
    return parentNode
  }
}

class WorkflowElement<TPrimary extends ComfyComponent = ComfyButton> {
  parent: ComfyWorkflowsContent
  workflow: ComfyWorkflow
  primary: TPrimary
  buttons: ComfyButton[]
  element: HTMLElement
  constructor(
    parent: ComfyWorkflowsContent,
    workflow: ComfyWorkflow,
    {
      tagName = 'li',
      primary,
      buttons
    }: { tagName?: string; primary: TPrimary; buttons: ComfyButton[] }
  ) {
    this.parent = parent
    this.workflow = workflow
    this.primary = primary
    this.buttons = buttons

    this.element = $el(
      tagName + '.comfyui-workflows-tree-file',
      {
        onclick: () => {
          workflow.load()
          this.parent.popup.open = false
        },
        title: this.workflow.path
      },
      [
        this.primary?.element,
        $el('span', workflow.name),
        ...buttons.map((b) => b.element)
      ]
    )
  }
}

type WidgetSelectionDialogOptions = Array<{
  widget: { name: string }
  node: { pos: [number, number]; title: string; id: string; type: string }
}>

class WidgetSelectionDialog extends ComfyAsyncDialog {
  #options: WidgetSelectionDialogOptions

  constructor(options: WidgetSelectionDialogOptions) {
    super()
    this.#options = options
  }

  show(app) {
    this.element.classList.add('comfy-widget-selection-dialog')
    return super.show(
      $el('div', [
        $el('h2', 'Select image target'),
        $el(
          'p',
          "This workflow has multiple image loader nodes, you can rename a node to include 'input' in the title for it to be automatically selected, or select one below."
        ),
        $el(
          'section',
          this.#options.map((opt) => {
            return $el('div.comfy-widget-selection-item', [
              $el(
                'span',
                { dataset: { id: opt.node.id } },
                `${opt.node.title ?? opt.node.type} ${opt.widget.name}`
              ),
              $el(
                'button.comfyui-button',
                {
                  onclick: () => {
                    app.canvas.ds.offset[0] = -opt.node.pos[0] + 50
                    app.canvas.ds.offset[1] = -opt.node.pos[1] + 50
                    app.canvas.selectNode(opt.node)
                    app.graph.setDirtyCanvas(true, true)
                  }
                },
                'Show'
              ),
              $el(
                'button.comfyui-button.primary',
                {
                  onclick: () => {
                    this.close(opt)
                  }
                },
                'Select'
              )
            ])
          })
        )
      ])
    )
  }
}