// `.swift` handling for a vite dev server: compile the project's swift sources
// with the swift.org toolchain, serve the artifact at a content-hashed url,
// and answer every `.swift` import with the module that mounts it.
//
// two hosts use this. the bundler that builds the app's native bundle runs
// resolveId/load, and the dev server runs the middleware that serves the
// artifact. both share the build below, so the artifact is compiled once.
//
// the hash is the hash of the built wasm, so a new hash means a genuinely
// different app: createSwiftEntry moves the @State store across it instead of
// remounting.

import { execFile } from 'node:child_process'
import { createHash } from 'node:crypto'
import fs from 'node:fs'
import os from 'node:os'
import path from 'node:path'
import { promisify } from 'node:util'
import {
  enumerateBundleImages,
  enumerateBundleJson,
} from '../../sootsim-swift/src/bundleAssets'
import { swiftEntryModuleEsm } from '../../sootsim-swift/src/entryModule'
import toolchain from '../../sootsim-swift/toolchain.json'
import type { Plugin, ViteDevServer } from 'vite'

const run = promisify(execFile)

// the swift.org release scripts/build.sh builds with, from the one place it is
// named: the toolchain is <release>.xctoolchain and its matching wasm sdk is
// <release>_wasm. xcode's swift has no wasm backend, so this toolchain is the
// only one that can build the target.
const TOOLCHAIN_BIN = path.join(
  os.homedir(),
  `Library/Developer/Toolchains/${toolchain.release}.xctoolchain/usr/bin`,
)
const WASM_SDK = `${toolchain.release}_wasm`
const WASM_TRIPLE = 'wasm32-unknown-wasip1'

const ARTIFACT_PATH = '/__rnx/swift/'
const RESOURCE_PATH = '/__rnx/swift-res/'
const VIRTUAL_PREFIX = '\0rnx-swift:'

interface SwiftArtifact {
  hash: string
  bytes: Buffer
  // the sources this artifact was built from. a watcher event and a `load`
  // reach the project in either order, so freshness is read from the sources
  // rather than from whichever of them ran first.
  stamp: string
}

interface SwiftProject {
  packagePath: string
  artifact: SwiftArtifact | null
  // the build in flight and the sources it was started from
  building: { stamp: string; promise: Promise<SwiftArtifact> } | null
  // content hash to absolute path for this load's bundle images: the served
  // urls name a hash, never a path, so there is nothing to traverse.
  resourceFiles: Map<string, string>
}

// a project is its swiftpm package: the build, the artifact, and the dev
// server that serves it are shared by every plugin instance in this process.
const projects = new Map<string, SwiftProject>()
let devServer: ViteDevServer | undefined

function projectFor(packagePath: string): SwiftProject {
  const existing = projects.get(packagePath)
  if (existing) return existing
  const created: SwiftProject = {
    packagePath,
    artifact: null,
    building: null,
    resourceFiles: new Map(),
  }
  projects.set(packagePath, created)
  return created
}

// the compile unit is the whole package: `swift build` compiles every source
// in it, so one file's edit can change the artifact for all of them. also
// used for a linked module below, whose own nested packages are separate
// units by the same rule.
function sourceFiles(root: string): string[] {
  const found: string[] = []
  const walk = (dir: string): void => {
    for (const entry of fs.readdirSync(dir, { withFileTypes: true })) {
      if (entry.name.startsWith('.') || entry.name === 'node_modules') continue
      const full = path.join(dir, entry.name)
      // a nested Package.swift is a separate package with its own compile:
      // its sources belong to that unit's stamp, not this one's
      if (entry.isDirectory()) {
        if (full !== root && fs.existsSync(path.join(full, 'Package.swift'))) continue
        walk(full)
      } else if (entry.name.endsWith('.swift')) found.push(full)
    }
  }
  walk(root)
  return found
}

