// PRE Web GUI — tinygrad LoRA training backend
// Manages fine-tuning jobs on eGPU using tinygrad. Consumes training exports
// from ~/.pre/training/exports/ and produces LoRA adapters.

const { spawn } = require('child_process');
const fs = require('fs');
const path = require('path');
const os = require('os');
const { gpuDetector } = require('./gpu-detect');
const { getPythonPath, isTinygradInstalled } = require('./tinygrad');

const PRE_DIR = path.join(os.homedir(), '.pre');
const ADAPTERS_DIR = path.join(PRE_DIR, 'training', 'adapters');
const JOBS_FILE = path.join(PRE_DIR, 'training', 'jobs.json');
const WORKER_SCRIPT = path.join(__dirname, 'tinygrad-train-worker.py');

let _jobs = new Map();
let _jobSeq = 0;
let _activeProcess = null;
let _broadcastWS = null;

// ── Persistence ──────────────────────────────────────────────────────────

function _loadJobs() {
  try {
    const data = JSON.parse(fs.readFileSync(JOBS_FILE, 'utf-8'));
    _jobSeq = data.seq || 0;
    for (const j of (data.jobs || [])) {
      _jobs.set(j.id, j);
    }
  } catch {}
}

function _saveJobs() {
  try {
    const dir = path.dirname(JOBS_FILE);
    fs.mkdirSync(dir, { recursive: true });
    const jobs = [..._jobs.values()];
    fs.writeFileSync(JOBS_FILE, JSON.stringify({ jobs, seq: _jobSeq }, null, 2));
  } catch {}
}

// ── Job Management ───────────────────────────────────────────────────────

/**
 * Start a new training job.
 * @param {Object} config
 * @returns {{ id: string, status: string }}
 */
function startTraining(config) {
  if (!config.base_model) throw new Error('base_model path is required');
  if (!config.training_data) throw new Error('training_data path is required');
  if (!fs.existsSync(config.base_model)) throw new Error(`Base model not found: ${config.base_model}`);
  if (!fs.existsSync(config.training_data)) throw new Error(`Training data not found: ${config.training_data}`);

  const gpuState = gpuDetector.getState();
  if (!gpuState || !gpuState.available) {
    throw new Error('GPU not available. Connect an eGPU before starting training.');
  }

  if (!isTinygradInstalled()) {
    throw new Error('tinygrad not installed. Run: pip install tinygrad');
  }

  const id = `train_${++_jobSeq}`;
  const name = config.name || `finetune-${new Date().toISOString().slice(0, 10)}`;
  const outputDir = path.join(ADAPTERS_DIR, name);

  const job = {
    id,
    name,
    status: 'queued',
    config: {
      base_model: config.base_model,
      training_data: config.training_data,
      format: config.format || 'chatml',
      output_dir: outputDir,
      lora_rank: parseInt(config.lora_rank, 10) || 16,
      lora_alpha: parseInt(config.lora_alpha, 10) || 32,
      learning_rate: parseFloat(config.learning_rate) || 2e-4,
      epochs: parseInt(config.epochs, 10) || 3,
      batch_size: parseInt(config.batch_size, 10) || 4,
      max_seq_len: parseInt(config.max_seq_len, 10) || 2048,
      quantization: config.quantization || 'int8',
    },
    progress: { epoch: 0, step: 0, loss: null, elapsed_ms: 0 },
    createdAt: Date.now(),
    startedAt: null,
    completedAt: null,
    error: null,
    adapter_path: null,
  };

  _jobs.set(id, job);
  _saveJobs();

  // Execute immediately if no other job is running
  if (!_activeProcess) {
    _executeJob(job);
  }

  return { id: job.id, status: job.status };
}

/**
 * Execute a training job as a child process.
 */
