import assert from 'node:assert/strict';
import { spawnSync } from 'node:child_process';
import { createHash } from 'node:crypto';
import {
  mkdtemp,
  readFile,
  rm,
  writeFile
} from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { basename, delimiter, dirname, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';

const scriptPath = fileURLToPath(import.meta.url);
const testRoot = resolve(dirname(scriptPath));
const projectRoot = resolve(testRoot, '../..');
const modelDirectory = resolve(projectRoot, 'docs/protocol');
const modelPath = resolve(modelDirectory, 'Nuthatch.tla');
const defaultConfigurationPaths = [
  'Nuthatch.cfg',
  'NuthatchContentIdentity.cfg',
  'NuthatchKinds.cfg',
  'NuthatchQuota.cfg',
  'NuthatchSequential.cfg'
].map((name) => resolve(modelDirectory, name));
const configurationPaths = process.env.NUTHATCH_MODEL_CONFIGS
  ? process.env.NUTHATCH_MODEL_CONFIGS.split(delimiter)
  : defaultConfigurationPaths;
const actionConfigurationPath = resolve(
  testRoot,
  'NuthatchCorrespondence.cfg'
);
const fixturePath =
  process.env.NUTHATCH_MODEL_FIXTURES ??
  resolve(testRoot, 'concrete-snapshots.json');
const evidenceExecutable = process.env.NUTHATCH_MODEL_EVIDENCE;
const toolsJar = process.env.TLA2TOOLS_JAR ?? '/usr/local/lib/tla2tools.jar';
const pins = parseAssignments(
  await readFile(resolve(testRoot, 'tools.env'), 'utf8')
);
const workDirectory = await mkdtemp(
  resolve(tmpdir(), 'nuthatch-model-correspondence-')
);

assert.ok(evidenceExecutable, 'NUTHATCH_MODEL_EVIDENCE is required');

try {
  await verifyToolPin();
  const graph = await buildSemanticGraph();
  const semanticLeafActions = discoverLeafActions(graph, 'Next');
  const executedLeafActions = await executeLeafActions();
  assertSetEqual(
    executedLeafActions,
    semanticLeafActions,
    'minimal TLC action coverage'
  );

  const configurations = await readConfigurations(graph);
  const configuredChecks = new Set([
    ...configurations[0].invariants,
    ...configurations[0].properties
  ]);
  const inventory = readInventory();
  const inventoryActionNames = new Set(inventory.actions.values());
  const inventoryCheckNames = new Set(inventory.checks.values());
  assertSetEqual(
    inventoryActionNames,
    semanticLeafActions,
    'typed implementation action inventory'
  );
  assertSetEqual(
    inventoryCheckNames,
    configuredChecks,
    'typed implementation check inventory'
  );

  await verifyTypedEvidenceUse(inventory);
  verifyScenarios(inventory);
  const snapshots = await verifyProjection();
  if (process.env.NUTHATCH_SKIP_MODEL_MUTATIONS !== '1') {
    await verifyNegativeMutations(snapshots);
  }
  await verifyNoSecondNext();

  console.log(
    [
      `PASS SANY and TLC agree on ${semanticLeafActions.size} executable leaf actions`,
      `PASS ${configurations.length} model configurations resolve ${configuredChecks.size} typed checks`,
      `PASS ${inventory.scenarios.length} typed scenarios cover every leaf action`,
      `PASS ${snapshots.length} concrete snapshots project without a transition oracle`,
      'PASS stale, missing, unknown, uncovered, and invalid-projection mutations fail'
    ].join('\n')
  );
} finally {
  await rm(workDirectory, { recursive: true, force: true });
}

async function verifyToolPin() {
  const jar = await readFile(toolsJar);
  const digest = createHash('sha256').update(jar).digest('hex');
  assert.equal(digest, pins.TLA_TOOLS_SHA256, 'TLA+ Tools checksum');
  const helpResult = spawnSync(
    'java',
    ['-cp', toolsJar, 'tlc2.TLC', '-help'],
    { encoding: 'utf8' }
  );
  assert.ok(
    helpResult.status === 0 || helpResult.status === 1,
    `TLA+ Tools version probe failed with ${helpResult.status}`
  );
  const help = `${helpResult.stdout ?? ''}${helpResult.stderr ?? ''}`;
  assert.match(
    help,
    new RegExp(`Version ${escapeExpression(pins.TLA_TOOLS_VERSION)}\\b`),
    'TLA+ Tools version'
  );
}

async function buildSemanticGraph() {
  run(
    'java',
    ['-cp', toolsJar, 'tla2sany.SANY', '-d', modelPath, 'dot'],
    { cwd: workDirectory }
  );
  const semanticGraphPath = resolve(workDirectory, 'Nuthatch.dot');
  return parseSemanticGraph(await readFile(semanticGraphPath, 'utf8'));
}

async function executeLeafActions() {
  const actionGraphPath = resolve(workDirectory, 'actions.dot');
  run(
    'java',
    [
      '-cp',
      toolsJar,
      'tlc2.TLC',
      '-cleanup',
      '-workers',
      '1',
      '-metadir',
      resolve(workDirectory, 'states'),
      '-config',
      actionConfigurationPath,
      '-dump',
      'dot,actionlabels',
      actionGraphPath,
      basename(modelPath)
    ],
    { cwd: modelDirectory }
  );
  return extractActionGraphLabels(await readFile(actionGraphPath, 'utf8'));
}

async function readConfigurations(graph) {
  const configurations = await Promise.all(
    configurationPaths.map(async (path) => ({
      name: basename(path),
      ...parseConfiguration(await readFile(path, 'utf8'))
    }))
  );
  for (const configuration of configurations) {
    for (const symbol of [
      ...configuration.invariants,
      ...configuration.properties
    ]) {
      assert.ok(
        graph.operatorNames.has(symbol),
        `${configuration.name} references unknown model operator ${symbol}`
      );
    }
  }
  for (const configuration of configurations.slice(1)) {
    assertSetEqual(
      configuration.invariants,
      configurations[0].invariants,
      `${configuration.name} invariants`
    );
    assertSetEqual(
      configuration.properties,
      configurations[0].properties,
      `${configuration.name} properties`
    );
  }
  return configurations;
}

function readInventory() {
  const lines = run(evidenceExecutable, ['inventory']).trim().split('\n');
  const actions = new Map();
  const checks = new Map();
  const scenarios = [];
  for (const line of lines) {
    const fields = line.split('\t');
    switch (fields[0]) {
      case 'ACTION':
        assert.equal(fields.length, 3, `invalid action inventory line: ${line}`);
        assert.ok(!actions.has(fields[1]), `duplicate action ${fields[1]}`);
        actions.set(fields[1], fields[2]);
        break;
      case 'CHECK':
        assert.equal(fields.length, 3, `invalid check inventory line: ${line}`);
        assert.ok(!checks.has(fields[1]), `duplicate check ${fields[1]}`);
        checks.set(fields[1], fields[2]);
        break;
      case 'SCENARIO':
        assert.equal(fields.length, 6, `invalid scenario line: ${line}`);
        scenarios.push({
          constructor: fields[1],
          action: fields[2],
          precondition: fields[3],
          observation: fields[4],
          expectedPost: fields[5]
        });
        break;
      default:
        assert.fail(`unknown inventory line: ${line}`);
    }
  }
  switch (process.env.NUTHATCH_INVENTORY_MUTATION) {
    case undefined:
      break;
    case 'stale':
      actions.set('ActionRemovedModelAction', 'RemovedModelAction');
      break;
    case 'missing':
      actions.delete('ActionFirstRequest');
      break;
    case 'uncovered':
      return {
        actions,
        checks,
        scenarios: scenarios.filter(
          (scenario) => scenario.constructor !== 'ActionFirstRequest'
        )
      };
    default:
      assert.fail(
        `unknown inventory mutation ${process.env.NUTHATCH_INVENTORY_MUTATION}`
      );
  }
  return { actions, checks, scenarios };
}

async function verifyTypedEvidenceUse(inventory) {
  const source = await readFile(
    resolve(projectRoot, 'test/ProtocolConformance.hs'),
    'utf8'
  );
  const usedActions = extractConstructors(source, 'coveredAction', 'Action');
  const usedChecks = extractConstructors(source, 'checkedModel', 'Check');
  assertSetEqual(
    usedActions,
    new Set(inventory.actions.keys()),
    'typed action evidence calls'
  );
  assertSetEqual(
    usedChecks,
    new Set(inventory.checks.keys()),
    'typed check evidence calls'
  );
}

function verifyScenarios(inventory) {
  const scenarioConstructors = new Set();
  for (const scenario of inventory.scenarios) {
    assert.ok(
      inventory.actions.has(scenario.constructor),
      `unknown scenario action ${scenario.constructor}`
    );
    assert.equal(
      scenario.action,
      inventory.actions.get(scenario.constructor),
      `${scenario.constructor} scenario name`
    );
    assert.ok(scenario.precondition, `${scenario.action} precondition`);
    assert.ok(scenario.observation, `${scenario.action} observation`);
    assert.ok(scenario.expectedPost, `${scenario.action} expected post-state`);
    assert.ok(
      !scenarioConstructors.has(scenario.constructor),
      `duplicate scenario ${scenario.constructor}`
    );
    scenarioConstructors.add(scenario.constructor);
  }
  assertSetEqual(
    scenarioConstructors,
    new Set(inventory.actions.keys()),
    'scenario coverage'
  );
}

async function verifyProjection() {
  run(evidenceExecutable, ['verify', fixturePath]);
  return JSON.parse(await readFile(fixturePath, 'utf8'));
}

async function verifyNegativeMutations(snapshots) {
  assertCheckerFails(
    { NUTHATCH_INVENTORY_MUTATION: 'stale' },
    /RemovedModelAction/,
    'stale action'
  );
  assertCheckerFails(
    { NUTHATCH_INVENTORY_MUTATION: 'missing' },
    /FirstRequest/,
    'missing action'
  );
  assertCheckerFails(
    { NUTHATCH_INVENTORY_MUTATION: 'uncovered' },
    /FirstRequest/,
    'uncovered scenario'
  );

  const unknownConfigurationPath = resolve(
    workDirectory,
    'unknown-check.cfg'
  );
  const firstConfiguration = await readFile(configurationPaths[0], 'utf8');
  await writeFile(
    unknownConfigurationPath,
    `${firstConfiguration}\nINVARIANTS\n    UnknownInvariant\n`
  );
  assertCheckerFails(
    {
      NUTHATCH_MODEL_CONFIGS: [
        unknownConfigurationPath,
        ...configurationPaths.slice(1)
      ].join(delimiter)
    },
    /UnknownInvariant/,
    'unknown configured check'
  );

  const invalidFixtures = structuredClone(snapshots);
  invalidFixtures[0].concrete.external_state = 'pending';
  const invalidFixturePath = resolve(workDirectory, 'invalid-snapshots.json');
  await writeFile(
    invalidFixturePath,
    `${JSON.stringify(invalidFixtures, null, 2)}\n`
  );
  assertCheckerFails(
    { NUTHATCH_MODEL_FIXTURES: invalidFixturePath },
    /ExternalStateConsistent/,
    'invalid projection'
  );
}

function assertCheckerFails(environment, expression, label) {
  const result = spawnSync(process.execPath, [scriptPath], {
    encoding: 'utf8',
    maxBuffer: 32 * 1024 * 1024,
    env: {
      ...process.env,
      NUTHATCH_SKIP_MODEL_MUTATIONS: '1',
      ...environment
    }
  });
  assert.notEqual(result.status, 0, `${label} mutation unexpectedly passed`);
  assert.match(
    `${result.stdout ?? ''}\n${result.stderr ?? ''}`,
    expression,
    `${label} mutation failure`
  );
}

async function verifyNoSecondNext() {
  const sourceFiles = run(
    'rg',
    ['--files', 'src', 'test', '-g', '*.hs'],
    { cwd: projectRoot }
  )
    .trim()
    .split('\n')
    .filter(Boolean);
  for (const sourcePath of sourceFiles) {
    const source = await readFile(resolve(projectRoot, sourcePath), 'utf8');
    assert.doesNotMatch(source, /\bNext\b/, `${sourcePath} defines model Next`);
    assert.doesNotMatch(
      source,
      /ModelAction\s*->\s*AbstractSnapshot|AbstractSnapshot\s*->\s*ModelAction/,
      `${sourcePath} defines an abstract action transition`
    );
  }
}

function parseSemanticGraph(dot) {
  const labels = new Map();
  const edges = new Map();
  const operatorIds = new Map();
  for (const line of dot.split('\n')) {
    const operator = line.match(
      /^(\d+) \[style=filled,shape=diamond,fillcolor="red",label="([A-Za-z][A-Za-z0-9_]*)"\];$/
    );
    if (operator) {
      const [, id, label] = operator;
      operatorIds.set(label, id);
    }
    const node = line.match(/^(\d+) \[.*label="(.*)"\];$/);
    if (node) {
      const [, id, label] = node;
      const nodeLabels = labels.get(id) ?? new Set();
      nodeLabels.add(label);
      labels.set(id, nodeLabels);
    }
    const edge = line.match(/^(\d+) -> (\d+)$/);
    if (edge) {
      const [, from, to] = edge;
      const targets = edges.get(from) ?? new Set();
      targets.add(to);
      edges.set(from, targets);
    }
  }
  return {
    labels,
    edges,
    operatorIds,
    operatorNames: new Set(operatorIds.keys())
  };
}

function discoverLeafActions(graph, rootName) {
  const memo = new Map();
  const visiting = new Set();
  const visitOperator = (name) => {
    if (memo.has(name)) return memo.get(name);
    if (visiting.has(name)) return new Set();
    visiting.add(name);
    const definition = definitionText(graph, name);
    const ownsPrime = /[A-Za-z][A-Za-z0-9_]*'/.test(definition);
    let leaves;
    if (ownsPrime) {
      leaves = new Set([name]);
    } else {
      leaves = new Set();
      const dependencies = new Set(
        [...definition.matchAll(/[A-Za-z][A-Za-z0-9_]*/g)]
          .map((match) => match[0])
          .filter(
            (dependency) =>
              dependency !== name && graph.operatorNames.has(dependency)
          )
      );
      for (const dependency of dependencies) {
        for (const leaf of visitOperator(dependency)) leaves.add(leaf);
      }
    }
    visiting.delete(name);
    memo.set(name, leaves);
    return leaves;
  };
  return visitOperator(rootName);
}

function definitionText(graph, operatorName) {
  const id = graph.operatorIds.get(operatorName);
  assert.ok(id, `missing semantic operator ${operatorName}`);
  const candidates = [...(graph.edges.get(id) ?? [])].flatMap(
    (target) => [...(graph.labels.get(target) ?? [])]
  );
  assert.ok(candidates.length > 0, `missing definition body for ${operatorName}`);
  return candidates.sort((left, right) => right.length - left.length)[0];
}

function extractActionGraphLabels(dot) {
  return new Set(
    [...dot.matchAll(/label="([A-Za-z][A-Za-z0-9_]*)"/g)].map(
      (match) => match[1]
    )
  );
}

function parseConfiguration(text) {
  const invariants = new Set();
  const properties = new Set();
  let section = null;
  for (const rawLine of text.split('\n')) {
    const line = rawLine.replace(/\\\\\*.*/, '').trim();
    if (line === 'INVARIANTS' || line === 'PROPERTIES') {
      section = line;
      continue;
    }
    if (/^[A-Z_]+(?:\s|$)/.test(line)) {
      section = null;
      continue;
    }
    if (!line || !section) continue;
    const symbol = line.match(/^([A-Za-z][A-Za-z0-9_]*)$/)?.[1];
    assert.ok(symbol, `unsupported ${section} entry: ${line}`);
    (section === 'INVARIANTS' ? invariants : properties).add(symbol);
  }
  return { invariants, properties };
}

function extractConstructors(source, functionName, prefix) {
  const expression = new RegExp(
    `${functionName}\\s+(${prefix}[A-Za-z][A-Za-z0-9_]*)`,
    'g'
  );
  return new Set([...source.matchAll(expression)].map((match) => match[1]));
}

function parseAssignments(text) {
  return Object.fromEntries(
    text
      .split('\n')
      .map((line) => line.trim())
      .filter(Boolean)
      .map((line) => {
        const separator = line.indexOf('=');
        assert.ok(separator > 0, `invalid assignment: ${line}`);
        return [line.slice(0, separator), line.slice(separator + 1)];
      })
  );
}

function run(command, arguments_, options = {}) {
  const result = spawnSync(command, arguments_, {
    encoding: 'utf8',
    maxBuffer: 32 * 1024 * 1024,
    ...options
  });
  assert.equal(
    result.status,
    0,
    `${command} failed\n${result.stdout ?? ''}\n${result.stderr ?? ''}`
  );
  return `${result.stdout ?? ''}${result.stderr ?? ''}`;
}

function assertSubset(expected, actual, label) {
  const missing = [...expected].filter((value) => !actual.has(value)).sort();
  assert.deepEqual(missing, [], `${label}: ${missing.join(', ')}`);
}

function assertSetEqual(actual, expected, label) {
  assertSubset(expected, actual, `${label} missing`);
  assertSubset(actual, expected, `${label} unexpected`);
}

function escapeExpression(value) {
  return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
}