#!/usr/bin/env node

import { createHash } from 'node:crypto';
import { createReadStream, createWriteStream } from 'node:fs';
import {
  chmod,
  copyFile,
  mkdir,
  readFile,
  readdir,
  rm,
  stat,
  utimes,
  writeFile,
} from 'node:fs/promises';
import { execFileSync } from 'node:child_process';
import path from 'node:path';
import { pipeline } from 'node:stream/promises';
import { pathToFileURL } from 'node:url';
import { parseArgs } from 'node:util';
import { createGzip } from 'node:zlib';

import {
  CONTRACT_DOCUMENT,
  documentVersion,
} from './contract-version-utils.js';
import {
  LOCATIONS,
  migrationFiles,
  migrationVersion,
} from './flyway-migration-utils.js';

const DISTRIBUTION_NAME = 'qwen-managed-agent-reference';
const PRERELEASE_IDENTIFIER =
  '(?:0|[1-9]\\d*|[0-9A-Za-z-]*[A-Za-z-][0-9A-Za-z-]*)';
const VERSION_PATTERN = new RegExp(
  `^(?:0|[1-9]\\d*)\\.(?:0|[1-9]\\d*)\\.(?:0|[1-9]\\d*)-${PRERELEASE_IDENTIFIER}(?:\\.${PRERELEASE_IDENTIFIER})*$`,
);
const REVISION_PATTERN = /^[0-9a-f]{40}$/;
const MAIN_SOURCE_ROOT = 'packages/managed-agent-server/server/src/main/java';
const PERSISTENCE = 'com.alibaba.qwen.code.managedagent.maintenance';
const PREFLIGHT = 'com.alibaba.qwen.code.managedagent.preflight';
const WORKSPACE_RUNTIME = 'com.alibaba.qwen.code.managedagent.maintenance';

// The offline maintenance entry points, each a main class in the server jar.
// bin/qwen-managed-agent-admin runs them through Spring Boot's
// PropertiesLauncher, so the distribution ships no jar per command.
export const ADMIN_COMMANDS = Object.freeze([
  [
    'storage-registration',
    `${PERSISTENCE}.WorkspaceStorageRegistrationMain`,
    'Register, fence, restore or inspect a native Workspace storage root.',
  ],
  [
    'workspace-recovery',
    `${PERSISTENCE}.WorkspaceRecoveryMain`,
    'Capture, verify or inspect a Workspace recovery bundle.',
  ],
  [
    'workspace-migration',
    `${WORKSPACE_RUNTIME}.WorkspaceMigrationMain`,
    'Retire, prepare, promote, abort or inspect a Workspace migration.',
  ],
  [
    'maintenance-recovery',
    `${PERSISTENCE}.WorkspaceMaintenanceRecoveryMain`,
    'Inspect or reconcile a Workspace held by maintenance.',
  ],
  [
    'operator-recovery',
    `${WORKSPACE_RUNTIME}.WorkspaceRecoveryCommand`,
    'Prepare, complete or inspect a Hosted Shell lease recovery.',
  ],
  [
    'csi-registration',
    `${PERSISTENCE}.WorkspaceCsiRegistrationMain`,
    'Register or inspect a reviewed CSI Workspace registration.',
  ],
  [
    'csi-protection',
    `${WORKSPACE_RUNTIME}.WorkspaceCsiProtectionMain`,
    'Render the protection policy for a CSI registration.',
  ],
  [
    'csi-session',
    `${PERSISTENCE}.WorkspaceCsiSessionMain`,
    'Create a reviewed CSI Session offline.',
  ],
  [
    'csi-runtime',
    `${PERSISTENCE}.WorkspaceCsiRuntimeMain`,
    'Serve, prompt or retire one reviewed CSI Session Runtime.',
  ],
  [
    'csi-retirement',
    `${PERSISTENCE}.WorkspaceCsiRetirementMain`,
    'Inspect or begin a CSI Workspace retirement.',
  ],
  [
    'csi-worker-ack',
    `${PERSISTENCE}.WorkspaceCsiWorkerAckMain`,
    'Acknowledge or inspect a CSI worker publication.',
  ],
  [
    'resident-retirement',
    `${PERSISTENCE}.ResidentWorkspaceRetirementMain`,
    'Inspect, begin or delete the pod of a resident Workspace retirement.',
  ],
  [
    'resolved-configuration',
    `${PREFLIGHT}.ResolvedConfiguration`,
    'Print the Runtime state directory and auth mode the configuration resolves to.',
  ],
]);