// every file that can carry bundle content: catalogs and flat Resources
// rasters anywhere under the package. `.build` is skipped where sourceFiles
// does not need to: a release tree holds tens of thousands of files and no
// bundle content, and walking it per import would tax every hot reload.
function resourceCandidateFiles(root: string): string[] {
  const found: string[] = []
  const walk = (dir: string): void => {
    for (const entry of fs.readdirSync(dir, { withFileTypes: true })) {
      if (
        entry.name.startsWith('.') ||
        entry.name === 'node_modules' ||
        entry.name === '.build'
      )
        continue
      const full = path.join(dir, entry.name)
      if (entry.isDirectory()) {
        if (full !== root && fs.existsSync(path.join(full, 'Package.swift'))) continue
        walk(full)
      } else {
        found.push(full)
      }
    }
  }
  walk(root)
  return found
}

const RESOURCE_CONTENT_TYPES: Record<string, string> = {
  '.avif': 'image/avif',
  '.bmp': 'image/bmp',
  '.gif': 'image/gif',
  '.ico': 'image/x-icon',
  '.jpeg': 'image/jpeg',
  '.jpg': 'image/jpeg',
  '.png': 'image/png',
  '.webp': 'image/webp',
}

export interface SwiftResourceTable {
  /** bundle asset name to absolute serving url. */
  table: Record<string, string>
  /** project-relative image files this load's table covers (for watching). */
  watched: string[]
}

// the producer for the local vehicle: enumerate the package's bundle images
// through the shared rule both hosts use, hash each raster's bytes into its
// url (an edited photo moves its url, so no stale paint survives a reload),
// and register hash to path for the serving middleware below.
export function buildSwiftResourceTable(
  project: SwiftProject,
  baseUrl: string,
): SwiftResourceTable {
  const absolute = resourceCandidateFiles(project.packagePath)
  const relative = absolute.map((file) =>
    path.relative(project.packagePath, file).split(path.sep).join(path.posix.sep),
  )
  const byRelative = new Map(absolute.map((file, index) => [relative[index], file]))
  const found = enumerateBundleImages(relative, (candidate) => {
    const file = byRelative.get(candidate)
    if (!file || path.basename(file) !== 'Contents.json') return null
    try {
      return fs.readFileSync(file, 'utf8')
    } catch {
      return null
    }
  })
  const table: Record<string, string> = {}
  const watched: string[] = []
  const registered = new Map<string, string>()
  for (const { name, file } of found) {
    const disk = byRelative.get(file)
    if (!disk) continue
    let bytes: Buffer
    try {
      bytes = fs.readFileSync(disk)
    } catch {
      continue
    }
    const hash = createHash('sha256').update(bytes).digest('hex').slice(0, 16)
    const ext = path.posix.extname(file).toLowerCase()
    table[name] = `${baseUrl}${RESOURCE_PATH}${hash}${ext}`
    watched.push(file)
    // a Contents.json edit (scale swap, file rename) re-resolves the same
    // way a raster edit does, so the catalog rides the watcher too.
    const contents = file.replace(/\/[^/]+$/, '/Contents.json')
    if (byRelative.has(contents)) watched.push(contents)
    registered.set(`${hash}${ext}`, disk)
  }
  project.resourceFiles = registered
  return { table, watched }
}

// the /res pack: every Resources/ json in the package, basename-keyed and
// base64, riding the entry module so the mount serves it without a fetch.
// the namespace is flat like the bundle the wasm side reads: shallower
// paths pack first (the Resources root is the bundle root) with a lexical
// tiebreak, and the first spelling of a basename wins a collision. caps are
// INFERRED (no app in either rung is near them; landmarkData-scale json is
// tens of KB): a file past them is skipped and named in `skipped` so the
// load warns instead of silently serving a partial bundle. binary bytes
// travel the same map when a port reaches past json; the worker half packs
// text only until then (see the gaps ledger).
const BUNDLE_MAX_FILES = 128
const BUNDLE_MAX_FILE_BYTES = 2 * 1024 * 1024
const BUNDLE_MAX_TOTAL_BYTES = 8 * 1024 * 1024