function _executeJob(job) {
  const python = getPythonPath();
  if (!python) {
    job.status = 'failed';
    job.error = 'Python 3.11+ not found';
    job.completedAt = Date.now();
    _saveJobs();
    return;
  }

  job.status = 'running';
  job.startedAt = Date.now();
  _saveJobs();
  _notify(job, 'running');

  const args = [
    WORKER_SCRIPT,
    '--base-model', job.config.base_model,
    '--training-data', job.config.training_data,
    '--format', job.config.format,
    '--output-dir', job.config.output_dir,
    '--lora-rank', String(job.config.lora_rank),
    '--lora-alpha', String(job.config.lora_alpha),
    '--lr', String(job.config.learning_rate),
    '--epochs', String(job.config.epochs),
    '--batch-size', String(job.config.batch_size),
    '--max-seq-len', String(job.config.max_seq_len),
    '--quantize', job.config.quantization,
  ];

  const env = { ...process.env, DEV: 'NV' };

  _activeProcess = spawn(python, args, {
    stdio: ['ignore', 'pipe', 'pipe'],
    env,
    detached: false,
  });

  let stderrBuf = '';

  _activeProcess.stdout.on('data', (data) => {
    // Parse progress lines (JSON per line)
    const lines = data.toString().split('\n').filter(l => l.trim());
    for (const line of lines) {
      try {
        const progress = JSON.parse(line);
        if (progress.epoch !== undefined) {
          job.progress = {
            epoch: progress.epoch,
            step: progress.step || 0,
            loss: progress.loss ?? null,
            elapsed_ms: Date.now() - job.startedAt,
          };
          _saveJobs();
          _notify(job, 'progress');
        }
        if (progress.status === 'completed') {
          job.adapter_path = progress.adapter_path || job.config.output_dir;
        }
      } catch {
        // Non-JSON output — ignore or log
      }
    }
  });

  _activeProcess.stderr.on('data', (data) => {
    stderrBuf += data.toString();
    // Keep only last 2KB of stderr
    if (stderrBuf.length > 2048) stderrBuf = stderrBuf.slice(-2048);
  });

  _activeProcess.on('exit', (code) => {
    _activeProcess = null;
    if (code === 0) {
      job.status = 'completed';
      job.completedAt = Date.now();
      job.progress.elapsed_ms = job.completedAt - job.startedAt;
      if (!job.adapter_path) job.adapter_path = job.config.output_dir;
    } else if (job.status !== 'cancelled') {
      job.status = 'failed';
      job.completedAt = Date.now();
      job.error = stderrBuf.trim().split('\n').pop() || `Process exited with code ${code}`;
    }
    _saveJobs();
    _notify(job, job.status);
  });

  _activeProcess.on('error', (err) => {
    _activeProcess = null;
    job.status = 'failed';
    job.completedAt = Date.now();
    job.error = err.message;
    _saveJobs();
    _notify(job, 'failed');
  });
}

/**
 * Get a training job by ID.
 */
function getJob(id) {
  return _jobs.get(id) || null;
}

/**
 * List all training jobs (newest first).
 */
function listJobs() {
  return [..._jobs.values()].sort((a, b) => b.createdAt - a.createdAt);
}

/**
 * Cancel a running or queued training job.
 */
function cancelJob(id) {
  const job = _jobs.get(id);
  if (!job) return { success: false, error: 'Job not found' };
  if (job.status !== 'running' && job.status !== 'queued') {
    return { success: false, error: `Cannot cancel job with status: ${job.status}` };
  }
  job.status = 'cancelled';
  job.completedAt = Date.now();
  if (_activeProcess && job.status === 'cancelled') {
    try { _activeProcess.kill('SIGTERM'); } catch {}
  }
  _saveJobs();
  _notify(job, 'cancelled');
  return { success: true };
}

/**
 * List saved LoRA adapters.
 */
function listAdapters() {
  try {
    fs.mkdirSync(ADAPTERS_DIR, { recursive: true });
    const entries = fs.readdirSync(ADAPTERS_DIR, { withFileTypes: true });
    return entries.filter(e => e.isDirectory()).map(e => {
      const adapterDir = path.join(ADAPTERS_DIR, e.name);
      let metadata = {};
      try {
        metadata = JSON.parse(fs.readFileSync(path.join(adapterDir, 'metadata.json'), 'utf-8'));
      } catch {}
      return {
        name: e.name,
        path: adapterDir,
        created: metadata.created || null,
        base_model: metadata.base_model || null,
        training_samples: metadata.training_samples || null,
        final_loss: metadata.final_loss || null,
      };
    });
  } catch {
    return [];
  }
}

// ── Notifications ────────────────────────────────────────────────────────

function _notify(job, event) {
  if (!_broadcastWS) return;
  _broadcastWS({
    type: 'training_job',
    event,
    id: job.id,
    name: job.name,
    status: job.status,
    progress: job.progress,
    error: job.error,
    adapter_path: job.adapter_path,
  });
}

function setBroadcast(fn) {
  _broadcastWS = fn;
}

/**
 * Initialize — load persisted jobs.
 */
function init() {
  _loadJobs();
}

/**
 * Graceful shutdown.
 */
function shutdown() {
  if (_activeProcess) {
    try { _activeProcess.kill('SIGTERM'); } catch {}
    _activeProcess = null;
  }
  _saveJobs();
}

module.exports = {
  startTraining,
  getJob,
  listJobs,
  cancelJob,
  listAdapters,
  setBroadcast,
  init,
  shutdown,
};