async function sha256(file) {
  return createHash('sha256')
    .update(await readFile(file))
    .digest('hex');
}

function match(text, pattern, label) {
  const result = text.match(pattern);
  if (!result) throw new Error(`Cannot read ${label}.`);
  return result[1];
}

function uniqueMatch(text, pattern, label) {
  const results = [...text.matchAll(pattern)];
  if (results.length !== 1) {
    throw new Error(`Cannot uniquely read ${label}.`);
  }
  return results[0][1];
}

function withoutComments(text) {
  let output = '';
  let block = false;
  for (let index = 0; index < text.length; index += 1) {
    if (block) {
      if (text.startsWith('*/', index)) {
        block = false;
        index += 1;
      } else if (text[index] === '\n') {
        output += '\n';
      }
    } else if (text.startsWith('/*', index)) {
      block = true;
      index += 1;
    } else if (text.startsWith('//', index)) {
      index = text.indexOf('\n', index) - 1;
      if (index < -1) break;
    } else {
      output += text[index];
    }
  }
  return output;
}

async function pomVersion(file, artifactId) {
  const text = await readFile(file, 'utf8');
  return match(
    text,
    new RegExp(
      `<artifactId>${artifactId}</artifactId>\\s*<version>([^<]+)</version>`,
    ),
    `${artifactId} version`,
  );
}

function compareMigrationVersions(left, right) {
  const a = left.split('.').map(Number);
  const b = right.split('.').map(Number);
  for (let index = 0; index < Math.max(a.length, b.length); index += 1) {
    const difference = (a[index] ?? 0) - (b[index] ?? 0);
    if (difference !== 0) return difference;
  }
  return 0;
}

function migrationRange(root) {
  const versions = LOCATIONS.flatMap(({ dir, suffix }) =>
    migrationFiles(
      path.join(root, 'packages/managed-agent-server', ...dir),
      suffix,
    ).map((file) => migrationVersion(path.basename(file))),
  );
  if (versions.length === 0) throw new Error('No Flyway migrations found.');
  versions.sort(compareMigrationVersions);
  return { minimum: versions[0], latest: versions.at(-1) };
}

async function listFiles(directory, prefix = '') {
  const files = [];
  for (const entry of await readdir(directory, { withFileTypes: true })) {
    const relative = path.posix.join(prefix, entry.name);
    if (entry.isDirectory()) {
      files.push(
        ...(await listFiles(path.join(directory, entry.name), relative)),
      );
    } else if (entry.isFile()) {
      files.push(relative);
    }
  }
  return files.sort();
}

async function writeChecksums(directory, files, destination) {
  const lines = [];
  for (const relative of files) {
    lines.push(`${await sha256(path.join(directory, relative))}  ${relative}`);
  }
  await writeFile(destination, `${lines.join('\n')}\n`);
}

async function verifyChecksums(directory, checksumFile) {
  const lines = (await readFile(checksumFile, 'utf8')).trim().split('\n');
  for (const line of lines) {
    const found = line.match(/^([0-9a-f]{64}) {2}(.+)$/);
    if (!found) throw new Error(`Invalid checksum line: ${line}`);
    const actual = await sha256(path.join(directory, found[2]));
    if (actual !== found[1]) {
      throw new Error(`Checksum mismatch for ${found[2]}.`);
    }
  }
}

// What the archive ships from the repository, and under which name.
export const DOCUMENTATION = Object.freeze([
  ['docs/users/managed-agent-reference-deployment.md', 'docs/deployment.md'],
  [
    'docs/users/managed-agent-reference-deployment.zh-CN.md',
    'docs/deployment.zh-CN.md',
  ],
  [
    'docs/users/managed-agent-configuration-reference.md',
    'docs/configuration.md',
  ],
  [
    'docs/users/managed-agent-configuration-reference.zh-CN.md',
    'docs/configuration.zh-CN.md',
  ],
  ['packages/managed-agent-server/README.md', 'docs/server-reference.md'],
  [
    'docs/users/hosted-workspace-recovery.md',
    'docs/hosted-workspace-recovery.md',
  ],
  [
    'docs/users/hosted-workspace-recovery.zh-CN.md',
    'docs/hosted-workspace-recovery.zh-CN.md',
  ],
  ['LICENSE', 'LICENSE'],
]);

