// Tests for training data export pipeline

const path = require('path');
const os = require('os');
const fs = require('fs');

const mockPreDir = path.join(os.tmpdir(), `pre-test-training-${process.pid}`);

jest.mock('../src/constants', () => ({
  PRE_DIR: mockPreDir,
  MODEL: 'pre-gemma4',
  MODEL_CTX: 8192,
  OLLAMA_PORT: 11434,
  MAX_TOOL_TURNS: 35,
  SESSIONS_DIR: require('path').join(mockPreDir, 'sessions'),
  MEMORY_DIR: require('path').join(mockPreDir, 'memory'),
  ARTIFACTS_DIR: require('path').join(mockPreDir, 'artifacts'),
  CONNECTIONS_FILE: require('path').join(mockPreDir, 'connections.json'),
  COMFYUI_FILE: require('path').join(mockPreDir, 'comfyui.json'),
  CRON_FILE: require('path').join(mockPreDir, 'cron.json'),
}));

const SESSIONS_DIR = path.join(mockPreDir, 'sessions');
const MEMORY_DIR = path.join(mockPreDir, 'memory');
const EXPERIENCE_DIR = path.join(mockPreDir, 'memory', 'experience');

beforeAll(() => {
  fs.mkdirSync(SESSIONS_DIR, { recursive: true });
  fs.mkdirSync(EXPERIENCE_DIR, { recursive: true });
});

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

// Helper: write a mock session JSONL file
function writeSession(filename, messages) {
  const lines = messages.map(m => JSON.stringify(m)).join('\n');
  fs.writeFileSync(path.join(SESSIONS_DIR, filename), lines + '\n');
}

// Helper: write a mock experience entry
function writeExperience(filename, content) {
  fs.writeFileSync(path.join(EXPERIENCE_DIR, filename), content);
}


describe('Training Data — Round Extraction', () => {
  test('extractRounds groups user/assistant/tool messages', () => {
    const training = require('../src/training');
    const messages = [
      { role: 'system', content: 'You are a helpful assistant.' },
      { role: 'user', content: 'List files in /tmp' },
      { role: 'assistant', content: 'Here are the files:', tool_calls: [{ function: { name: 'bash', arguments: { command: 'ls /tmp' } } }] },
      { role: 'tool', content: '<tool_response name="bash">\nfile1.txt\nfile2.txt\n</tool_response>' },
      { role: 'assistant', content: 'Found 2 files in /tmp.' },
      { role: 'user', content: 'Thanks!' },
    ];

    const rounds = training.extractRounds(messages);
    // "Thanks!" has no assistant response, so only the first round is captured
    expect(rounds.length).toBe(1);
    expect(rounds[0].userMessage).toBe('List files in /tmp');
    expect(rounds[0].toolCalls.length).toBe(1);
    expect(rounds[0].toolCalls[0].name).toBe('bash');
  });

  test('detects user corrections', () => {
    const training = require('../src/training');
    const messages = [
      { role: 'user', content: 'Read the config file' },
      { role: 'assistant', content: 'Here is the config.', tool_calls: [{ function: { name: 'read_file', arguments: { path: '/etc/config' } } }] },
      { role: 'tool', content: '<tool_response name="read_file">\nconfig data\n</tool_response>' },
      { role: 'assistant', content: 'The config says X.' },
      { role: 'user', content: "No, that's not what I wanted. Try again." },
      { role: 'assistant', content: 'Let me re-read the file.' },
    ];

    const rounds = training.extractRounds(messages);
    expect(rounds[1].isCorrection).toBe(true);
  });
});


describe('Training Data — Quality Filtering', () => {
  test('filters out low-quality rounds', () => {
    const training = require('../src/training');
    const rounds = [
      { userMessage: 'Hi', assistantResponse: 'Hello!', toolCalls: [], toolResults: [], isCorrection: false },
      { userMessage: 'Search for X', assistantResponse: 'A'.repeat(100), toolCalls: [{ name: 'bash' }, { name: 'read_file' }], toolResults: ['result1', 'result2'], isCorrection: false },
      { userMessage: '', assistantResponse: 'empty user', toolCalls: [], toolResults: [], isCorrection: false },
    ];

    const filtered = training.filterQualityRounds(rounds);
    expect(filtered.length).toBe(1);
    expect(filtered[0].userMessage).toBe('Search for X');
  });

  test('includes corrections when requested', () => {
    const training = require('../src/training');
    const rounds = [
      { userMessage: 'Do X', assistantResponse: 'A'.repeat(100), toolCalls: [{ name: 'a' }, { name: 'b' }], toolResults: ['ok', 'ok'], isCorrection: false },
      { userMessage: "No that's wrong", assistantResponse: 'B'.repeat(100), toolCalls: [{ name: 'a' }, { name: 'b' }], toolResults: ['ok', 'ok'], isCorrection: true },
    ];

    expect(training.filterQualityRounds(rounds).length).toBe(1);
    expect(training.filterQualityRounds(rounds, { includeCorrections: true }).length).toBe(2);
  });
});


