import * as cache from '@actions/cache'
import * as core from '@actions/core'
import * as io from '@actions/io'
import assert from 'assert'
import * as fs from 'fs'
import PQueue from 'p-queue'
import * as path from 'path'
import * as cachix from './cachix.js'
import * as nix from './nix.js'
import { execSync } from 'child_process'
import { writeFileSync } from 'fs';


interface BuildGraphNode {
  drvPath: string
  state: 'pending' | 'failed'
  references: Set<BuildGraphNode>
  referrers: Set<BuildGraphNode>
}

class BuildGraph {
  private graph: Map<string, BuildGraphNode> = new Map()
  private ready_nodes: BuildGraphNode[] = []
  private resolve_ready: (null | ((value: BuildGraphNode | PromiseLike<BuildGraphNode>) => void)) = null

  public add(drvPath: string, references: Iterable<string>, rosDerivations: Map<string, Derivation>) {
    let node = this.graph.get(drvPath)
    if (node === undefined) {
      node = {
        drvPath,
        state: 'pending',
        references: new Set(),
        referrers: new Set()
      }
      this.graph.set(drvPath, node)
    }

    for (const reference of references) {
      if (rosDerivations.get(reference) === undefined)
        continue; // don't add sources and non-ROS packages

      let referenceNode = this.graph.get(reference)
      if (referenceNode === undefined) {
        referenceNode = {
          drvPath: reference,
          state: 'pending',
          references: new Set(),
          referrers: new Set()
        }
        this.graph.set(reference, referenceNode)
      }

      referenceNode.referrers.add(node)
      node.references.add(referenceNode)
    }
  }

  /// Find the initial set of packages that are ready to build and add them to
  /// the ready queue. Must be called after all packages are added.
  public init() {
    for (const node of this.graph.values()) {
      if (node.references.size == 0) {
        this.putReady(node)
      }
    }
  }

  public empty(): boolean {
    return this.graph.size === 0 && this.ready_nodes.length == 0
  }

  public drvPaths(): Iterable<string> {
    return this.graph.keys()
  }

  private putReady(node: BuildGraphNode) {
    assert(this.graph.delete(node.drvPath))
    if (this.resolve_ready !== null) {
      this.resolve_ready(node)
      this.resolve_ready = null
    } else {
      this.ready_nodes.push(node)
    }
  }

  public takeReady(): Promise<BuildGraphNode> {
    const node = this.ready_nodes.pop()
    if (node === undefined) {
      return new Promise(resolve => this.resolve_ready = resolve)
    } else {
      return Promise.resolve(node)
    }
  }

  public succeeded(node: BuildGraphNode) {
    assert.strictEqual(node.state, 'pending')
    assert.strictEqual(node.references.size, 0)
    for (const referrer of node.referrers) {
      assert(referrer.references.delete(node))
      if (referrer.references.size == 0) {
        this.putReady(referrer)
      }
    }
  }

  public failed(node: BuildGraphNode) {
    assert.strictEqual(node.state, 'pending')
    node.state = 'failed'
    // Disconnect from graph so diamond dependencies don't cause it to added
    // twice
    for (const reference of node.references) {
      assert(reference.referrers.delete(node))
    }
    for (const referrer of node.referrers) {
      // Don't remove reference from referrer so we can print them out later
      this.failed(referrer)
      this.putReady(referrer)
    }
  }
}

interface Derivation {
  attr: string,
  drvPath: string,
  references: Set<string>,
}

interface EvalSuccess {
  readonly status: 'success',
  readonly derivation: Derivation,
}

interface EvalFailure {
  readonly status: 'failure',
  readonly attr: string,
  readonly stderr: string,
}

interface EvalError {
  readonly status: 'error',
  readonly attr: string,
  readonly error: unknown,
}

type EvalResult = EvalSuccess | EvalFailure | EvalError;


async function instantiate(nixFile: string, rootAttribute: string, drvDir: string, system?: string, parallelism = 4): Promise<EvalResult[]> {
  const jobs = await nix.evalJobs(nixFile, rootAttribute)

  const queue = new PQueue({ concurrency: parallelism })

  return await queue.addAll(jobs.map(job => async () => {
    const attr = job.attr
    try {
      const drvPath = await fs.promises.realpath(job.drvPath)

      const references = new Set(await nix.getReferences(drvPath));

      return {
        status: 'success' as const,
        derivation: {
          attr,
          drvPath,
          references
        }
      }
    } catch (error: unknown) {
      core.warning(`${job.attr} failed to evaluate`)
      if (typeof error === 'string') {
        return {
          attr,
          status: 'failure' as const,
          stderr: error
        }
      } else {
        return {
          attr,
          status: 'error' as const,
          error
        }
      }
    }
  }
  ))
}

interface BuildSuccess {
  readonly status: 'success',
  readonly drvPath: string,
}

interface BuildCachedSuccess {
  readonly status: 'cached_success',
  readonly drvPath: string,
}

