creativetimofficial--ui/apps/www/lib/registry.ts

import { promises as fs } from "fs"
import { tmpdir } from "os"
import path from "path"
import { registryItemFileSchema, registryItemSchema } from "shadcn/schema"
import { Project, ScriptKind } from "ts-morph"
import { z } from "zod"

import { Index } from "@/registry/__index__"

export function getRegistryComponent(name: string) {
  return Index[name]?.component
}

export async function getRegistryItem(name: string) {
  const item = Index[name]

  if (!item) {
    return null
  }

  // Convert all file paths to object.
  // TODO: remove when we migrate to new registry.
  item.files = item.files.map((file: unknown) =>
    typeof file === "string" ? { path: file } : file
  )

  // Fail early before doing expensive file operations.
  const result = registryItemSchema.safeParse(item)
  if (!result.success) {
    return null
  }

  // Build path mappings from all files for import rewriting
  const pathMappings = buildPathMappings(item.files)

  let files: typeof result.data.files = []
  for (const file of item.files) {
    const content = await getFileContent(file, pathMappings)
    const relativePath = path.relative(process.cwd(), file.path)

    files.push({
      ...file,
      path: relativePath,
      content,
    })
  }

  // Fix file paths.
  files = fixFilePaths(files)

  const parsed = registryItemSchema.safeParse({
    ...result.data,
    files,
  })

  if (!parsed.success) {
    console.error(parsed.error.message)
    return null
  }

  return parsed.data
}

function buildPathMappings(
  files: Array<{ path?: string; target?: string }>
): Map<string, string> {
  const mappings = new Map<string, string>()

  files.forEach((file) => {
    if (file.path && file.target) {
      // Extract the source path relative to registry/creative-tim/
      const match = file.path.match(/registry\/creative-tim\/(.+)$/)
      if (match) {
        mappings.set(match[1], file.target)
      }
    }
  })

  return mappings
}

async function getFileContent(
  file: z.infer<typeof registryItemFileSchema>,
  pathMappings: Map<string, string>
) {
  const raw = await fs.readFile(file.path, "utf-8")

  const project = new Project({
    compilerOptions: {},
  })

  const tempFile = await createTempSourceFile(file.path)
  const sourceFile = project.createSourceFile(tempFile, raw, {
    scriptKind: ScriptKind.TSX,
  })

  // Remove meta variables.
  // removeVariable(sourceFile, "iframeHeight")
  // removeVariable(sourceFile, "containerClassName")
  // removeVariable(sourceFile, "description")

  let code = sourceFile.getFullText()

  // Some registry items uses default export.
  // We want to use named export instead.
  // TODO: do we really need this? - @shadcn.
  if (file.type !== "registry:page") {
    code = code.replaceAll("export default", "export")
  }

  // Fix imports using path mappings.
  code = fixImport(code, pathMappings)

  return code
}

function getFileTarget(file: z.infer<typeof registryItemFileSchema>) {
  let target = file.target

  if (!target || target === "") {
    const fileName = file.path.split("/").pop()
    if (
      file.type === "registry:block" ||
      file.type === "registry:component" ||
      file.type === "registry:example"
    ) {
      target = `components/${fileName}`
    }

    if (file.type === "registry:ui") {
      target = `components/ui/${fileName}`
    }

    if (file.type === "registry:hook") {
      target = `hooks/${fileName}`
    }

    if (file.type === "registry:lib") {
      target = `lib/${fileName}`
    }
  }

  return target ?? ""
}

async function createTempSourceFile(filename: string) {
  const dir = await fs.mkdtemp(path.join(tmpdir(), "shadcn-"))
  return path.join(dir, filename)
}

function fixFilePaths(files: z.infer<typeof registryItemSchema>["files"]) {
  if (!files) {
    return []
  }

  // Resolve all paths relative to the first file's directory.
  const firstFilePath = files[0].path
  const firstFilePathDir = path.dirname(firstFilePath)

  return files.map((file) => {
    return {
      ...file,
      path: path.relative(firstFilePathDir, file.path),
      target: getFileTarget(file),
    }
  })
}

export function fixImport(content: string, pathMappings: Map<string, string>) {
  // Convert registry imports to their target paths using the provided mappings
  // This ensures imports match the actual target locations

  // Sort mappings by path length (longest first) to handle more specific paths first
  const sortedMappings = Array.from(pathMappings.entries()).sort(
    ([a], [b]) => b.length - a.length
  )

  // Handle block-specific imports using the exact path mappings
  sortedMappings.forEach(([sourcePath, targetPath]) => {
    if (sourcePath.startsWith("blocks/")) {
      // Remove the file extension from sourcePath to match import statements
      const sourcePathNoExt = sourcePath.replace(/\.(tsx?|jsx?)$/, "")

      // Remove the file extension from targetPath
      const targetPathNoExt = targetPath.replace(/\.(tsx?|jsx?)$/, "")

      // Create regex to match this specific import path
      const escapedPath = sourcePathNoExt.replace(/\//g, "\\/")
      const importRegex = new RegExp(
        `@\\/registry\\/creative-tim\\/${escapedPath}`,
        "g"
      )

      content = content.replace(importRegex, `@/${targetPathNoExt}`)
    }
  })

  // Fix UI component imports
  content = content.replaceAll(
    "@/components/ui/",
    "@/components/ui/"
  )

  // Fix example imports
  content = content.replaceAll(
    "@/registry/creative-tim/examples/",
    "@/components/examples/"
  )

  // Fix hook imports
  content = content.replaceAll("@/registry/creative-tim/hooks/", "@/hooks/")

  // Fix lib imports
  content = content.replaceAll("@/registry/creative-tim/lib/", "@/lib/")

  return content
}

export type FileTree = {
  name: string
  path?: string
  children?: FileTree[]
}

export function createFileTreeForRegistryItemFiles(
  files: Array<{ path: string; target?: string }>
) {
  const root: FileTree[] = []

  for (const file of files) {
    const path = file.target ?? file.path
    const parts = path.split("/")
    let currentLevel = root

    for (let i = 0; i < parts.length; i++) {
      const part = parts[i]
      const isFile = i === parts.length - 1
      const existingNode = currentLevel.find((node) => node.name === part)

      if (existingNode) {
        if (isFile) {
          // Update existing file node with full path
          existingNode.path = path
        } else {
          // Move to next level in the tree
          currentLevel = existingNode.children!
        }
      } else {
        const newNode: FileTree = isFile
          ? { name: part, path }
          : { name: part, children: [] }

        currentLevel.push(newNode)

        if (!isFile) {
          currentLevel = newNode.children!
        }
      }
    }
  }

  return root
}