// Tests for web/src/tinygrad-train.js — tinygrad LoRA training backend

const path = require('path');
const fs = require('fs');
const os = require('os');
const cp = require('child_process');
const EventEmitter = require('events');

// Mock os.homedir before any require that uses it
const mockHome = path.join(os.tmpdir(), `pre-train-test-${process.pid}`);
jest.spyOn(os, 'homedir').mockReturnValue(mockHome);

// Mock execSync (used by getPythonPath / isTinygradInstalled)
const mockExecSync = jest.spyOn(cp, 'execSync');

// Mock spawn (used by _executeJob)
const mockSpawn = jest.spyOn(cp, 'spawn');

// Mock gpu-detect module
jest.mock('../src/gpu-detect', () => {
  const emitter = new (require('events').EventEmitter)();
  emitter.getState = jest.fn(() => ({ available: true, vendor: 'nvidia', name: 'RTX 4090', vram_gb: 24 }));
  return { gpuDetector: emitter };
});

// Mock tinygrad module (for isTinygradInstalled / getPythonPath)
jest.mock('../src/tinygrad', () => ({
  getPythonPath: jest.fn(() => '/usr/bin/python3'),
  isTinygradInstalled: jest.fn(() => true),
}));

const tinygradTrain = require('../src/tinygrad-train');
const { gpuDetector } = require('../src/gpu-detect');
const { getPythonPath, isTinygradInstalled } = require('../src/tinygrad');

const PRE_DIR = path.join(mockHome, '.pre');
const TRAINING_DIR = path.join(PRE_DIR, 'training');
const ADAPTERS_DIR = path.join(TRAINING_DIR, 'adapters');
const JOBS_FILE = path.join(TRAINING_DIR, 'jobs.json');

// Test fixtures
const MODEL_PATH = path.join(mockHome, 'test-model.safetensors');
const DATA_PATH = path.join(mockHome, 'training-data.jsonl');

// Helper: create a mock child process
function createMockProcess() {
  const proc = new EventEmitter();
  proc.stdout = new EventEmitter();
  proc.stderr = new EventEmitter();
  proc.kill = jest.fn();
  proc.pid = 12345;
  return proc;
}

beforeAll(() => {
  fs.mkdirSync(TRAINING_DIR, { recursive: true });
  fs.mkdirSync(ADAPTERS_DIR, { recursive: true });
  // Create dummy model file
  fs.writeFileSync(MODEL_PATH, 'dummy model data');
  // Create dummy training data
  const trainingData = [
    '{"messages":[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]}',
    '{"messages":[{"role":"user","content":"How are you?"},{"role":"assistant","content":"Good, thanks!"}]}',
  ].join('\n');
  fs.writeFileSync(DATA_PATH, trainingData);
  // Initialize the module
  tinygradTrain.init();
});

afterAll(() => {
  tinygradTrain.shutdown();
  fs.rmSync(mockHome, { recursive: true, force: true });
});

beforeEach(() => {
  // Reset module state — kill any lingering _activeProcess
  tinygradTrain.shutdown();
  mockExecSync.mockReset();
  mockSpawn.mockReset();
  getPythonPath.mockReturnValue('/usr/bin/python3');
  isTinygradInstalled.mockReturnValue(true);
  gpuDetector.getState.mockReturnValue({ available: true, vendor: 'nvidia', name: 'RTX 4090', vram_gb: 24 });
});

// ── Validation ──────────────────────────────────────────────────────────

describe('startTraining — validation', () => {
  test('throws if base_model missing', () => {
    expect(() => tinygradTrain.startTraining({
      training_data: DATA_PATH,
    })).toThrow('base_model path is required');
  });

  test('throws if training_data missing', () => {
    expect(() => tinygradTrain.startTraining({
      base_model: MODEL_PATH,
    })).toThrow('training_data path is required');
  });

  test('throws if base_model file does not exist', () => {
    expect(() => tinygradTrain.startTraining({
      base_model: '/nonexistent/model.safetensors',
      training_data: DATA_PATH,
    })).toThrow('Base model not found');
  });

  test('throws if training_data file does not exist', () => {
    expect(() => tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: '/nonexistent/data.jsonl',
    })).toThrow('Training data not found');
  });

  test('throws if GPU not available', () => {
    gpuDetector.getState.mockReturnValue({ available: false, vendor: null });
    expect(() => tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
    })).toThrow('GPU not available');
  });

  test('throws if tinygrad not installed', () => {
    isTinygradInstalled.mockReturnValue(false);
    expect(() => tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
    })).toThrow('tinygrad not installed');
  });
});

// ── Job Creation ────────────────────────────────────────────────────────

