'use strict';

const { spawn } = require('node:child_process');
const path = require('node:path');
const fs = require('node:fs');

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const SUPPORTED_BASES = {
  'llama3.1-8b':    { ollamaName: 'llama3.1:8b-instruct-q4_K_M', huggingFace: 'unsloth/Meta-Llama-3.1-8B-Instruct' },
  'phi-4':          { ollamaName: 'phi4:latest',                   huggingFace: 'unsloth/Phi-4' },
  'qwen2.5-7b':    { ollamaName: 'qwen2.5:7b-instruct-q4_K_M',   huggingFace: 'unsloth/Qwen2.5-7B-Instruct' },
  'gemma4-4b':      { ollamaName: 'gemma4:e4b',                    huggingFace: 'unsloth/gemma-4-E4B-it' },
};

const DEFAULT_HYPERPARAMS = {
  epochs:       3,
  lr:           2e-4,
  loraRank:     16,
  loraAlpha:    32,
  batchSize:    4,
  maxSeqLength: 2048,
};

const TRAIN_SCRIPT = path.join(__dirname, 'train.py');

// ---------------------------------------------------------------------------
// Command builder (exported for testing)
// ---------------------------------------------------------------------------

/**
 * Build the CLI args array for the Python training script.
 * @param {Object} opts - Same shape as trainModel options
 * @returns {string[]} args array (without 'python3' prefix)
 */
function buildTrainCommand(opts) {
  const hp = { ...DEFAULT_HYPERPARAMS, ...opts.hyperparams };
  const outputName = opts.outputName || `walle-${opts.baseModel}-v1`;

  const args = [
    TRAIN_SCRIPT,
    '--base',           opts.baseModel,
    '--data',           opts.dataPath,
    '--output',         opts.outputDir,
    '--name',           outputName,
    '--epochs',         String(hp.epochs),
    '--lr',             String(hp.lr),
    '--lora-rank',      String(hp.loraRank),
    '--lora-alpha',     String(hp.loraAlpha),
    '--batch-size',     String(hp.batchSize),
    '--max-seq-length', String(hp.maxSeqLength),
  ];

  if (opts.format) {
    args.push('--format', opts.format);
  }

  return args;
}

// ---------------------------------------------------------------------------
// Main entry point
// ---------------------------------------------------------------------------

/**
 * Run fine-tuning for a base model.
 * @param {Object} options
 * @param {string} options.baseModel      - Base model name: 'llama3.1-8b', 'phi-4', 'qwen2.5-7b', 'gemma4-4b'
 * @param {string} options.dataPath       - Path to JSONL training data
 * @param {string} options.outputDir      - Where to save the model
 * @param {string} [options.outputName]   - Model name (default: 'walle-{base}-v1')
 * @param {Object} [options.hyperparams]  - Override training hyperparams
 * @param {Function} [options.onProgress] - Progress callback ({ stage, progress, message })
 * @returns {Promise<{ success: boolean, modelPath: string, metrics: Object, error?: string }>}
 */
const DEFAULT_TIMEOUT = 7200000; // 2 hours
const SAFE_NAME_RE = /^[a-zA-Z0-9._-]+$/;
const MAX_STDERR = 10240; // 10 KB

async function trainModel(options) {
  // --- Validate --------------------------------------------------------
  if (!options || !options.dataPath) {
    throw new Error('dataPath is required');
  }
  if (!options.baseModel) {
    throw new Error('baseModel is required');
  }
  if (!options.outputDir) {
    throw new Error('outputDir is required');
  }
  if (!SUPPORTED_BASES[options.baseModel]) {
    throw new Error(
      `Unsupported base model "${options.baseModel}". ` +
      `Supported: ${Object.keys(SUPPORTED_BASES).join(', ')}`
    );
  }
  if (options.outputName && !SAFE_NAME_RE.test(options.outputName)) {
    throw new Error(
      `Invalid outputName "${options.outputName}". Must match ${SAFE_NAME_RE}`
    );
  }
  if (!fs.existsSync(options.dataPath)) {
    throw new Error(`Training data not found: ${options.dataPath}`);
  }

  // --- Build command ---------------------------------------------------
  const args = buildTrainCommand(options);
  const onProgress = options.onProgress || (() => {});

  // --- Spawn Python subprocess -----------------------------------------
  const timeout = options.timeout || DEFAULT_TIMEOUT;

  return new Promise((resolve, reject) => {
    const proc = spawn('python3', args, {
      stdio: ['ignore', 'pipe', 'pipe'],
      env: { ...process.env },
      timeout,
    });

    let stderr = '';
    let lastResult = null;
    let stdoutBuffer = '';

    proc.stdout.on('data', (chunk) => {
      stdoutBuffer += chunk.toString();
      const lines = stdoutBuffer.split('\n');
      stdoutBuffer = lines.pop(); // keep incomplete trailing line
      for (const line of lines) {
        if (!line.trim()) continue;
        try {
          const msg = JSON.parse(line);
          if (msg.stage) {
            onProgress(msg);
          }
          // Keep the last JSON object as the final result
          lastResult = msg;
        } catch {
          // Non-JSON output, ignore
        }
      }
    });

    proc.stderr.on('data', (chunk) => {
      stderr += chunk.toString();
      if (stderr.length > MAX_STDERR) stderr = stderr.slice(-MAX_STDERR);
    });

    proc.on('error', (err) => {
      reject(new Error(`Failed to spawn python3: ${err.message}`));
    });

    proc.on('close', (code) => {
      // Parse any remaining buffered stdout
      if (stdoutBuffer.trim()) {
        try {
          const msg = JSON.parse(stdoutBuffer);
          if (msg.stage) onProgress(msg);
          lastResult = msg;
        } catch {
          // Non-JSON output, ignore
        }
      }

      if (code !== 0) {
        const errMsg = stderr.trim() || (lastResult && lastResult.message) || `Process exited with code ${code}`;
        resolve({
          success: false,
          modelPath: '',
          metrics: {},
          error: errMsg,
        });
        return;
      }

      if (lastResult && lastResult.success) {
        resolve({
          success: true,
          modelPath: lastResult.model_path || '',
          modelName: lastResult.model_name || '',
          metrics: {
            samplesTrained: lastResult.samples_trained,
            epochs: lastResult.epochs,
            ...(lastResult.final_loss != null && { finalLoss: lastResult.final_loss }),
          },
        });
      } else {
        resolve({
          success: false,
          modelPath: '',
          metrics: {},
          error: lastResult?.message || 'Training completed but no success result received',
        });
      }
    });
  });
}

// ---------------------------------------------------------------------------
// Version helper
// ---------------------------------------------------------------------------

/**
 * Get the next version number for a model based on existing training runs.
 * @param {Object|null} brain - Brain module (for querying training_runs)
 * @param {string} baseModel - Base model key (e.g. 'phi-4')
 * @returns {number} Next version number (starting from 1)
 */
function getNextVersion(brain, baseModel) {
  if (!brain) return 1;
  const runs = brain.getTrainingRuns({});
  const versions = runs
    .filter(r => r.base_model === baseModel)
    .map(r => {
      const match = r.model_name.match(/-v(\d+)$/);
      return match ? parseInt(match[1]) : 0;
    });
  return versions.length > 0 ? Math.max(...versions) + 1 : 1;
}

module.exports = {
  trainModel,
  buildTrainCommand,
  getNextVersion,
  SUPPORTED_BASES,
  DEFAULT_HYPERPARAMS,
  DEFAULT_TIMEOUT,
  SAFE_NAME_RE,
  MAX_STDERR,
  TRAIN_SCRIPT,
};