export interface SwiftBundlePack {
  /** basename to base64 bytes for the mount's /res. */
  files: Record<string, string>
  /** project-relative json files this load's pack covers (for watching). */
  watched: string[]
  /** project-relative files skipped by a cap, for the load warning. */
  skipped: string[]
}

export function buildSwiftBundleFiles(packagePath: string): SwiftBundlePack {
  // resourceCandidateFiles walks every non-ignored file with the nested
  // package exclusion; the shared enumeration decides namespace, order and
  // collisions, identically for every host that packs.
  const relative = resourceCandidateFiles(packagePath).map((file) =>
    path.relative(packagePath, file).split(path.sep).join(path.posix.sep),
  )
  const files: Record<string, string> = {}
  const watched: string[] = []
  const skipped: string[] = []
  let total = 0
  for (const { base, file } of enumerateBundleJson(relative)) {
    watched.push(file)
    let bytes: Buffer
    try {
      bytes = fs.readFileSync(path.join(packagePath, ...file.split('/')))
    } catch {
      continue
    }
    if (bytes.length > BUNDLE_MAX_FILE_BYTES) {
      skipped.push(`${file} (${bytes.length} bytes past the per-file cap)`)
      continue
    }
    if (
      Object.keys(files).length >= BUNDLE_MAX_FILES ||
      total + bytes.length > BUNDLE_MAX_TOTAL_BYTES
    ) {
      skipped.push(
        `${file} (pack caps: ${BUNDLE_MAX_FILES} files / ${BUNDLE_MAX_TOTAL_BYTES} bytes)`,
      )
      continue
    }
    files[base] = bytes.toString('base64')
    total += bytes.length
  }
  return { files, watched, skipped }
}

// the middleware's lookup, exported for the producer's own proof: a hit is
// known bytes for a known hash, anything else is null, never a file.
export function serveSwiftResource(
  project: SwiftProject,
  key: string,
): { bytes: Buffer; contentType: string } | null {
  const disk = project.resourceFiles.get(key)
  if (!disk) return null
  try {
    return {
      bytes: fs.readFileSync(disk),
      contentType:
        RESOURCE_CONTENT_TYPES[path.extname(disk).toLowerCase()] ??
        'application/octet-stream',
    }
  } catch {
    return null
  }
}

// the SwiftUI module the project links: its manifest names it as a path
// dependency, and `swift build` compiles it with the project, so an edit
// there changes the artifact without moving any of the project's own mtimes.
export function linkedModuleRoots(packagePath: string): string[] {
  let manifest: string
  try {
    manifest = fs.readFileSync(path.join(packagePath, 'Package.swift'), 'utf8')
  } catch {
    return []
  }
  const roots = new Set<string>()
  const dependency = /\.package\(\s*path:\s*"([^"]+)"\s*\)/g
  for (const match of manifest.replace(/^\s*\/\/.*$/gm, '').matchAll(dependency)) {
    const root = path.resolve(packagePath, match[1])
    if (fs.existsSync(path.join(root, 'Package.swift'))) roots.add(root)
  }
  return [...roots].sort()
}

// the stamp covers the project plus the modules it links: a running dev
// server rebuilds only when this moves, so leaving the linked sources out
// would keep serving the artifact built before a SwiftUI module edit for as
// long as the project's own sources stayed the same. it names the same
// closure the compile service's artifact hash names, but it is not the same
// function: the stamp hashes mtimes, which is the cheap rebuild signal for a
// watcher, while the artifact address hashes content, which is what stays
// stable when a checkout restores byte-identical files. a restore rebuilds
// here and reuses the stored artifact there, by design.
export function sourceStamp(packagePath: string): string {
  const hash = createHash('sha256')
  for (const root of [packagePath, ...linkedModuleRoots(packagePath)]) {
    hash.update(root)
    for (const file of sourceFiles(root).sort()) {
      hash.update(path.relative(root, file))
      hash.update(String(fs.statSync(file).mtimeMs))
    }
  }
  return hash.digest('hex').slice(0, 16)
}