describe('startTraining — job creation', () => {
  test('creates a job with correct defaults and starts it immediately', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
    });

    expect(result.id).toMatch(/^train_\d+$/);
    // Job starts immediately when no other job is running
    expect(result.status).toBe('running');

    const job = tinygradTrain.getJob(result.id);
    expect(job).not.toBeNull();
    expect(job.config.lora_rank).toBe(16);
    expect(job.config.lora_alpha).toBe(32);
    expect(job.config.learning_rate).toBe(2e-4);
    expect(job.config.epochs).toBe(3);
    expect(job.config.batch_size).toBe(4);
    expect(job.config.max_seq_len).toBe(2048);
    expect(job.config.quantization).toBe('int8');
    expect(job.config.format).toBe('chatml');

    // Clean up
    mockProc.emit('exit', 0);
  });

  test('respects custom config values', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'custom-finetune',
      format: 'alpaca',
      lora_rank: '32',
      lora_alpha: '64',
      learning_rate: '1e-3',
      epochs: '5',
      batch_size: '8',
      max_seq_len: '4096',
      quantization: 'nf4',
    });

    const job = tinygradTrain.getJob(result.id);
    expect(job.name).toBe('custom-finetune');
    expect(job.config.format).toBe('alpaca');
    expect(job.config.lora_rank).toBe(32);
    expect(job.config.lora_alpha).toBe(64);
    expect(job.config.learning_rate).toBe(1e-3);
    expect(job.config.epochs).toBe(5);
    expect(job.config.batch_size).toBe(8);
    expect(job.config.max_seq_len).toBe(4096);
    expect(job.config.quantization).toBe('nf4');

    mockProc.emit('exit', 0);
  });
});

// ── Job Execution ───────────────────────────────────────────────────────

describe('_executeJob — process management', () => {
  test('spawns Python with correct arguments', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'spawn-test',
    });

    expect(mockSpawn).toHaveBeenCalledTimes(1);
    const [pythonPath, args, opts] = mockSpawn.mock.calls[0];
    expect(pythonPath).toBe('/usr/bin/python3');
    expect(args).toContain('--base-model');
    expect(args).toContain(MODEL_PATH);
    expect(args).toContain('--training-data');
    expect(args).toContain(DATA_PATH);
    expect(args).toContain('--lora-rank');
    expect(args).toContain('16');
    expect(args).toContain('--epochs');
    expect(args).toContain('3');
    expect(opts.env.DEV).toBe('NV');

    mockProc.emit('exit', 0);
  });

  test('marks job completed on exit code 0', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'complete-test',
    });

    expect(result.status).toBe('running');
    mockProc.emit('exit', 0);

    const job = tinygradTrain.getJob(result.id);
    expect(job.status).toBe('completed');
    expect(job.completedAt).toBeGreaterThan(0);
  });

  test('marks job failed on non-zero exit code', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'fail-test',
    });

    mockProc.stderr.emit('data', Buffer.from('RuntimeError: CUDA out of memory'));
    mockProc.emit('exit', 1);

    const job = tinygradTrain.getJob(result.id);
    expect(job.status).toBe('failed');
    expect(job.error).toContain('CUDA out of memory');
  });

  test('marks job failed on process error', () => {
    const mockProc = createMockProcess();
    // Add an error listener to prevent unhandled error from EventEmitter
    mockProc.on('error', () => {});
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'error-test',
    });

    mockProc.emit('error', new Error('spawn ENOENT'));

    const job = tinygradTrain.getJob(result.id);
    expect(job.status).toBe('failed');
    expect(job.error).toBe('spawn ENOENT');
  });

  test('handles Python not found during execution', () => {
    // getPythonPath returns valid during validation (isTinygradInstalled mock),
    // but null during _executeJob
    getPythonPath.mockReturnValue(null);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'no-python-test',
    });

    const job = tinygradTrain.getJob(result.id);
    expect(job.status).toBe('failed');
    expect(job.error).toContain('Python');
  });
});

// ── Progress Parsing ────────────────────────────────────────────────────

describe('progress parsing', () => {
  test('updates job progress from stdout JSON', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'progress-test',
    });

    // Simulate progress output
    mockProc.stdout.emit('data', Buffer.from(
      '{"epoch":1,"step":10,"loss":0.543}\n{"epoch":1,"step":20,"loss":0.321}\n'
    ));

    const job = tinygradTrain.getJob(result.id);
    expect(job.progress.epoch).toBe(1);
    expect(job.progress.step).toBe(20);
    expect(job.progress.loss).toBe(0.321);

    mockProc.emit('exit', 0);
  });

  test('sets adapter_path from completion signal', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'adapter-path-test',
    });

    // Simulate completion with adapter path
    mockProc.stdout.emit('data', Buffer.from(
      '{"status":"completed","adapter_path":"/path/to/adapter"}\n'
    ));
    mockProc.emit('exit', 0);

    const job = tinygradTrain.getJob(result.id);
    expect(job.adapter_path).toBe('/path/to/adapter');
  });

  test('ignores non-JSON stdout lines', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'non-json-test',
    });

    // Should not throw
    mockProc.stdout.emit('data', Buffer.from('Loading model...\nSome debug output\n'));

    const job = tinygradTrain.getJob(result.id);
    expect(job.progress.epoch).toBe(0); // unchanged

    mockProc.emit('exit', 0);
  });
});

