// @ts-check

const { probePython, probeWhl, resolveLoginName } = require('./utils');

let bundleVersion = '0.5.1';
let bundleHash = '11d418d';
const CANCEL_INTERRUPT_GRACE_MS = 3000;
const CANCEL_KILL_WAIT_GRACE_MS = 1500;

/**
 * Construct the ASCT bundle path based on the version and hash defined above.
 * @param {string} toolsRoot
 * @returns {string}
 */
function buildBundlePath(toolsRoot) {
  return `${toolsRoot}/asct/${bundleVersion}/asct-${bundleVersion}+${bundleHash}.tar.gz`;
}

/**
 * Stop a running ASCT process with interrupt-first behavior and bounded waits.
 * @param {import("../recipes/docs/jsdocs").Engine} engine
 * @param {import("../recipes/docs/jsdocs").ToolContext} ctx
 * @param {string} toolName
 * @param {string} triggerName
 */
async function stopAsctRun(engine, ctx, toolName, triggerName) {
  const runHandle = ctx.metadata.runHandle;
  if (!runHandle) {
    engine.log(
      'info',
      `${triggerName} requested for ${toolName}, but no running process handle was found`,
    );
    return;
  }

  /** @param {string} reason */
  const killAndWait = async (reason) => {
    engine.log('warn', reason);
    await runHandle.kill();
    engine.log('info', `${toolName} kill signal sent to PID ${runHandle.pid}`);

    const killWaitResult = await Promise.race([
      runHandle
        .wait()
        .then(() => 'exited')
        .catch(() => 'exited'),
      new Promise((resolve) =>
        setTimeout(() => resolve('timed_out'), CANCEL_KILL_WAIT_GRACE_MS),
      ),
    ]);

    if (killWaitResult === 'timed_out') {
      engine.log(
        'warn',
        `${toolName} did not exit within ${CANCEL_KILL_WAIT_GRACE_MS}ms after kill. The process may need to be manually killed on the target side (PID ${runHandle.pid}).`,
      );
    }
  };

  try {
    engine.log(
      'info',
      `${triggerName} requested for ${toolName}; sending interrupt signal to PID ${runHandle.pid}`,
    );
    await runHandle.interrupt();
    engine.log(
      'info',
      `${toolName} interrupt signal sent to PID ${runHandle.pid}`,
    );

    const waitResult = await Promise.race([
      runHandle.wait().then(() => 'exited'),
      new Promise((resolve) =>
        setTimeout(() => resolve('timed_out'), CANCEL_INTERRUPT_GRACE_MS),
      ),
    ]);

    if (waitResult === 'exited') {
      engine.log(
        'info',
        `${toolName} exited after interrupt for PID ${runHandle.pid}`,
      );
      return;
    }

    await killAndWait(
      `${toolName} did not exit within ${CANCEL_INTERRUPT_GRACE_MS}ms after interrupt; falling back to kill for PID ${runHandle.pid}`,
    );
  } catch {
    await killAndWait(
      `Interrupt failed for ${toolName} PID ${runHandle.pid}; falling back to kill`,
    );
  }
}

/**
 * @param {import("../recipes/docs/jsdocs").ToolContext} ctx
 * @returns {boolean}
 */
function isCancellationRequested(ctx) {
  return ctx.metadata.cancelRequested === true;
}

const ASCT_BENCHMARK_OPTIONS = [
  'idle-latency',
  'peak-bandwidth',
  'cross-numa-bandwidth',
  'latency-sweep',
  'bandwidth-sweep',
  'loaded-latency',
  'c2c-latency',
  'storage-request-size-sweep',
  'storage-io-depth-sweep',
  'storage-process-count-sweep',
  'storage-access-pattern-sweep',
];

/** @type {Record<string, number>} */
const ASCT_BENCHMARK_MAX_ESTIMATED_MINUTES = {
  'idle-latency': 3,
  'peak-bandwidth': 1,
  'cross-numa-bandwidth': 1,
  'latency-sweep': 3,
  'bandwidth-sweep': 2,
  'loaded-latency': 5,
  'c2c-latency': 9,
  'storage-request-size-sweep': 15,
  'storage-io-depth-sweep': 15,
  'storage-process-count-sweep': 15,
  'storage-access-pattern-sweep': 15,
};

/**
 * Calculate a max estimated runtime in minutes from selected benchmarks.
 * Unknown benchmark names are ignored.
 * @param {string[]} benchmarks
 * @returns {number}
 */
function getMaxEstimatedBenchmarkRuntimeMinutes(benchmarks) {
  const baseEstimateMinutes = benchmarks.reduce((total, benchmarkName) => {
    const benchmarkEstimate =
      ASCT_BENCHMARK_MAX_ESTIMATED_MINUTES[benchmarkName] ?? 0;
    return total + benchmarkEstimate;
  }, 0);

  return Math.ceil(baseEstimateMinutes * 1.5);
}