interface BuildFailure {
  readonly status: 'failure',
  readonly drvPath: string,
  readonly stderr: string,
}

interface BuildCachedFailure {
  readonly status: 'cached_failure',
  readonly drvPath: string,
}

interface BuildError {
  readonly status: 'error',
  readonly drvPath: string,
  readonly error: unknown,
}

type BuildResult = BuildSuccess | BuildCachedSuccess | BuildFailure | BuildCachedFailure | BuildError;

// Regular expression matching debug output of a derivation. We don't
// want the debug output to pollute the binary cache and consume its limited capacity.
const excludedOutput = new RegExp("-debug$")

async function build(drvPath: string, resultDir: string, cachixCache: string, cacheDir: string, progress: string): Promise<BuildResult> {
  const cacheKey = "failed-" + path.basename(drvPath, ".drv");

  try {
    if (await nix.isDrvCached(drvPath, excludedOutput)) { // skip debug output
      core.info(`${progress} found cached: ${drvPath}`)
      return {
        status: 'cached_success',
        drvPath
      }
    }
  } catch (e: unknown) {
    core.warning(`failed to query binary cache for ${drvPath}: ${e}`)
  }

  try {
    if (cache.isFeatureAvailable()) {
      // We don't actually store anything in the cache, just lookup the
      // derivation name
      if (await cache.restoreCache([cacheDir], cacheKey, [], { lookupOnly: true }) !== undefined) {
        core.debug(`found cached failure: ${drvPath}`)
        return {
          status: 'cached_failure',
          drvPath
        }
      }
    }
  } catch (e: unknown) {
    core.warning(`failed to query failed build cache for ${drvPath}: ${e}`)
  }

  try {
    core.info(`${progress} building: ${drvPath}`)

    const resultPath = path.join(resultDir, path.basename(drvPath, ".drv"))
    const outputs = await nix.realize(drvPath, resultPath)

    await cachix.push(cachixCache, outputs.filter((output) => !excludedOutput.test(output)))

    return {
      status: 'success',
      drvPath
    }
  } catch (error: unknown) {
    core.debug(`${drvPath} failed to build`)
    if (typeof error === 'string') {
      try {
        if (cache.isFeatureAvailable()) {
          await cache.saveCache([cacheDir], cacheKey)
        }
      } catch (e: unknown) {
        core.warning(`failed to cache failed build for ${drvPath}: ${e}`)
      }

      return {
        status: 'failure',
        drvPath,
        stderr: error
      }
    } else {
      return {
        status: 'error',
        drvPath,
        error
      }
    }
  }
}

interface SuccessResult {
  attr: string,
  drvPath: string,
}

interface CachedSuccessResult {
  attr: string,
  drvPath: string
}

interface CachedFailureResult {
  attr: string,
  drvPath: string
}

interface BuildFailureResult {
  attr: string,
  drvPath: string
  stderr: string
}

interface DependencyFailureResult {
  attr: string,
  drvPath: string,
  failedDependencies: string[]
}

interface ErrorResult {
  attr: string,
  drvPath?: string
  error: unknown
}

