'use strict';

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

const ALL_SOURCES = ['knowledge', 'chat', 'patterns'];

/**
 * Export training data from brain DB to JSONL format.
 * @param {Object} options
 * @param {Object} options.brain - Brain module (for dependency injection)
 * @param {string} [options.outputPath] - Where to write JSONL file
 * @param {string[]} [options.sources] - Which sources: ['knowledge', 'chat', 'patterns']
 * @returns {{ samples: Array, stats: { total, bySource } }}
 */
async function exportTrainingData({ brain, outputPath, sources } = {}) {
  const activeSources = sources || ALL_SOURCES;
  const allSamples = [];
  const bySource = {};

  for (const src of ALL_SOURCES) {
    bySource[src] = 0;
  }

  // 1. Knowledge triples -> Q&A pairs
  if (activeSources.includes('knowledge')) {
    const entries = brain.findKnowledge ? brain.findKnowledge() : [];
    for (const entry of entries) {
      if (entry.confidence < 0.5) continue;
      if (entry.status !== 'active') continue;
      const instruction = `What ${entry.predicate} ${entry.subject}?`;
      const output = entry.object;
      allSamples.push({ instruction, output, source: 'knowledge' });
    }
  }

  // 2. Chat messages -> conversation pairs
  if (activeSources.includes('chat')) {
    const messages = brain.listChatMessages ? brain.listChatMessages() : [];
    // Group by session
    const sessions = new Map();
    for (const msg of messages) {
      const sid = msg.session_id || '__default__';
      if (!sessions.has(sid)) sessions.set(sid, []);
      sessions.get(sid).push(msg);
    }
    // Create pairs from consecutive user/assistant turns
    for (const [, msgs] of sessions) {
      for (let i = 0; i < msgs.length - 1; i++) {
        const userMsg = msgs[i];
        const assistantMsg = msgs[i + 1];
        if (userMsg.role !== 'user' || assistantMsg.role !== 'assistant') continue;
        // Filter short messages
        if (!userMsg.content || userMsg.content.length < 20) continue;
        if (!assistantMsg.content || assistantMsg.content.length < 20) continue;
        // Filter tool-use-only messages
        if (isToolUseOnly(assistantMsg.content)) continue;
        allSamples.push({
          instruction: userMsg.content,
          output: assistantMsg.content,
          source: 'chat',
        });
      }
    }
  }

  // 3. Patterns -> instruction pairs
  if (activeSources.includes('patterns')) {
    const patterns = (brain.listPatterns ? brain.listPatterns() : []);
    for (const p of patterns) {
      const instruction = `Describe the pattern: ${p.pattern_type}`;
      const output = p.description;
      allSamples.push({ instruction, output, source: 'patterns' });
    }
  }

  // 4. Quality filtering
  const filtered = [];
  const seen = new Set();
  for (const sample of allSamples) {
    // Skip empty instruction/output
    if (!sample.instruction || !sample.instruction.trim()) continue;
    if (!sample.output || sample.output.length < 10) continue;
    // Dedup by instruction
    if (seen.has(sample.instruction)) continue;
    seen.add(sample.instruction);
    filtered.push(sample);
  }

  // 5. Compute stats
  for (const sample of filtered) {
    bySource[sample.source] = (bySource[sample.source] || 0) + 1;
  }
  const stats = { total: filtered.length, bySource };

  // 6. Write JSONL if outputPath provided
  if (outputPath) {
    const dir = path.dirname(outputPath);
    if (!fs.existsSync(dir)) {
      fs.mkdirSync(dir, { recursive: true });
    }
    const lines = filtered.map(s => JSON.stringify(s));
    fs.writeFileSync(outputPath, lines.join('\n') + (lines.length ? '\n' : ''), 'utf8');
  }

  return { samples: filtered, stats };
}

/**
 * Detect if a message is tool-use only (JSON array of tool_use blocks).
 */
function isToolUseOnly(content) {
  if (!content) return false;
  const trimmed = content.trim();
  if (!trimmed.startsWith('[')) return false;
  try {
    const parsed = JSON.parse(trimmed);
    if (Array.isArray(parsed) && parsed.every(item => item.type === 'tool_use')) {
      return true;
    }
  } catch {
    // Not JSON, not tool-use only
  }
  return false;
}

/**
 * Export harvested training data from JSONL files.
 * @param {Object} options
 * @param {string} options.harvestDir - Directory containing harvested/*.jsonl files
 * @param {string} [options.outputPath] - Where to write output JSONL
 * @param {number} [options.minQuality=0.6] - Minimum quality_label to include
 * @param {string} [options.format='alpaca'] - Output format: 'alpaca' or 'chat-ml'
 * @param {number} [options.maxPerTaskType=0] - Max samples per task type (0=unlimited)
 * @returns {{ samples: Array, stats: { total, bySource, byTaskType, filtered } }}
 */
async function exportHarvestedData({ harvestDir, outputPath, minQuality = 0.6, format = 'alpaca', maxPerTaskType = 0 } = {}) {
  const harvestedDir = path.join(harvestDir, 'training', 'harvested');
  if (!fs.existsSync(harvestedDir)) return { samples: [], stats: { total: 0, bySource: {}, byTaskType: {}, filtered: 0 } };

  // Read all JSONL files
  const files = fs.readdirSync(harvestedDir).filter(f => f.endsWith('.jsonl')).sort();
  let allSamples = [];
  for (const file of files) {
    const lines = fs.readFileSync(path.join(harvestedDir, file), 'utf8').split('\n').filter(Boolean);
    for (const line of lines) {
      try { allSamples.push(JSON.parse(line)); } catch {}
    }
  }

  // Filter by quality
  const filtered = allSamples.length;
  allSamples = allSamples.filter(s => (s.quality_label || 0) >= minQuality);

  // Deduplicate by id
  const seen = new Set();
  allSamples = allSamples.filter(s => {
    if (seen.has(s.id)) return false;
    seen.add(s.id);
    return true;
  });

  // Balance by task_type if maxPerTaskType > 0
  if (maxPerTaskType > 0) {
    const byType = {};
    const balanced = [];
    for (const s of allSamples) {
      const t = s.task_type || 'unknown';
      byType[t] = (byType[t] || 0) + 1;
      if (byType[t] <= maxPerTaskType) balanced.push(s);
    }
    allSamples = balanced;
  }

  // Convert to output format
  const outputSamples = allSamples.map(s => {
    if (format === 'chat-ml') {
      return {
        messages: [
          { role: 'system', content: 'You are a helpful AI assistant.' },
          { role: 'user', content: s.prompt },
          { role: 'assistant', content: s.response },
        ],
      };
    }
    // Default: alpaca format
    return {
      instruction: s.prompt,
      input: '',
      output: s.response,
    };
  });

  // Build stats
  const stats = { total: outputSamples.length, bySource: {}, byTaskType: {}, filtered: filtered - allSamples.length };
  for (const s of allSamples) {
    stats.bySource[s.source] = (stats.bySource[s.source] || 0) + 1;
    stats.byTaskType[s.task_type] = (stats.byTaskType[s.task_type] || 0) + 1;
  }

  // Write to file if outputPath specified
  if (outputPath) {
    const dir = path.dirname(outputPath);
    if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
    const jsonl = outputSamples.map(s => JSON.stringify(s)).join('\n') + '\n';
    fs.writeFileSync(outputPath, jsonl);
  }

  return { samples: outputSamples, stats };
}

module.exports = { exportTrainingData, exportHarvestedData };