// ── Cancel ──────────────────────────────────────────────────────────────

describe('cancelJob', () => {
  test('cancels a running job', () => {
    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'cancel-test',
    });

    const cancelResult = tinygradTrain.cancelJob(result.id);
    expect(cancelResult.success).toBe(true);

    const job = tinygradTrain.getJob(result.id);
    expect(job.status).toBe('cancelled');
    expect(job.completedAt).toBeGreaterThan(0);

    // Simulate process exit after cancel
    mockProc.emit('exit', 137);
  });

  test('returns error for non-existent job', () => {
    const result = tinygradTrain.cancelJob('nonexistent_id');
    expect(result.success).toBe(false);
    expect(result.error).toBe('Job not found');
  });
});

// ── List Jobs ───────────────────────────────────────────────────────────

describe('listJobs', () => {
  test('returns jobs sorted by creation date (newest first)', () => {
    const jobs = tinygradTrain.listJobs();
    expect(Array.isArray(jobs)).toBe(true);
    expect(jobs.length).toBeGreaterThan(0);
    for (let i = 1; i < jobs.length; i++) {
      expect(jobs[i - 1].createdAt).toBeGreaterThanOrEqual(jobs[i].createdAt);
    }
  });
});

// ── Adapter Listing ─────────────────────────────────────────────────────

describe('listAdapters', () => {
  test('returns empty array when no adapters exist', () => {
    // Clear any adapters from previous tests
    const entries = fs.readdirSync(ADAPTERS_DIR, { withFileTypes: true });
    for (const e of entries) {
      if (e.isDirectory()) {
        fs.rmSync(path.join(ADAPTERS_DIR, e.name), { recursive: true, force: true });
      }
    }
    const adapters = tinygradTrain.listAdapters();
    expect(Array.isArray(adapters)).toBe(true);
    expect(adapters.length).toBe(0);
  });

  test('lists adapter directories with metadata', () => {
    const adapterDir = path.join(ADAPTERS_DIR, 'test-adapter');
    fs.mkdirSync(adapterDir, { recursive: true });
    fs.writeFileSync(path.join(adapterDir, 'metadata.json'), JSON.stringify({
      created: '2025-05-15T00:00:00Z',
      base_model: '/path/to/model',
      training_samples: 100,
      final_loss: 0.25,
    }));

    const adapters = tinygradTrain.listAdapters();
    const testAdapter = adapters.find(a => a.name === 'test-adapter');
    expect(testAdapter).toBeDefined();
    expect(testAdapter.base_model).toBe('/path/to/model');
    expect(testAdapter.training_samples).toBe(100);
    expect(testAdapter.final_loss).toBe(0.25);
  });
});

// ── Broadcast / WebSocket Notifications ─────────────────────────────────

describe('setBroadcast', () => {
  test('sends training_job events via broadcast function', () => {
    const broadcast = jest.fn();
    tinygradTrain.setBroadcast(broadcast);

    const mockProc = createMockProcess();
    mockSpawn.mockReturnValue(mockProc);

    const result = tinygradTrain.startTraining({
      base_model: MODEL_PATH,
      training_data: DATA_PATH,
      name: 'broadcast-test',
    });

    // Should have been called with 'running' event
    expect(broadcast).toHaveBeenCalledWith(
      expect.objectContaining({
        type: 'training_job',
        event: 'running',
        id: result.id,
      })
    );

    // Simulate completion
    mockProc.emit('exit', 0);

    expect(broadcast).toHaveBeenCalledWith(
      expect.objectContaining({
        type: 'training_job',
        event: 'completed',
        id: result.id,
      })
    );

    // Reset broadcast
    tinygradTrain.setBroadcast(null);
  });
});

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

describe('persistence', () => {
  test('saves jobs to disk', () => {
    expect(fs.existsSync(JOBS_FILE)).toBe(true);
    const data = JSON.parse(fs.readFileSync(JOBS_FILE, 'utf-8'));
    expect(data.jobs).toBeDefined();
    expect(Array.isArray(data.jobs)).toBe(true);
    expect(data.seq).toBeGreaterThan(0);
  });
});