const PYTHON_VER_MAJOR = 3;
const PYTHON_VER_MINOR = 10;
const readinessMessageCode =
  'engine.recipeparser.js_recipe_stage.READINESS_MESSAGE';

/**
 * @param {import("../recipes/docs/jsdocs").Engine} engine
 * @returns {Promise<import("../recipes/docs/jsdocs").ProbeAdvice>}
 */
async function probeNumactl(engine) {
  const check = await engine.execCommand(
    ['bash', '-c', 'command -v numactl'],
    {},
  );
  if (check.rc !== 0) {
    return {
      level: 'error',
      messageCode: readinessMessageCode,
      metadata: {
        message:
          'numactl is not available on the target machine. Install the numactl package on the target machine.',
      },
    };
  }

  return {
    level: 'ready',
    messageCode: '',
  };
}

/**
 * @type {import("../recipes/docs/jsdocs").ToolIntegration}}
 */
let tool = {
  name: 'asct',
  version: bundleVersion,
  supportsWorkloadLaunch: true,
  description: {
    short: 'ASCT runs system characterization benchmarks and diagnostics.',
    long: 'ASCT (Arm System Characterization Tool) is a standalone command-line utility for running low-level benchmarks, diagnostic scripts, and system tests on Arm-based platforms. It provides a standardized environment for measuring hardware characteristics such as memory latency/bandwidth and storage performance, supporting platform bring-up, tuning, and architectural comparisons.',
  },
  parameters: [
    {
      id: 'systemInfoOnly',
      label: 'System info only',
      description: 'Run only system-info (no benchmarks).',
      config: {
        type: 'checkbox',
        defaultValue: false,
      },
    },
    {
      id: 'benchmarks',
      label: 'Benchmarks',
      description: 'List of benchmarks to run when not in system-info mode.',
      config: {
        type: 'select',
        options: ASCT_BENCHMARK_OPTIONS,
        defaultValue: [],
      },
    },
  ],
  deployments: [
    {
      appliesTo: [
        { architecture: 'aarch64', os: 'Linux' },
        { architecture: 'x86_64', os: 'Linux' },
      ],
      dependencies: [
        {
          type: 'tool_bundle',
          name: 'asct',
          version: bundleVersion,
          requiredWhen: { type: 'always' },
        },
      ],
    },
  ],
  probe: async (engine, ctx) => {
    /** @type {import("../recipes/docs/jsdocs").ProbeAdvice[]} */
    let advice = [];

    let py = await probePython(
      engine,
      PYTHON_VER_MAJOR,
      PYTHON_VER_MINOR,
      tool.name,
    );
    if (py.level !== 'ready') advice.push(py);

    let whlAdvice = await probeWhl(
      engine,
      buildBundlePath(ctx.toolsRoot),
      tool.name,
    );
    if (whlAdvice.level !== 'ready') {
      advice.push(whlAdvice);
    }

    let numactlAdvice = await probeNumactl(engine);
    if (numactlAdvice.level !== 'ready') {
      advice.push(numactlAdvice);
    }

    return {
      available: advice.length === 0,
      capabilities: {},
      advice,
    };
  },

  reformat: async (engine, ctx) => {},

  run: async (engine, ctx) => {
    ctx.metadata.cancelRequested = false;
    let deployPath = buildBundlePath(ctx.toolsRoot);

    let outputDir = await engine.createTempDir();

    engine.emitOutput(
      outputDir + '/data/**/*',
      'output/*',
      /** @type {any} */ ({ name: 'asct-data', version: '1.0' }),
    );
    engine.emitOutput(
      outputDir + '/data/asct.log',
      'asct.log',
      /** @type {any} */ ({ name: 'log-text', version: '1.0' }),
    );

    let createVenvResult = await engine.execCommand(
      ['python3', '-m', 'venv', outputDir + '/venv'],
      {},
    );
    if (createVenvResult.rc !== 0) {
      throw {
        code: 'tool_integrations.common.CREATE_PYTHON_VENV',
        metadata: {
          tool: tool.name,
          pythonVersion: '3.10',
          exitCode: createVenvResult.rc,
        },
      };
    }

    let wheelExistsResult = await engine.execCommand(['stat', deployPath], {});
    if (wheelExistsResult.rc !== 0) {
      throw {
        code: 'tool_integrations.common.TOOL_NOT_DEPLOYED',
        metadata: { tool: tool.name, deployPath: deployPath },
      };
    }

    let installResult = await engine.execCommand(
      [outputDir + '/venv/bin/pip', 'install', deployPath],
      {},
    );
    if (installResult.rc !== 0) {
      throw {
        code: 'tool_integrations.common.INSTALL_MODULE',
        metadata: {
          tool: tool.name,
          exitCode: installResult.rc,
          deployPath: deployPath,
        },
      };
    }

    let testCmd = [outputDir + '/venv/bin/asct', '--help'];
    let testResult = await engine.execCommand(testCmd, {});
    if (testResult.rc !== 0) {
      // This throw is fine as it will be scooped up into a SCRIPTED_STAGE_ERROR later (and there's no advice here anyway)
      throw `asct module installed successfully, but ${testCmd.join(' ')} exited with code ${testResult.rc}`;
    }

    let benchmarks = Array.isArray(ctx.params?.benchmarks)
      ? ctx.params.benchmarks
      : [];
    let systemInfoOnly = ctx.params?.systemInfoOnly === true;

    let processArgs = [outputDir + '/venv/bin/asct'];
    if (systemInfoOnly) {
      processArgs.push('system-info');
    } else {
      processArgs.push('run');
      if (benchmarks.length > 0) {
        processArgs.push(...benchmarks);
      }
    }

    processArgs.push(
      '--format',
      'csv',
      '--output-dir',
      outputDir + '/data',
      '--force',
      '--log-file',
      outputDir + '/data/asct.log',
    );

    let runAsRoot = true;
    // Collect the current PATH information in the same setting as we run ASCT later.
    let currentPath = await engine.execCommand(['bash', '-c', 'echo $PATH'], {
      asRoot: runAsRoot,
    });
    let currentPythonPath = await engine.execCommand(
      ['bash', '-c', 'echo ${PYTHONPATH:-}'],
      {
        asRoot: runAsRoot,
      },
    );
    const venvPath = outputDir + '/venv';
    const existingPath = currentPath.stdout.trim();
    const existingPythonPath = currentPythonPath.stdout.trim();
    const updatePath = [`${venvPath}/bin`, ...existingPath.split(':')]
      .filter(Boolean)
      .join(':');
    const updatePythonPath = [
      `${venvPath}/lib/python${PYTHON_VER_MAJOR}.${PYTHON_VER_MINOR}/site-packages`,
      ...existingPythonPath.split(':'),
    ]
      .filter(Boolean)
      .join(':');

    engine.log('info', `PATH for ${tool.name} run: ${updatePath}`);
    engine.log('info', `PYTHONPATH for ${tool.name} run: ${updatePythonPath}`);
    engine.log(
      'info',
      `Starting ${tool.name} with args: ${processArgs.join(' ')}`,
    );

    if (isCancellationRequested(ctx)) {
      engine.log(
        'info',
        `${tool.name} launch skipped because cancellation was requested before process start`,
      );
      return;
    }

    const benchmarkSelectionSummary =
      benchmarks.length > 0
        ? `${benchmarks.length} selected`
        : 'default benchmarks';
    const maxEstimatedRuntimeMinutes =
      getMaxEstimatedBenchmarkRuntimeMinutes(benchmarks);
    const runtimeEstimateSummary =
      maxEstimatedRuntimeMinutes > 0
        ? `may take ${maxEstimatedRuntimeMinutes}+ min`
        : 'may take 20+ min';
    const progressTrackerId = systemInfoOnly
      ? `Collecting ${tool.name} system information`
      : `Collecting ${tool.name} benchmarks (${benchmarkSelectionSummary}; ${runtimeEstimateSummary} depending on target size and system load)`;
    engine.startProgressTracker(progressTrackerId);

    let runHandle = await engine.startProcess(processArgs, {
      asRoot: runAsRoot,
      environment: {
        PATH: updatePath,
        PYTHONPATH: updatePythonPath,
        VIRTUAL_ENV: venvPath,
      },
    });
    ctx.metadata.runHandle = runHandle;

    if (isCancellationRequested(ctx)) {
      await stopAsctRun(engine, ctx, tool.name, 'Cancellation');
    }

    let runResult;
    try {
      runResult = await runHandle.wait();
    } finally {
      if (runAsRoot) {
        try {
          const loginName = await resolveLoginName(engine);
          await engine.execCommand(
            [
              'bash',
              '-c',
              `chown -R "${loginName}":"${loginName}" "${outputDir}"`,
            ],
            { asRoot: true },
          );
          await engine.execCommand(
            ['bash', '-c', `chmod -R u+rwX,go+rX "${outputDir}"`],
            { asRoot: true },
          );
        } catch (err) {
          engine.log(
            'warn',
            `Failed to fix ownership/permissions for ${tool.name} output dir: ${String(err)}`,
          );
        }
      }
    }

    if (runResult.exitCode === 0) {
      engine.log('info', `${tool.name} module executed successfully`);
    } else {
      engine.log(
        'error',
        `${tool.name} module failed with exit code ${runResult.exitCode}`,
      );
      throw {
        code: 'tool_integrations.asct.UNEXPECTED_ERROR',
        metadata: { exitCode: runResult.exitCode, log: 'asct.log' },
      };
    }

    engine.endProgress(progressTrackerId);
  },

  onCancel: async (engine, ctx) => {
    ctx.metadata.cancelRequested = true;
    await stopAsctRun(engine, ctx, tool.name, 'Cancellation');
  },

  onStop: async (engine, ctx) => {
    ctx.metadata.cancelRequested = true;
    await stopAsctRun(engine, ctx, tool.name, 'Stop');
  },
};

// ASCT does not require objdump; no additional probe checks are needed.