/**
 * `text`, a document shipped from `source`, with each link to another shipped
 * document pointed at the name the archive gives it: the archive renames the
 * documents, so a link written against the repository would name a file the
 * archive does not have. Every other link is left as written: a URL or an
 * anchor names no shipped path, and an absolute path is the reader's own.
 */
export function shippedDocumentLinks(
  text,
  source,
  documentation = DOCUMENTATION,
) {
  const shipped = new Map(documentation);
  const directory = path.posix.dirname(shipped.get(source));
  return text.replace(
    /\]\(([^()\s#]+)(#[^()\s]*)?\)/g,
    (link, target, anchor = '') => {
      // path.posix.join would read it as relative.
      if (target.startsWith('/')) {
        return link;
      }
      const renamed = shipped.get(
        path.posix.join(path.posix.dirname(source), target),
      );
      return renamed === undefined
        ? link
        : `](${path.posix.relative(directory, renamed)}${anchor})`;
    },
  );
}

async function copy(root, stage, source, destination) {
  const target = path.join(stage, destination);
  await mkdir(path.dirname(target), { recursive: true });
  await copyFile(path.join(root, source), target);
  return target;
}

// Every launcher reads config/application.yml the same way, so the
// operator's configuration cannot depend on which one they ran.
const LAUNCHER_PREAMBLE = `#!/usr/bin/env sh
set -eu
base="$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)"
config="$base/config/application.yml"
if [ -f "$config" ] && [ -z "\${SPRING_CONFIG_ADDITIONAL_LOCATION:-}" ]; then
  SPRING_CONFIG_ADDITIONAL_LOCATION="file:\${config}"
  export SPRING_CONFIG_ADDITIONAL_LOCATION
fi
`;

function launcher(jar) {
  return `${LAUNCHER_PREAMBLE}exec java -jar "$base/lib/${jar}" "$@"
`;
}

function adminLauncher(jar) {
  const width = Math.max(...ADMIN_COMMANDS.map(([name]) => name.length));
  const usage = ADMIN_COMMANDS.map(
    ([name, , summary]) => `  ${name.padEnd(width)}  ${summary}`,
  );
  const cases = ADMIN_COMMANDS.map(
    ([name, main]) => `  ${name}) main=${main} ;;`,
  );
  return `${LAUNCHER_PREAMBLE}usage() {
  cat <<'USAGE'
Usage: qwen-managed-agent-admin <command> [arguments]

Offline maintenance commands. Run one without arguments for its usage.

${usage.join('\n')}
USAGE
}
case "\${1:-}" in
${cases.join('\n')}
  ''|help|-h|--help) usage; exit 0 ;;
  *) echo "qwen-managed-agent-admin: unknown command: $1" >&2; usage >&2; exit 2 ;;
esac
shift
exec java "-Dloader.main=$main" -cp "$base/lib/${jar}" \\
  org.springframework.boot.loader.launch.PropertiesLauncher "$@"
`;
}

// A renamed or removed entry point fails packaging instead of shipping a
// command that cannot start.
async function verifyAdminCommands(root) {
  for (const [name, main] of ADMIN_COMMANDS) {
    const source = path.join(
      root,
      MAIN_SOURCE_ROOT,
      `${main.replaceAll('.', '/')}.java`,
    );
    let text;
    try {
      text = await readFile(source, 'utf8');
    } catch {
      throw new Error(`Admin command ${name} names a missing class ${main}.`);
    }
    if (!/public static void main\(String/.test(text)) {
      throw new Error(`Admin command ${name}: ${main} has no main method.`);
    }
  }
}

async function createDeterministicArchive(output, rootName, stage, archive) {
  const files = await listFiles(stage);
  const timestamp = new Date(0);
  for (const relative of files) {
    const file = path.join(stage, relative);
    await chmod(file, relative.startsWith('bin/') ? 0o755 : 0o644);
    await utimes(file, timestamp, timestamp);
  }

  const tarVersion = execFileSync('tar', ['--version'], { encoding: 'utf8' });
  const ownershipArguments = tarVersion.toLowerCase().includes('bsdtar')
    ? ['--uid', '0', '--gid', '0', '--uname', 'root', '--gname', 'root']
    : ['--owner=0', '--group=0'];
  const tarFile = `${archive}.tmp`;
  try {
    execFileSync(
      'tar',
      [
        '-c',
        '--format=ustar',
        ...ownershipArguments,
        '-f',
        tarFile,
        '-C',
        output,
        ...files.map((file) => path.posix.join(rootName, file)),
      ],
      { stdio: 'inherit' },
    );
    await pipeline(
      createReadStream(tarFile),
      createGzip({ level: 9, mtime: 0 }),
      createWriteStream(archive),
    );
  } finally {
    await rm(tarFile, { force: true });
  }
}

export async function createReferenceDistribution({
  repoRoot,
  outputDirectory,
  version,
  sourceRevision,
}) {
  if (!VERSION_PATTERN.test(version)) {
    throw new Error(
      'Version must be an experimental SemVer with a prerelease.',
    );
  }
  if (!REVISION_PATTERN.test(sourceRevision)) {
    throw new Error(
      'Source revision must be a lowercase 40-character Git SHA.',
    );
  }

  const root = path.resolve(repoRoot);
  const output = path.resolve(outputDirectory);
  const rootName = `${DISTRIBUTION_NAME}-${version}`;
  const stage = path.join(output, rootName);
  const archive = path.join(output, `${rootName}.tar.gz`);
  const detachedManifest = path.join(output, `${rootName}.manifest.json`);
  const detachedChecksums = path.join(output, 'SHA256SUMS.txt');

  await mkdir(output, { recursive: true });
  await Promise.all([
    rm(stage, { recursive: true, force: true }),
    rm(archive, { force: true }),
    rm(detachedManifest, { force: true }),
    rm(detachedChecksums, { force: true }),
  ]);
  await mkdir(stage, { recursive: true });

  const serverPom = path.join(root, 'packages/managed-agent-server/pom.xml');
  const serverVersion = await pomVersion(
    serverPom,
    'qwen-managed-agent-parent',
  );
  const packageJson = JSON.parse(
    await readFile(path.join(root, 'package.json')),
  );
  const publicApi = documentVersion(
    await readFile(path.join(root, CONTRACT_DOCUMENT), 'utf8'),
  );
  if (publicApi.error) {
    throw new Error(`${CONTRACT_DOCUMENT} ${publicApi.error}.`);
  }
  const capabilities = await readFile(
    path.join(root, 'packages/cli/src/serve/capabilities.ts'),
    'utf8',
  );
  const brokerClient = await readFile(
    path.join(root, 'packages/cli/src/serve/managed-runtime-broker-client.ts'),
    'utf8',
  );
  const applicationConfiguration = await readFile(
    path.join(
      root,
      'packages/managed-agent-server/server/src/main/resources/application.yml',
    ),
    'utf8',
  );
  const hostedHarnessContract = await readFile(
    path.join(root, 'packages/cli/src/serve/hosted-harness-contract.ts'),
    'utf8',
  );
  const serverPomText = await readFile(serverPom, 'utf8');
  const schema = migrationRange(root);
  const publicShellDefault = uniqueMatch(
    applicationConfiguration,
    /^[ \t]*workspace-shell-enabled:[ \t]*'\$\{QWEN_MANAGED_AGENT_WORKSPACE_SHELL_ENABLED:([^}\r\n]+)\}'[ \t]*$/gm,
    'public Shell default',
  );
  if (publicShellDefault !== 'false') {
    throw new Error('Public Shell must default to false.');
  }
  const lifecycleV2ProfilesDefault = uniqueMatch(
    withoutComments(hostedHarnessContract),
    /^[ \t]*process\.env\[HOSTED_LIFECYCLE_V2_PROFILES_ENV\][ \t]*\?\?[ \t]*'([^'\r\n]*)'[ \t]*$/gm,
    'Hosted lifecycle v2 profile default',
  );
  if (lifecycleV2ProfilesDefault !== 'hosted-workspace-shell/1') {
    throw new Error(
      'Hosted lifecycle v2 profiles must default to canonical Shell only.',
    );
  }

  const serverJar = `qwen-managed-agent-server-${serverVersion}.jar`;
  const recoveryJar = `qwen-managed-agent-server-${serverVersion}-operator-recovery.jar`;
  const componentSources = [
    {
      name: 'qwen-managed-agent-server',
      version: serverVersion,
      source: `packages/managed-agent-server/server/target/qwen-managed-agent-server-${serverVersion}-exec.jar`,
      destination: `lib/${serverJar}`,
      role: 'control-plane',
    },
    {
      name: 'qwen-managed-agent-operator-recovery',
      version: serverVersion,
      source: `packages/managed-agent-server/server/target/${recoveryJar}`,
      destination: `lib/${recoveryJar}`,
      role: 'operator-recovery',
    },
  ];

  for (const component of componentSources) {
    await copy(root, stage, component.source, component.destination);
  }
  await copy(
    root,
    stage,
    'packages/managed-agent-server/distribution/config/application.example.yml',
    'config/application.example.yml',
  );
  const systemdFiles = [
    'qwen-hosted-harness.service',
    'qwen-managed-agent-server.service',
    'harness.env.example',
    'server.env.example',
    'README.md',
    'README.zh-CN.md',
  ];
  for (const file of systemdFiles) {
    await copy(
      root,
      stage,
      `packages/managed-agent-server/distribution/systemd/${file}`,
      `systemd/${file}`,
    );
  }
  for (const file of [
    'qwen-managed-agent-bootstrap',
    'qwen-managed-agent-bootstrap.mjs',
    'qwen-managed-agent-upgrade',
    'qwen-managed-agent-upgrade.mjs',
  ]) {
    await copy(
      root,
      stage,
      `packages/managed-agent-server/distribution/bin/${file}`,
      `bin/${file}`,
    );
  }
  for (const [source, destination] of DOCUMENTATION) {
    if (!source.endsWith('.md')) {
      await copy(root, stage, source, destination);
      continue;
    }
    const target = path.join(stage, destination);
    await mkdir(path.dirname(target), { recursive: true });
    await writeFile(
      target,
      shippedDocumentLinks(
        await readFile(path.join(root, source), 'utf8'),
        source,
      ),
    );
  }

  const serverLauncher = path.join(stage, 'bin/qwen-managed-agent-server');
  const recoveryLauncher = path.join(
    stage,
    'bin/qwen-managed-agent-operator-recovery',
  );
  const adminCommand = path.join(stage, 'bin/qwen-managed-agent-admin');
  await verifyAdminCommands(root);
  await mkdir(path.dirname(serverLauncher), { recursive: true });
  await Promise.all([
    writeFile(serverLauncher, launcher(serverJar)),
    writeFile(recoveryLauncher, launcher(recoveryJar)),
    writeFile(adminCommand, adminLauncher(serverJar)),
  ]);
  await Promise.all([
    chmod(serverLauncher, 0o755),
    chmod(recoveryLauncher, 0o755),
    chmod(adminCommand, 0o755),
  ]);

  const components = [];
  for (const component of componentSources) {
    const file = path.join(stage, component.destination);
    components.push({
      name: component.name,
      version: component.version,
      role: component.role,
      path: component.destination,
      size: (await stat(file)).size,
      sha256: await sha256(file),
    });
  }

  const manifest = {
    schemaVersion: 1,
    name: DISTRIBUTION_NAME,
    version,
    status: 'experimental',
    source: {
      repository: 'https://github.com/QwenLM/qwen-code',
      revision: sourceRevision,
    },
    components,
    compatibility: {
      qwenCodeVersion: packageJson.version,
      hostedHarnessProfile: 'hosted-harness',
      serveProtocolVersion: match(
        capabilities,
        /SERVE_PROTOCOL_VERSION = '([^']+)'/,
        'Serve protocol version',
      ),
      managedRuntimeBrokerProtocolVersion: Number(
        match(
          brokerClient,
          /MANAGED_RUNTIME_BROKER_PROTOCOL_VERSION = (\d+)/,
          'Runtime Broker protocol version',
        ),
      ),
      managedAgentPublicApiVersion: publicApi.version,
      java: `>=${match(
        serverPomText,
        /<java\.version>([^<]+)<\/java\.version>/,
        'Java version',
      )}`,
      node: packageJson.engines.node,
      databaseSchema: {
        migrationTool: 'Flyway',
        minimumMigrationVersion: schema.minimum,
        latestMigrationVersion: schema.latest,
      },
    },
    referenceTopology: {
      operatingSystem: 'Linux',
      javaControlPlanes: 1,
      hostTrust: 'trusted-single-host',
      runtimeBroker: 'embedded',
      database: 'MySQL 8',
    },
    defaults: {
      publicShellEnabled: publicShellDefault === 'true',
      v2CapabilityAdvertisementEnabled: lifecycleV2ProfilesDefault !== '',
    },
    documentation: {
      deployment: 'docs/deployment.md',
      deploymentChinese: 'docs/deployment.zh-CN.md',
      configuration: 'docs/configuration.md',
      configurationChinese: 'docs/configuration.zh-CN.md',
      serverReference: 'docs/server-reference.md',
      operatorRecovery: 'docs/hosted-workspace-recovery.md',
    },
  };
  const manifestFile = path.join(stage, 'manifest.json');
  await writeFile(manifestFile, `${JSON.stringify(manifest, null, 2)}\n`);

  const payload = await listFiles(stage);
  await writeChecksums(stage, payload, path.join(stage, 'SHA256SUMS'));
  await verifyChecksums(stage, path.join(stage, 'SHA256SUMS'));
  await copyFile(manifestFile, detachedManifest);

  await createDeterministicArchive(output, rootName, stage, archive);
  const entries = execFileSync('tar', ['-tzf', archive], { encoding: 'utf8' });
  for (const required of [
    `${rootName}/manifest.json`,
    `${rootName}/SHA256SUMS`,
    `${rootName}/bin/qwen-managed-agent-server`,
    `${rootName}/bin/qwen-managed-agent-operator-recovery`,
    `${rootName}/bin/qwen-managed-agent-admin`,
    // The two launchers exec these modules.
    `${rootName}/bin/qwen-managed-agent-bootstrap`,
    `${rootName}/bin/qwen-managed-agent-bootstrap.mjs`,
    `${rootName}/bin/qwen-managed-agent-upgrade`,
    `${rootName}/bin/qwen-managed-agent-upgrade.mjs`,
    `${rootName}/config/application.example.yml`,
    ...systemdFiles.map((file) => `${rootName}/systemd/${file}`),
  ]) {
    if (!entries.split('\n').includes(required)) {
      throw new Error(`Archive is missing ${required}.`);
    }
  }

  await writeChecksums(
    output,
    [path.basename(archive), path.basename(detachedManifest)],
    detachedChecksums,
  );
  await verifyChecksums(output, detachedChecksums);
  await rm(stage, { recursive: true, force: true });

  return { archive, detachedManifest, detachedChecksums, manifest };
}

async function main() {
  const { values } = parseArgs({
    options: {
      version: { type: 'string' },
      'source-revision': { type: 'string' },
      output: { type: 'string' },
    },
  });
  if (!values.version || !values['source-revision'] || !values.output) {
    throw new Error(
      'Usage: package-managed-agent-reference.mjs --version <experimental-semver> --source-revision <sha> --output <directory>',
    );
  }
  const result = await createReferenceDistribution({
    repoRoot: process.cwd(),
    outputDirectory: values.output,
    version: values.version,
    sourceRevision: values['source-revision'],
  });
  console.log(result.archive);
  console.log(result.detachedManifest);
  console.log(result.detachedChecksums);
}

if (import.meta.url === pathToFileURL(process.argv[1]).href) {
  main().catch((error) => {
    console.error(error instanceof Error ? error.message : error);
    process.exitCode = 1;
  });
}