// the registry slot for an entry: the owning package's own name plus the
// source within it. two packages in one app can hold the same filename, and
// a bare relative path would share one slot, one instance, and its state.
// the name travels instead of the path so identical source in two checkouts
// emits the identical module; two same-named packages sharing one app still
// collide, and take the rootId override for that.
export function entryRootId(packagePath: string, file: string): string {
  const relative = path.relative(packagePath, file).split(path.sep).join(path.posix.sep)
  return `${path.basename(packagePath)}${path.posix.sep}${relative}`
}

// the app fetches from a worker whose origin is the shell, so urls have to
// name this dev server rather than resolve relative to the shell.
function serverOrigin(): string {
  const address = devServer?.httpServer?.address()
  const port =
    address && typeof address === 'object' ? address.port : devServer?.config.server.port
  if (!port) throw new Error('the swift urls need the dev server port')
  return `http://localhost:${port}`
}

function artifactUrl(project: SwiftProject, artifact: SwiftArtifact): string {
  return `${serverOrigin()}${ARTIFACT_PATH}${artifact.hash}.wasm`
}

// one executable is the app. a package with several is ambiguous: the host
// cannot know which one the page is, so it says so instead of guessing.
function builtArtifact(packagePath: string): string {
  const binDir = path.join(packagePath, '.build', WASM_TRIPLE, 'release')
  const artifacts = fs.existsSync(binDir)
    ? fs
        .readdirSync(binDir)
        .filter((name) => name.endsWith('.wasm'))
        .sort()
    : []
  if (artifacts.length !== 1) {
    throw new Error(
      `swift project ${packagePath} built ${artifacts.length} wasm artifacts; exactly one executable is the app`,
    )
  }
  return path.join(binDir, artifacts[0])
}

function streamText(value: unknown): string {
  if (typeof value === 'string') return value.trim()
  if (Buffer.isBuffer(value)) return value.toString('utf8').trim()
  return ''
}

async function compile(packagePath: string, stamp: string): Promise<SwiftArtifact> {
  const swift = path.join(TOOLCHAIN_BIN, 'swift')
  if (!fs.existsSync(swift)) {
    throw new Error(
      `no swift.org toolchain at ${TOOLCHAIN_BIN}. install it and the ${WASM_SDK} sdk; see packages/sootsim-swift/README.md`,
    )
  }
  try {
    await run(
      swift,
      [
        'build',
        '--package-path',
        packagePath,
        '--swift-sdk',
        WASM_SDK,
        '-c',
        'release',
        '-Xswiftc',
        '-Osize',
      ],
      { maxBuffer: 64 * 1024 * 1024 },
    )
  } catch (error) {
    // `swift build` prints the compiler's diagnostics on stdout and swiftpm's
    // own on stderr; the failure's message holds only the command line, so
    // both streams are what names the line that broke. stdout leads because
    // compiler diagnostics are what the editor needs to show.
    const streams =
      error && typeof error === 'object'
        ? [
            streamText(Reflect.get(error, 'stdout')),
            streamText(Reflect.get(error, 'stderr')),
          ].filter(Boolean)
        : []
    const detail =
      streams.length > 0
        ? streams.join('\n')
        : error instanceof Error
          ? error.message
          : String(error)
    throw new Error(`swift build failed for ${packagePath}\n${detail}`)
  }
  const bytes = fs.readFileSync(builtArtifact(packagePath))
  return {
    hash: createHash('sha256').update(bytes).digest('hex').slice(0, 16),
    bytes,
    stamp,
  }
}

