import type { TextareaRenderable } from "@opentui/core"
import { produce, type SetStoreFunction } from "solid-js/store"
import type { PromptInfo } from "./history"
import type { StoreShape } from "./use-prompt-store"

export type ExtmarkSyncDeps = {
  getInput: () => TextareaRenderable
  setStore: SetStoreFunction<StoreShape>
  fileStyleId: number
  agentStyleId: number
  pasteStyleId: number
  promptPartTypeId: number
}

export function restoreExtmarksFromParts(deps: ExtmarkSyncDeps, parts: PromptInfo["parts"]) {
  const input = deps.getInput()
  input.extmarks.clear()
  deps.setStore("extmarkToPartIndex", new Map())

  parts.forEach((part, partIndex) => {
    let start = 0
    let end = 0
    let virtualText = ""
    let styleId: number | undefined

    if (part.type === "file" && part.source?.text) {
      start = part.source.text.start
      end = part.source.text.end
      virtualText = part.source.text.value
      styleId = deps.fileStyleId
    } else if (part.type === "agent" && part.source) {
      start = part.source.start
      end = part.source.end
      virtualText = part.source.value
      styleId = deps.agentStyleId
    } else if (part.type === "text" && part.source?.text) {
      start = part.source.text.start
      end = part.source.text.end
      virtualText = part.source.text.value
      styleId = deps.pasteStyleId
    }

    if (virtualText) {
      const extmarkId = input.extmarks.create({
        start,
        end,
        virtual: true,
        styleId,
        typeId: deps.promptPartTypeId,
      })
      deps.setStore("extmarkToPartIndex", (map: Map<number, number>) => {
        const newMap = new Map(map)
        newMap.set(extmarkId, partIndex)
        return newMap
      })
    }
  })
}

export function syncExtmarksWithPromptParts(deps: ExtmarkSyncDeps) {
  const input = deps.getInput()
  const allExtmarks = input.extmarks.getAllForTypeId(deps.promptPartTypeId)
  deps.setStore(
    produce((draft) => {
      const newMap = new Map<number, number>()
      const newParts: typeof draft.prompt.parts = []

      for (const extmark of allExtmarks) {
        const partIndex = draft.extmarkToPartIndex.get(extmark.id)
        if (partIndex !== undefined) {
          const part = draft.prompt.parts[partIndex]
          if (part) {
            if (part.type === "agent" && part.source) {
              part.source.start = extmark.start
              part.source.end = extmark.end
            } else if (part.type === "file" && part.source?.text) {
              part.source.text.start = extmark.start
              part.source.text.end = extmark.end
            } else if (part.type === "text" && part.source?.text) {
              part.source.text.start = extmark.start
              part.source.text.end = extmark.end
            }
            newMap.set(extmark.id, newParts.length)
            newParts.push(part)
          }
        }
      }

      draft.extmarkToPartIndex = newMap
      draft.prompt.parts = newParts
    }),
  )
}
