import type { ComfyApp } from './app'
import { api } from './api'
import { ChangeTracker } from './changeTracker'
import { ComfyAsyncDialog } from './ui/components/asyncDialog'
import { getStorageValue, setStorageValue } from './utils'
import { LGraphCanvas, LGraph } from '@comfyorg/litegraph'

function appendJsonExt(path: string) {
  if (!path.toLowerCase().endsWith('.json')) {
    path += '.json'
  }
  return path
}

export function trimJsonExt(path: string) {
  return path?.replace(/\.json$/, '')
}

export class ComfyWorkflowManager extends EventTarget {
  #activePromptId: string | null = null
  #unsavedCount = 0
  #activeWorkflow: ComfyWorkflow

  workflowLookup: Record<string, ComfyWorkflow> = {}
  workflows: Array<ComfyWorkflow> = []
  openWorkflows: Array<ComfyWorkflow> = []
  queuedPrompts: Record<
    string,
    { workflow?: ComfyWorkflow; nodes?: Record<string, boolean> }
  > = {}
  app: ComfyApp

  get activeWorkflow() {
    return this.#activeWorkflow ?? this.openWorkflows[0]
  }

  get activePromptId() {
    return this.#activePromptId
  }

  get activePrompt() {
    return this.queuedPrompts[this.#activePromptId]
  }

  constructor(app: ComfyApp) {
    super()
    this.app = app
    ChangeTracker.init(app)

    this.#bindExecutionEvents()
  }

  #bindExecutionEvents() {
    // TODO: on reload, set active prompt based on the latest ws message

    const emit = () =>
      this.dispatchEvent(
        new CustomEvent('execute', { detail: this.activePrompt })
      )
    let executing = null
    api.addEventListener('execution_start', (e) => {
      this.#activePromptId = e.detail.prompt_id

      // This event can fire before the event is stored, so put a placeholder
      this.queuedPrompts[this.#activePromptId] ??= { nodes: {} }
      emit()
    })
    api.addEventListener('execution_cached', (e) => {
      if (!this.activePrompt) return
      for (const n of e.detail.nodes) {
        this.activePrompt.nodes[n] = true
      }
      emit()
    })
    api.addEventListener('executed', (e) => {
      if (!this.activePrompt) return
      this.activePrompt.nodes[e.detail.node] = true
      emit()
    })
    api.addEventListener('executing', (e) => {
      if (!this.activePrompt) return

      if (executing) {
        // Seems sometimes nodes that are cached fire executing but not executed
        this.activePrompt.nodes[executing] = true
      }
      executing = e.detail
      if (!executing) {
        delete this.queuedPrompts[this.#activePromptId]
        this.#activePromptId = null
      }
      emit()
    })
  }

  async loadWorkflows() {
    try {
      let favorites
      const resp = await api.getUserData('workflows/.index.json')
      let info
      if (resp.status === 200) {
        info = await resp.json()
        favorites = new Set(info?.favorites ?? [])
      } else {
        favorites = new Set()
      }

      const workflows = (await api.listUserData('workflows', true, true)).map(
        (w) => {
          let workflow = this.workflowLookup[w[0]]
          if (!workflow) {
            workflow = new ComfyWorkflow(
              this,
              w[0],
              w.slice(1),
              favorites.has(w[0])
            )
            this.workflowLookup[workflow.path] = workflow
          }
          return workflow
        }
      )

      this.workflows = workflows
    } catch (error) {
      alert('Error loading workflows: ' + (error.message ?? error))
      this.workflows = []
    }
  }

  async saveWorkflowMetadata() {
    await api.storeUserData('workflows/.index.json', {
      favorites: [
        ...this.workflows.filter((w) => w.isFavorite).map((w) => w.path)
      ]
    })
  }

  /**
   * @param {string | ComfyWorkflow | null} workflow
   */
  setWorkflow(workflow) {
    if (workflow && typeof workflow === 'string') {
      // Selected by path, i.e. on reload of last workflow
      const found = this.workflows.find((w) => w.path === workflow)
      if (found) {
        workflow = found
        workflow.unsaved =
          !workflow ||
          getStorageValue('Comfy.PreviousWorkflowUnsaved') === 'true'
      }
    }

    if (!(workflow instanceof ComfyWorkflow)) {
      // Still not found, either reloading a deleted workflow or blank
      workflow = new ComfyWorkflow(
        this,
        workflow ||
          'Unsaved Workflow' +
            (this.#unsavedCount++ ? ` (${this.#unsavedCount})` : '')
      )
    }

    const index = this.openWorkflows.indexOf(workflow)
    if (index === -1) {
      // Opening a new workflow
      this.openWorkflows.push(workflow)
    }

    this.#activeWorkflow = workflow

    setStorageValue('Comfy.PreviousWorkflow', this.activeWorkflow.path ?? '')
    this.dispatchEvent(new CustomEvent('changeWorkflow'))
  }

  storePrompt({ nodes, id }) {
    this.queuedPrompts[id] ??= {}
    this.queuedPrompts[id].nodes = {
      ...nodes.reduce((p, n) => {
        p[n] = false
        return p
      }, {}),
      ...this.queuedPrompts[id].nodes
    }
    this.queuedPrompts[id].workflow = this.activeWorkflow
  }

  /**
   * @param {ComfyWorkflow} workflow
   */
  async closeWorkflow(workflow, warnIfUnsaved = true) {
    if (!workflow.isOpen) {
      return true
    }
    if (workflow.unsaved && warnIfUnsaved) {
      const res = await ComfyAsyncDialog.prompt({
        title: 'Save Changes?',
        message: `Do you want to save changes to "${workflow.path ?? workflow.name}" before closing?`,
        actions: ['Yes', 'No', 'Cancel']
      })
      if (res === 'Yes') {
        const active = this.activeWorkflow
        if (active !== workflow) {
          // We need to switch to the workflow to save it
          await workflow.load()
        }

        if (!(await workflow.save())) {
          // Save was canceled, restore the previous workflow
          if (active !== workflow) {
            await active.load()
          }
          return
        }
      } else if (res === 'Cancel') {
        return
      }
    }
    workflow.changeTracker = null
    this.openWorkflows.splice(this.openWorkflows.indexOf(workflow), 1)
    if (this.openWorkflows.length) {
      this.#activeWorkflow = this.openWorkflows[0]
      await this.#activeWorkflow.load()
    } else {
      // Load default
      await this.app.loadGraphData()
    }
  }
}

export class ComfyWorkflow {
  #name
  #path
  #pathParts
  #isFavorite = false
  changeTracker: ChangeTracker | null = null
  unsaved = false
  manager: ComfyWorkflowManager

  get name() {
    return this.#name
  }

  get path() {
    return this.#path
  }

  get pathParts() {
    return this.#pathParts
  }

  get isFavorite() {
    return this.#isFavorite
  }

  get isOpen() {
    return !!this.changeTracker
  }

  constructor(
    manager: ComfyWorkflowManager,
    path: string,
    pathParts?: string[],
    isFavorite?: boolean
  ) {
    this.manager = manager
    if (pathParts) {
      this.#updatePath(path, pathParts)
      this.#isFavorite = isFavorite
    } else {
      this.#name = path
      this.unsaved = true
    }
  }

  #updatePath(path: string, pathParts: string[]) {
    this.#path = path

    if (!pathParts) {
      if (!path.includes('\\')) {
        pathParts = path.split('/')
      } else {
        pathParts = path.split('\\')
      }
    }

    this.#pathParts = pathParts
    this.#name = trimJsonExt(pathParts[pathParts.length - 1])
  }

  async getWorkflowData() {
    const resp = await api.getUserData('workflows/' + this.path)
    if (resp.status !== 200) {
      alert(
        `Error loading workflow file '${this.path}': ${resp.status} ${resp.statusText}`
      )
      return
    }
    return await resp.json()
  }

  load = async () => {
    if (this.isOpen) {
      await this.manager.app.loadGraphData(
        this.changeTracker.activeState,
        true,
        true,
        this
      )
    } else {
      const data = await this.getWorkflowData()
      if (!data) return
      await this.manager.app.loadGraphData(data, true, true, this)
    }
  }

  async save(saveAs = false) {
    if (!this.path || saveAs) {
      return !!(await this.#save(null, false))
    } else {
      return !!(await this.#save(this.path, true))
    }
  }

  async favorite(value: boolean) {
    try {
      if (this.#isFavorite === value) return
      this.#isFavorite = value
      await this.manager.saveWorkflowMetadata()
      this.manager.dispatchEvent(new CustomEvent('favorite', { detail: this }))
    } catch (error) {
      alert(
        'Error favoriting workflow ' +
          this.path +
          '\n' +
          (error.message ?? error)
      )
    }
  }

  async rename(path: string) {
    path = appendJsonExt(path)
    let resp = await api.moveUserData(
      'workflows/' + this.path,
      'workflows/' + path
    )

    if (resp.status === 409) {
      if (
        !confirm(
          `Workflow '${path}' already exists, do you want to overwrite it?`
        )
      )
        return resp
      resp = await api.moveUserData(
        'workflows/' + this.path,
        'workflows/' + path,
        { overwrite: true }
      )
    }

    if (resp.status !== 200) {
      alert(
        `Error renaming workflow file '${this.path}': ${resp.status} ${resp.statusText}`
      )
      return
    }

    const isFav = this.isFavorite
    if (isFav) {
      await this.favorite(false)
    }
    path = (await resp.json()).substring('workflows/'.length)
    this.#updatePath(path, null)
    if (isFav) {
      await this.favorite(true)
    }
    this.manager.dispatchEvent(new CustomEvent('rename', { detail: this }))
    setStorageValue('Comfy.PreviousWorkflow', this.path ?? '')
  }

  async insert() {
    const data = await this.getWorkflowData()
    if (!data) return

    const old = localStorage.getItem('litegrapheditor_clipboard')
    // @ts-ignore
    const graph = new LGraph(data)
    const canvas = new LGraphCanvas(null, graph, {
      // @ts-ignore
      skip_events: true,
      skip_render: true
    })
    canvas.selectNodes()
    canvas.copyToClipboard()
    this.manager.app.canvas.pasteFromClipboard()
    localStorage.setItem('litegrapheditor_clipboard', old)
  }

  async delete() {
    // TODO: fix delete of current workflow - should mark workflow as unsaved and when saving use old name by default

    try {
      if (this.isFavorite) {
        await this.favorite(false)
      }
      await api.deleteUserData('workflows/' + this.path)
      this.unsaved = true
      this.#path = null
      this.#pathParts = null
      this.manager.workflows.splice(this.manager.workflows.indexOf(this), 1)
      this.manager.dispatchEvent(new CustomEvent('delete', { detail: this }))
    } catch (error) {
      alert(`Error deleting workflow: ${error.message || error}`)
    }
  }

  track() {
    if (this.changeTracker) {
      this.changeTracker.restore()
    } else {
      this.changeTracker = new ChangeTracker(this)
    }
  }

  async #save(path: string | null, overwrite: boolean) {
    if (!path) {
      path = prompt(
        'Save workflow as:',
        trimJsonExt(this.path) ?? this.name ?? 'workflow'
      )
      if (!path) return
    }

    path = appendJsonExt(path)

    const p = await this.manager.app.graphToPrompt()
    const json = JSON.stringify(p.workflow, null, 2)
    let resp = await api.storeUserData('workflows/' + path, json, {
      stringify: false,
      throwOnError: false,
      overwrite
    })
    if (resp.status === 409) {
      if (
        !confirm(
          `Workflow '${path}' already exists, do you want to overwrite it?`
        )
      )
        return
      resp = await api.storeUserData('workflows/' + path, json, {
        stringify: false
      })
    }

    if (resp.status !== 200) {
      alert(
        `Error saving workflow '${this.path}': ${resp.status} ${resp.statusText}`
      )
      return
    }

    path = (await resp.json()).substring('workflows/'.length)

    if (!this.path) {
      // Saved new workflow, patch this instance
      this.#updatePath(path, null)
      await this.manager.loadWorkflows()
      this.unsaved = false
      this.manager.dispatchEvent(new CustomEvent('rename', { detail: this }))
      setStorageValue('Comfy.PreviousWorkflow', this.path ?? '')
    } else if (path !== this.path) {
      // Saved as, open the new copy
      await this.manager.loadWorkflows()
      const workflow = this.manager.workflowLookup[path]
      await workflow.load()
    } else {
      // Normal save
      this.unsaved = false
      this.manager.dispatchEvent(new CustomEvent('save', { detail: this }))
    }

    return true
  }
}