function artifactFor(project: SwiftProject): Promise<SwiftArtifact> {
  const stamp = sourceStamp(project.packagePath)
  if (project.artifact?.stamp === stamp) return Promise.resolve(project.artifact)
  if (project.building) {
    // a second edit during a build: the build in flight was started from
    // older sources, so its artifact is not the answer for these. wait it out
    // and ask again, which starts the build these sources need; joining it
    // would emit the older artifact and nothing would trigger another load.
    const running = project.building
    if (running.stamp === stamp) return running.promise
    const again = () => artifactFor(project)
    return running.promise.then(again, again)
  }
  const promise = compile(project.packagePath, stamp).then(
    (artifact) => {
      project.artifact = artifact
      project.building = null
      return artifact
    },
    (error: unknown) => {
      project.building = null
      throw error
    },
  )
  project.building = { stamp, promise }
  return promise
}

// the swiftpm package that owns an imported file: the nearest directory at
// or above the file that holds a Package.swift, stopping at the vite root
// once the dev server has named it. a swift-only project keeps its sources at
// the root so that is its package; a React Native app keeps its in a
// subdirectory (a `native/` package beside the js), and each such package
// compiles on its own.
function packageFor(file: string): string {
  let dir = path.dirname(file)
  const stop = devServer?.config.root
  for (;;) {
    if (fs.existsSync(path.join(dir, 'Package.swift'))) return dir
    if ((stop !== undefined && dir === stop) || path.dirname(dir) === dir) break
    dir = path.dirname(dir)
  }
  throw new Error(
    `no Package.swift above ${file}${stop === undefined ? '' : ` (searched up to ${stop})`}; a .swift import needs a swiftpm package like the one in packages/sootsim-swift/example`,
  )
}

export interface SwiftPluginOptions {
  // swiftpm package that owns the project's .swift sources. defaults to the
  // nearest Package.swift above each imported file, which is the vite root
  // for the layout a swift-only project uses.
  packagePath?: string
}