async function run() {
  try {
    const nixFile = core.getInput('nix-file', { required: true })
    const rootAttribute = core.getInput('root-attribute', { required: true })
    const system = core.getInput('system') || undefined
    const evalJobs = parseInt(core.getInput('eval-jobs', { required: true }))
    const buildJobs = parseInt(core.getInput('build-jobs', { required: true }))
    const cachixCache = core.getInput('cachix-cache', { required: true })

    const buildDir = "build"
    const drvDir = path.join(buildDir, 'drvs')
    const resultDir = path.join(buildDir, 'results')
    // Empty directory to make GitHub Actions cache happy
    const cacheDir = path.join(buildDir, 'cache')

    await io.mkdirP(drvDir)
    await io.mkdirP(resultDir)
    await io.mkdirP(cacheDir)

    const successes: SuccessResult[] = []
    const cachedSuccesses: CachedSuccessResult[] = []
    const evalFailures: EvalFailure[] = []
    const buildFailures: BuildFailureResult[] = []
    const cachedFailures: CachedFailureResult[] = []
    const dependencyFailures: DependencyFailureResult[] = []
    const errors: ErrorResult[] = []

    core.startGroup("installing nix-eval-jobs...")
    execSync("nix-env --install --file default.nix -A nix-eval-jobs");
    core.endGroup()

    core.startGroup(`evaluating packages in ${rootAttribute}...`)
    const evalResults = await instantiate(nixFile, rootAttribute, drvDir, system, evalJobs)
    const derivations: Map<string, Derivation> = new Map()
    for (const r of evalResults) {
      if (r.status === 'success') {
        derivations.set(r.derivation.drvPath, r.derivation)
      } else if (r.status === 'failure') {
        evalFailures.push(r)
      } else if (r.status === 'error') {
        errors.push({
          attr: r.attr,
          error: r.error
        })
      }
    }
    core.endGroup()

    // Create graph of references used to order the builds
    const buildGraph = new BuildGraph();
    for (const [drvPath, drv] of derivations) {
      buildGraph.add(drvPath, drv.references, derivations)
    }
    buildGraph.init()

    core.startGroup(`building ${derivations.size} packages...`)

    const queue = new PQueue({ concurrency: buildJobs })
    let num_done = 0
    do {
      const node = await buildGraph.takeReady()

      queue.add(async () => {
        if (node.state === 'failed') {
          const drv = derivations.get(node.drvPath)
          if (drv === undefined) {
            num_done++
            return;
          }

          const failedDependencies: string[] = [];
          for (const reference of node.references) {
            if (reference.state === 'failed') {
              failedDependencies.push(reference.drvPath)
            }
          }

          assert.notStrictEqual(failedDependencies.length, 0)

          dependencyFailures.push({
            attr: drv.attr,
            drvPath: drv.drvPath,
            failedDependencies
          })

          num_done++
          return;
        }

        const progress = `[${(num_done / derivations.size * 100).toFixed(1)}%]`
        const r = await build(node.drvPath, resultDir, cachixCache, cacheDir, progress)
        num_done++

        if (r.status === 'cached_success' || r.status === 'success') {
          buildGraph.succeeded(node)
        } else {
          buildGraph.failed(node)
        }

        const drv = derivations.get(r.drvPath)
        if (drv === undefined) {
          return;
        }

        switch (r.status) {
          case 'success':
            successes.push({
              attr: drv.attr,
              drvPath: drv.drvPath,
            })
            break;
          case 'cached_success':
            cachedSuccesses.push({
              attr: drv.attr,
              drvPath: drv.drvPath,
            })
            break;
          case 'failure':
            buildFailures.push({
              attr: drv.attr,
              drvPath: r.drvPath,
              stderr: r.stderr
            })
            break;
          case 'cached_failure':
            cachedFailures.push({
              attr: drv.attr,
              drvPath: r.drvPath,
            })
            break;
          case 'error':
            errors.push({
              attr: drv.attr,
              drvPath: drv.drvPath,
              error: r.error
            })
        }
      })
    } while (!buildGraph.empty())
    core.endGroup()

    await queue.onIdle()

    core.info(`Creating linkFarm.nix from successful builds`)
    const successfulAttrs: string[] = [];
    function successfulAttrsPush(attr: string) {
      if (attr != "buildEnv")   // buildEnv is a function, not derivation
        successfulAttrs.push(attr);
    }
    successes.forEach((element) => successfulAttrsPush(element.attr));
    cachedSuccesses.forEach((element) => successfulAttrsPush(element.attr));

    writeFileSync("linkFarm.nix", `
      { name }:
      with import ./. { system = "${system}"; };
      linkFarmFromDrvs name (with ${rootAttribute}; [
        ${successfulAttrs.join("\n")}
      ])`)

    core.startGroup("Results")
    core.info(`Successes: ${successes.length}`)
    core.info(`Cached: ${cachedSuccesses.length}`)
    core.info(`Cached failures: ${cachedFailures.length}`)
    core.info(`Evaluation failures: ${evalFailures.length}`)
    core.info(`Dependency failures: ${dependencyFailures.length}`)
    core.info(`Build failures: ${buildFailures.length}`)
    core.info(`Unknown errors: ${errors.length}`)
    core.endGroup()

    for (const r of errors) {
      core.startGroup(`Unknown error building ${r.attr} (${r.drvPath})`)
      core.error(`${r.error}`)
      core.endGroup()
    }

    for (const r of buildFailures) {
      await core.group(
        `Failed to build ${r.attr} (${r.drvPath})`,
        async () => core.warning(r.stderr)
      )
    }

    for (const r of dependencyFailures) {
      core.startGroup(`Dependency of ${r.attr} (${r.drvPath}) failed to build`)
      for (const failedDependency of r.failedDependencies) {
        core.warning(`dependency: ${failedDependency}`)
      }
      core.endGroup()
    }

    for (const r of evalFailures) {
      core.startGroup(`Failed to evaluate ${r.attr}`)
      core.warning(r.stderr)
      core.endGroup()
    }

    core.startGroup("Cached build failures")
    cachedFailures.forEach(r => core.info(`${r.attr} (${r.drvPath})`));
    core.endGroup()

    for (const r of successes) {
      await core.group(
        `Sucessfully built ${r.attr} (${r.drvPath})`,
        () => nix.printLog(r.drvPath!).catch(() => undefined)
      )
    }

    core.startGroup("Already built packages")
    cachedSuccesses.forEach(r => core.info(`${r.attr} (${r.drvPath})`));
    core.endGroup()
  } catch (e: unknown) {
    assert(e instanceof Error)
    core.setFailed(e.message)
  }
}

run()
