import { parseFragment as parse, serialize } from "parse5"

type TagDefinitions = Record<string, [string, string[]]>

type Node = {
  nodeName: string
  tagName: string
  parentNode: Node | null
  childNodes: Node[]
  value: string
  attrs: { name: string, value: string }[]
}

// search for <tags> node
function findTagsNode(node: Node): Node | null {
  if (node.nodeName === "tags" && node.childNodes) {
    return node
  }

  if (node.childNodes) {
    for (const child of node.childNodes) {
      const result = findTagsNode(child)
      if (result) return result
    }
  }

  return null
}

// pull out <tags>'s inner text
function tagsNodeText(node: Node): string | null {
  if (node && node.childNodes) {
    const textNode = node.childNodes.find((n: Node) => n.nodeName === "#text")
    if (textNode)
      return textNode.value
  }
  return null
}

// turn what's inside <tags> into a map of tag definitions
// ignores //comment lines
// format is `tag: parent class1 class2`
function parseTagDefinitions(text: string): TagDefinitions {
  const tags: TagDefinitions = {}
  text.split("\n")
    .forEach(line => {
      line = line.trim()
      if (line === "") return
      if (line.startsWith("//")) return

      const [name, definition] = line.split(/:(.+)/)
      if (!name || !definition) {
        console.error(`Invalid tag definition: ${line}`)
        return
      }

      const [parent, ...classes] = definition.trim().split(" ")
      if (!parent) {
        console.error(`Invalid tag definition: ${line}`)
        return
      }

      tags[name.trim()] = [parent.trim(), classes]
    })

  return tags
}

// using our tag definitions, replace the tags in the AST with the real tags w/ classes
function replaceTags(ast: Node, tagDefinitions: TagDefinitions): string {
  const tags = Object.keys(tagDefinitions)

  for (const tag of tags) {
    const [parent, classes] = tagDefinitions[tag] as [string, string[]]
    const parentNodes = findTagsByName(ast, tag)
    for (const parentNode of parentNodes) {
      if (parentNode) {
        parentNode.tagName = parent
        parentNode.nodeName = parent
        if (parentNode.attrs.length === 0) {
          parentNode.attrs = [{ name: "class", value: classes.join(" ") }]
        } else {
          const classAttr = parentNode.attrs.find((a: any) => a.name === "class")
          if (classAttr) {
            classAttr.value += " " + classes.join(" ")
          } else {
            parentNode.attrs.push({ name: "class", value: classes.join(" ") })
          }
        }
      }
    }
  }

  // delete <tags> node
  const tagsIndex = ast.childNodes.findIndex((n: Node) => n.nodeName === "tags")
  if (tagsIndex !== -1) {
    ast.childNodes.splice(tagsIndex, 1)
  }

  return serialize(ast as any)
}

// find tag by name in the AST
function findTagByName(ast: Node, tag: string): Node | null {
  if (ast.nodeName === tag) {
    return ast
  }

  if (ast.childNodes) {
    for (const child of ast.childNodes) {
      const result = findTagByName(child, tag)
      if (result) return result
    }
  }

  return null
}

// find tags by name in the AST
function findTagsByName(ast: Node, tag: string): Node[] {
  const tags: Node[] = []
  if (ast.nodeName === tag) {
    tags.push(ast)
  }
  if (ast.childNodes) {
    for (const child of ast.childNodes) {
      tags.push(...findTagsByName(child, tag))
    }
  }
  return tags
}

// transform string HTML w/ <tags> into new HTML w/ <tags> + friends replaced by new HTML
function transform(content: string): string {
  const ast = parse(content) as Node

  const tagsNode = findTagsNode(ast)
  if (!tagsNode) {
    console.error("No <tags> element found")
    return content
  }

  const tagsText = tagsNodeText(tagsNode)
  if (!tagsText) {
    console.error("No <tags> text found")
    return content
  }

  const tagDefinitions = parseTagDefinitions(tagsText)

  return replaceTags(ast, tagDefinitions)
}

async function main() {
  const content = await Bun.file("example.tsx").text()
  const transformed = transform(content)
  console.log(transformed.trim())
}

if (import.meta.main) {
  main()
}