export function swiftPlugin(options: SwiftPluginOptions = {}): Plugin {
  return {
    name: 'rnx-swift',
    enforce: 'pre',

    async resolveId(source, importer) {
      if (source.startsWith(VIRTUAL_PREFIX)) return source
      // import-analysis re-resolves the served URL form of the id
      // (/@id/__x00__..., with the null byte encoded) while transforming the
      // importer. it still names this module, so answer the decoded id; any
      // other /@id/ url is already resolved and must not fall through to the
      // filesystem resolve below, which would throw for it.
      if (source.startsWith('/@id/')) {
        const decoded = source.slice('/@id/'.length).replace(/^__x00__/, '\0')
        return decoded.startsWith(VIRTUAL_PREFIX) ? decoded : null
      }
      if (!source.endsWith('.swift')) return null
      const resolved = await this.resolve(source, importer, { skipSelf: true })
      if (!resolved) throw new Error(`cannot resolve swift source ${source}`)
      return VIRTUAL_PREFIX + resolved.id
    },

    async load(id) {
      if (!id.startsWith(VIRTUAL_PREFIX)) return null
      const file = id.slice(VIRTUAL_PREFIX.length)
      // the module is the whole package's artifact, so an import of a file
      // that was deleted would keep answering with it. a deleted .tsx fails
      // its importer; so does a deleted .swift.
      if (!fs.existsSync(file)) throw new Error(`swift source ${file} no longer exists`)
      const packagePath = options.packagePath ?? packageFor(file)
      const project = projectFor(packagePath)
      // every source in the package is a dependency of this module, so a
      // bundler watching this module learns about a `.swift` edit and re-runs
      // load. that is the whole hot reload: the new hash in the emitted module
      // is what tells the mount it is looking at a different app.
      for (const source of sourceFiles(packagePath)) this.addWatchFile(source)
      // an edit to the linked SwiftUI module rebuilds the artifact without
      // touching the project's own sources, so those files gate the reload
      // too: the new hash in the emitted module is what tells the mount it
      // is looking at a different app.
      for (const root of linkedModuleRoots(packagePath)) {
        for (const source of sourceFiles(root)) this.addWatchFile(source)
      }
      let artifact: SwiftArtifact
      try {
        artifact = await artifactFor(project)
      } catch (error) {
        // a compile error must not fail this module. a module whose load threw
        // leaves the graph unwatched, so the edit that fixes the error would
        // never rebuild it and the app would sit on a stale frame until the
        // dev server restarted. the previous artifact keeps rendering, and the
        // error is reported here instead.
        if (!project.artifact) throw error
        this.warn(error instanceof Error ? error.message : String(error))
        artifact = project.artifact
      }
      // the bundle image table rides the same module: a photo edit moves
      // its urls without rebuilding the artifact, and the watched files
      // re-run this load so the mount sees the new table on hot reload.
      let resources: Record<string, string> = {}
      try {
        const built = buildSwiftResourceTable(project, serverOrigin())
        resources = built.table
        for (const watched of built.watched) this.addWatchFile(watched)
      } catch (error) {
        this.warn(
          `swift resources for ${packagePath}: ${error instanceof Error ? error.message : String(error)}`,
        )
      }
      // the /res pack rides the same module as bytes: a json edit moves the
      // pack, the watched files re-run this load, and the mount's
      // bundleFiles compare treats it as a new app and remounts through it.
      let bundleFiles: Record<string, string> = {}
      try {
        const packed = buildSwiftBundleFiles(packagePath)
        bundleFiles = packed.files
        for (const watched of packed.watched) {
          this.addWatchFile(path.join(packagePath, ...watched.split(path.posix.sep)))
        }
        if (packed.skipped.length > 0) {
          this.warn(
            `swift bundle files skipped for ${packagePath}: ${packed.skipped.join('; ')}`,
          )
        }
      } catch (error) {
        this.warn(
          `swift bundle files for ${packagePath}: ${error instanceof Error ? error.message : String(error)}`,
        )
      }
      // every import mounts through the one entry: the artifact answers
      // whether it is a view or an app root at runtime, so the transform
      // never re-derives the conformance from source. the entry is rooted by
      // its package and source path so a hot reload finds the same instance.
      return swiftEntryModuleEsm({
        url: artifactUrl(project, artifact),
        hash: artifact.hash,
        rootId: entryRootId(packagePath, file),
        resources,
        bundleFiles,
      })
    },

    configureServer(server) {
      devServer = server

      server.middlewares.use((req, res, next) => {
        const url = (req.url ?? '').split('?')[0]
        if (url.startsWith(RESOURCE_PATH)) {
          // hash-addressed: the table this load emitted names these hashes,
          // so a hit serves known bytes and anything else is a 404, never a
          // neighboring file. the tenant worker is cross-origin like the
          // artifact fetch, hence the same cors header.
          for (const project of projects.values()) {
            const served = serveSwiftResource(project, url.slice(RESOURCE_PATH.length))
            if (served) {
              res.setHeader('access-control-allow-origin', '*')
              res.setHeader('content-type', served.contentType)
              res.end(served.bytes)
              return
            }
          }
          res.statusCode = 404
          res.end('swift resource is stale')
          return
        }
        if (!url.startsWith(ARTIFACT_PATH)) {
          next()
          return
        }
        // the url names a content hash, not a package: the request is answered
        // by whichever known project built it, so a React Native app whose
        // swift package sits below its root serves the same path its imports
        // compiled. served from what is already built, never rebuilt here: the
        // page only knows a hash that `load` emitted, and `load` under a compile
        // error emits the last good hash, whose bytes are still the project's
        // artifact. rebuilding would fail that fetch with the compile error and
        // fail every other project's artifact with it.
        for (const project of projects.values()) {
          const artifact = project.artifact
          if (artifact && url === `${ARTIFACT_PATH}${artifact.hash}.wasm`) {
            // the tenant worker's origin is the shell, never this dev server
            res.setHeader('access-control-allow-origin', '*')
            res.setHeader('content-type', 'application/wasm')
            res.end(artifact.bytes)
            return
          }
        }
        res.statusCode = 404
        res.end('swift artifact is stale')
      })
    },
  }
}