describe('Training Data — Format Conversion', () => {
  const rounds = [
    { userMessage: 'Help me', assistantResponse: 'Sure, here is the answer.', toolCalls: [], toolResults: [], isCorrection: false },
    { userMessage: 'Run a test', assistantResponse: 'Tests passed.', toolCalls: [{ name: 'bash', args: { command: 'npm test' } }], toolResults: ['ok'], isCorrection: false },
  ];

  test('toChatML produces correct structure', () => {
    const training = require('../src/training');
    const result = training.toChatML(rounds, 'System prompt');
    expect(result.length).toBe(2);
    expect(result[0].messages[0].role).toBe('system');
    expect(result[0].messages[1].role).toBe('user');
    expect(result[0].messages[2].role).toBe('assistant');
  });

  test('toAlpaca produces instruction/output format', () => {
    const training = require('../src/training');
    const result = training.toAlpaca(rounds);
    expect(result.length).toBe(2);
    expect(result[0]).toHaveProperty('instruction');
    expect(result[0]).toHaveProperty('input');
    expect(result[0]).toHaveProperty('output');
  });

  test('toShareGPT produces conversations format', () => {
    const training = require('../src/training');
    const result = training.toShareGPT(rounds);
    expect(result.length).toBeGreaterThanOrEqual(1);
    expect(result[0].conversations[0].from).toBe('human');
    expect(result[0].conversations[1].from).toBe('gpt');
  });

  test('toDPO extracts correction pairs', () => {
    const training = require('../src/training');
    const dpoRounds = [
      { userMessage: 'Do X', assistantResponse: 'Wrong answer', toolCalls: [], toolResults: [], isCorrection: false },
      { userMessage: "No that's wrong", assistantResponse: 'Correct answer', toolCalls: [], toolResults: [], isCorrection: true },
    ];
    const result = training.toDPO(dpoRounds);
    expect(result.length).toBe(1);
    expect(result[0].prompt).toBe('Do X');
    expect(result[0].rejected).toBe('Wrong answer');
    expect(result[0].chosen).toBe('Correct answer');
  });
});


describe('Training Data — Experience Conversion', () => {
  test('converts experience entries to training pairs', () => {
    writeExperience('test_lesson.md', `---
name: test_lesson
description: A test lesson
type: experience
created: 2025-01-01
---

**Task:** Build a REST API

**Approach:** Used Express with middleware

**Outcome:** API worked correctly

**Lesson:** Always add error handling middleware early
`);

    const training = require('../src/training');
    const pairs = training.experiencesToTraining();
    expect(pairs.length).toBeGreaterThanOrEqual(1);
    const pair = pairs.find(p => p.file === 'test_lesson.md');
    expect(pair).toBeTruthy();
    expect(pair.instruction).toContain('Build a REST API');
    expect(pair.output).toContain('error handling middleware');
  });
});


describe('Training Data — Full Export', () => {
  test('exports JSONL file from session data', () => {
    writeSession('web:test-export.jsonl', [
      { role: 'system', content: 'You are helpful.' },
      { role: 'user', content: 'Search for config files' },
      { role: 'assistant', content: 'A'.repeat(100), tool_calls: [
        { function: { name: 'bash', arguments: { command: 'find . -name "*.conf"' } } },
        { function: { name: 'read_file', arguments: { path: 'app.conf' } } },
      ]},
      { role: 'tool', content: '<tool_response name="bash">\napp.conf\n</tool_response>\n<tool_response name="read_file">\nport=8080\n</tool_response>' },
      { role: 'assistant', content: 'B'.repeat(100) },
    ]);

    const training = require('../src/training');
    const result = training.exportTrainingData({ format: 'chatml' });
    expect(result.count).toBeGreaterThanOrEqual(1);
    expect(result.format).toBe('chatml');
    expect(fs.existsSync(result.path)).toBe(true);

    // Verify JSONL structure
    const content = fs.readFileSync(result.path, 'utf-8');
    const lines = content.trim().split('\n');
    for (const line of lines) {
      const parsed = JSON.parse(line);
      expect(parsed).toHaveProperty('messages');
    }
  });

  test('listExports returns exported files', () => {
    const training = require('../src/training');
    const exports = training.listExports();
    expect(exports.length).toBeGreaterThanOrEqual(1);
    expect(exports[0]).toHaveProperty('filename');
    expect(exports[0]).toHaveProperty('size');
  });
});
