// PRE Web GUI — Training Data Export Pipeline
// Extracts curated instruction/response pairs from session history and
// experience ledger for external LoRA fine-tuning (unsloth, TRL, MLX).
//
// Supports three export formats:
//   - ChatML: OpenAI-style messages (for TRL/unsloth SFTTrainer)
//   - Alpaca: instruction/input/output triples (for alpaca-lora)
//   - ShareGPT: multi-turn conversations (for Axolotl)
//
// Quality filtering: only exports turns with 2+ successful tool calls,
// non-error outcomes, and reasonable response length. User corrections
// (detected by negation patterns) are tagged as DPO pairs.

const fs = require('fs');
const path = require('path');
const { PRE_DIR, SESSIONS_DIR } = require('./constants');
const { MEMORY_DIR } = require('./constants');

const TRAINING_DIR = path.join(PRE_DIR, 'training');
const EXPORTS_DIR = path.join(TRAINING_DIR, 'exports');

// Ensure directories exist
try { fs.mkdirSync(EXPORTS_DIR, { recursive: true }); } catch (_e) {}

// ── Quality Filters ──────────────────────────────────────────────────────────

const MIN_TOOL_CALLS = 2;        // Minimum tool calls for a quality turn
const MIN_RESPONSE_LENGTH = 50;  // Skip very short assistant responses
const MAX_RESPONSE_LENGTH = 8000; // Truncate extremely long responses

// Patterns that suggest user correction/dissatisfaction
const CORRECTION_PATTERNS = [
  /\bno[,.]?\s+(that'?s|not|wrong|incorrect)/i,
  /\bdon'?t\b.*\bthat\b/i,
  /\bstop\b/i,
  /\btry again\b/i,
  /\bactually[,.]?\s+/i,
  /\binstead[,.]?\s+/i,
  /\bthat'?s not what i/i,
];

// ── Session Parsing ──────────────────────────────────────────────────────────

/**
 * Parse a JSONL session file into structured conversation turns.
 * Groups messages into (user, assistant+tools, tool_results) rounds.
 */
function parseSession(filePath) {
  try {
    const content = fs.readFileSync(filePath, 'utf-8');
    const lines = content.split('\n').filter(l => l.trim());
    return lines.map(l => {
      try { return JSON.parse(l); } catch (_e) { return null; }
    }).filter(Boolean);
  } catch (_e) {
    return [];
  }
}

/**
 * Extract training-quality conversation rounds from session messages.
 * Each round: { userMessage, assistantResponse, toolCalls, toolResults, isCorrection }
 */
function extractRounds(messages) {
  const rounds = [];
  let current = null;

  for (const msg of messages) {
    if (msg.role === 'system') continue;

    if (msg.role === 'user') {
      // Finish previous round
      if (current && current.assistantResponse) {
        rounds.push(current);
      }
      current = {
        userMessage: msg.content || '',
        assistantResponse: '',
        toolCalls: [],
        toolResults: [],
        isCorrection: false,
      };
      // Check if this user message is a correction of the previous response
      if (rounds.length > 0 && current.userMessage) {
        current.isCorrection = CORRECTION_PATTERNS.some(p => p.test(current.userMessage));
      }
    } else if (msg.role === 'assistant' && current) {
      current.assistantResponse += (msg.content || '');
      if (msg.tool_calls) {
        for (const tc of msg.tool_calls) {
          current.toolCalls.push({
            name: tc.function?.name || '',
            args: tc.function?.arguments || {},
          });
        }
      }
    } else if (msg.role === 'tool' && current) {
      current.toolResults.push(msg.content || '');
    }
  }

  // Don't forget the last round
  if (current && current.assistantResponse) {
    rounds.push(current);
  }

  return rounds;
}

/**
 * Filter rounds for training quality.
 */
function filterQualityRounds(rounds, opts = {}) {
  const minTools = opts.minToolCalls || MIN_TOOL_CALLS;
  const includeCorrections = opts.includeCorrections || false;

  return rounds.filter(r => {
    // Must have both user message and assistant response
    if (!r.userMessage.trim() || !r.assistantResponse.trim()) return false;
    // Minimum response length
    if (r.assistantResponse.length < MIN_RESPONSE_LENGTH) return false;
    // Skip corrections unless explicitly included (for DPO pairs)
    if (r.isCorrection && !includeCorrections) return false;
    // Minimum tool calls (0 = include pure chat turns too)
    if (minTools > 0 && r.toolCalls.length < minTools) return false;
    // Skip error-heavy results
    const errorCount = r.toolResults.filter(tr => tr.includes('Error:')).length;
    if (r.toolResults.length > 0 && errorCount / r.toolResults.length > 0.5) return false;
    return true;
  });
}


// ── Format Converters ────────────────────────────────────────────────────────

/**
 * Convert rounds to ChatML format (for TRL SFTTrainer).
 * { messages: [{ role: "user", content: "..." }, { role: "assistant", content: "..." }] }
 */
function toChatML(rounds, systemPrompt) {
  return rounds.map(r => {
    const messages = [];
    if (systemPrompt) {
      messages.push({ role: 'system', content: systemPrompt });
    }
    messages.push({ role: 'user', content: r.userMessage });

    // Include tool calls in assistant content for training
    let response = r.assistantResponse;
    if (r.toolCalls.length > 0) {
      const toolSection = r.toolCalls.map(tc =>
        `<tool_call>${JSON.stringify({ name: tc.name, arguments: tc.args })}</tool_call>`
      ).join('\n');
      response = toolSection + '\n\n' + response;
    }
    messages.push({ role: 'assistant', content: truncate(response) });

    return { messages };
  });
}

/**
 * Convert rounds to Alpaca format (for alpaca-lora).
 * { instruction: "...", input: "", output: "..." }
 */
function toAlpaca(rounds) {
  return rounds.map(r => ({
    instruction: r.userMessage,
    input: '',
    output: truncate(r.assistantResponse),
  }));
}

/**
 * Convert rounds to ShareGPT format (for Axolotl).
 * { conversations: [{ from: "human", value: "..." }, { from: "gpt", value: "..." }] }
 */
function toShareGPT(rounds) {
  // Group consecutive rounds into multi-turn conversations
  const conversations = [];
  let current = [];

  for (const r of rounds) {
    current.push({ from: 'human', value: r.userMessage });
    current.push({ from: 'gpt', value: truncate(r.assistantResponse) });

    // Start a new conversation every 6 turns or on correction
    if (current.length >= 12 || r.isCorrection) {
      conversations.push({ conversations: current });
      current = [];
    }
  }
  if (current.length > 0) {
    conversations.push({ conversations: current });
  }

  return conversations;
}

/**
 * Extract DPO (Direct Preference Optimization) pairs from corrections.
 * Returns: { prompt, chosen (corrected response), rejected (original response) }
 */
function toDPO(rounds) {
  const pairs = [];
  for (let i = 1; i < rounds.length; i++) {
    if (rounds[i].isCorrection && rounds[i - 1]) {
      pairs.push({
        prompt: rounds[i - 1].userMessage,
        chosen: rounds[i].assistantResponse ? truncate(rounds[i].assistantResponse) : null,
        rejected: truncate(rounds[i - 1].assistantResponse),
      });
    }
  }
  return pairs.filter(p => p.chosen); // Only include if we have the corrected version
}

function truncate(text) {
  if (!text) return '';
  return text.length > MAX_RESPONSE_LENGTH
    ? text.slice(0, MAX_RESPONSE_LENGTH) + '\n...(truncated)'
    : text;
}


// ── Experience-Based Training Data ───────────────────────────────────────────

/**
 * Convert experience ledger entries to training pairs.
 * Each experience becomes a "teach me about X" → structured response pair.
 */
function experiencesToTraining() {
  const expDir = path.join(MEMORY_DIR, 'experience');
  if (!fs.existsSync(expDir)) return [];

  const files = fs.readdirSync(expDir).filter(f => f.endsWith('.md') && f !== 'embeddings.json');
  const pairs = [];

  for (const file of files) {
    try {
      const content = fs.readFileSync(path.join(expDir, file), 'utf-8');
      const { frontmatter, body } = parseFrontmatterLocal(content);
      if (!body || !frontmatter) continue;

      // Parse structured fields
      const task = body.match(/\*\*Task:\*\*\s*(.*)/)?.[1]?.trim();
      const approach = body.match(/\*\*Approach:\*\*\s*([\s\S]*?)(?=\*\*|$)/)?.[1]?.trim();
      const outcome = body.match(/\*\*Outcome:\*\*\s*([\s\S]*?)(?=\*\*|$)/)?.[1]?.trim();
      const lesson = body.match(/\*\*Lesson:\*\*\s*([\s\S]*?)(?=\*\*|$)/)?.[1]?.trim();

      if (!task || !lesson) continue;

      pairs.push({
        instruction: `How should I approach: ${task}`,
        input: '',
        output: [
          approach ? `**Approach:** ${approach}` : null,
          outcome ? `**Outcome:** ${outcome}` : null,
          `**Key lesson:** ${lesson}`,
        ].filter(Boolean).join('\n\n'),
        source: 'experience',
        file,
      });
    } catch (_e) {}
  }

  return pairs;
}

/**
 * Simple frontmatter parser (avoid circular dependency on memory.js).
 */
function parseFrontmatterLocal(content) {
  const match = content.match(/^---\n([\s\S]*?)\n---\n([\s\S]*)$/);
  if (!match) return { frontmatter: null, body: content };
  const fm = {};
  for (const line of match[1].split('\n')) {
    const idx = line.indexOf(':');
    if (idx > 0) {
      fm[line.slice(0, idx).trim()] = line.slice(idx + 1).trim().replace(/^['"]|['"]$/g, '');
    }
  }
  return { frontmatter: fm, body: match[2].trim() };
}


// ── Export Pipeline ──────────────────────────────────────────────────────────

/**
 * Export training data from sessions and experiences.
 *
 * @param {Object} opts
 * @param {string} opts.format - 'chatml' | 'alpaca' | 'sharegpt' | 'dpo'
 * @param {string} [opts.project] - Filter sessions by project slug
 * @param {number} [opts.minToolCalls] - Minimum tool calls per round (default 2)
 * @param {boolean} [opts.includeExperience] - Include experience ledger (default true)
 * @param {string} [opts.systemPrompt] - System prompt for ChatML format
 * @param {number} [opts.maxSamples] - Maximum samples to export
 * @returns {{ path: string, count: number, format: string, stats: Object }}
 */
function exportTrainingData(opts = {}) {
  const format = opts.format || 'chatml';
  const project = opts.project || null;
  const minToolCalls = opts.minToolCalls ?? MIN_TOOL_CALLS;
  const includeExperience = opts.includeExperience !== false;
  const maxSamples = opts.maxSamples || 10000;

  // Collect all sessions
  const sessionFiles = [];
  if (fs.existsSync(SESSIONS_DIR)) {
    for (const f of fs.readdirSync(SESSIONS_DIR)) {
      if (!f.endsWith('.jsonl')) continue;
      if (project && !f.startsWith(project + ':')) continue;
      // Skip system sessions
      if (f.startsWith('cron:') || f.startsWith('mcp:')) continue;
      sessionFiles.push(path.join(SESSIONS_DIR, f));
    }
  }

  // Extract quality rounds from all sessions
  let allRounds = [];
  let totalSessions = 0;
  let totalMessages = 0;

  for (const sf of sessionFiles) {
    const messages = parseSession(sf);
    if (messages.length < 4) continue; // Skip very short sessions
    totalSessions++;
    totalMessages += messages.length;

    const rounds = extractRounds(messages);
    const quality = filterQualityRounds(rounds, {
      minToolCalls,
      includeCorrections: format === 'dpo',
    });
    allRounds.push(...quality);
  }

  // Add experience-derived pairs
  let experienceCount = 0;
  if (includeExperience && format !== 'dpo') {
    const expPairs = experiencesToTraining();
    experienceCount = expPairs.length;
    // Convert to round format
    for (const ep of expPairs) {
      allRounds.push({
        userMessage: ep.instruction,
        assistantResponse: ep.output,
        toolCalls: [],
        toolResults: [],
        isCorrection: false,
      });
    }
  }

  // Cap samples
  if (allRounds.length > maxSamples) {
    allRounds = allRounds.slice(-maxSamples); // Keep most recent
  }

  // Convert to target format
  let data;
  switch (format) {
    case 'chatml':
      data = toChatML(allRounds, opts.systemPrompt || null);
      break;
    case 'alpaca':
      data = toAlpaca(allRounds);
      break;
    case 'sharegpt':
      data = toShareGPT(allRounds);
      break;
    case 'dpo':
      data = toDPO(allRounds);
      break;
    default:
      throw new Error(`Unknown format: ${format}`);
  }

  // Write to file
  const timestamp = new Date().toISOString().replace(/[:.]/g, '-').slice(0, 19);
  const filename = `training-${format}-${timestamp}.jsonl`;
  const outPath = path.join(EXPORTS_DIR, filename);

  const lines = data.map(d => JSON.stringify(d)).join('\n');
  fs.writeFileSync(outPath, lines + '\n');

  return {
    path: outPath,
    filename,
    count: data.length,
    format,
    stats: {
      sessionsScanned: totalSessions,
      totalMessages,
      qualityRounds: allRounds.length,
      experienceEntries: experienceCount,
      dpoPairs: format === 'dpo' ? data.length : 0,
    },
  };
}

/**
 * List previously exported training files.
 */
function listExports() {
  if (!fs.existsSync(EXPORTS_DIR)) return [];
  return fs.readdirSync(EXPORTS_DIR)
    .filter(f => f.endsWith('.jsonl'))
    .map(f => {
      const stat = fs.statSync(path.join(EXPORTS_DIR, f));
      return {
        filename: f,
        path: path.join(EXPORTS_DIR, f),
        size: stat.size,
        created: stat.mtime.toISOString(),
      };
    })
    .sort((a, b) => new Date(b.created) - new Date(a.created));
}

/**
 * Get a preview of what would be exported (dry run).
 */
function previewExport(opts = {}) {
  const minToolCalls = opts.minToolCalls ?? MIN_TOOL_CALLS;
  const project = opts.project || null;

  let totalSessions = 0;
  let totalRounds = 0;
  let qualityRounds = 0;
  let correctionRounds = 0;
  const toolCounts = {};

  if (fs.existsSync(SESSIONS_DIR)) {
    for (const f of fs.readdirSync(SESSIONS_DIR)) {
      if (!f.endsWith('.jsonl')) continue;
      if (project && !f.startsWith(project + ':')) continue;
      if (f.startsWith('cron:') || f.startsWith('mcp:')) continue;

      const messages = parseSession(path.join(SESSIONS_DIR, f));
      if (messages.length < 4) continue;
      totalSessions++;

      const rounds = extractRounds(messages);
      totalRounds += rounds.length;

      for (const r of rounds) {
        if (r.isCorrection) correctionRounds++;
        for (const tc of r.toolCalls) {
          toolCounts[tc.name] = (toolCounts[tc.name] || 0) + 1;
        }
      }

      const quality = filterQualityRounds(rounds, { minToolCalls });
      qualityRounds += quality.length;
    }
  }

  const expCount = experiencesToTraining().length;

  return {
    sessions: totalSessions,
    totalRounds,
    qualityRounds,
    correctionRounds,
    experienceEntries: expCount,
    estimatedTotal: qualityRounds + expCount,
    topTools: Object.entries(toolCounts)
      .sort((a, b) => b[1] - a[1])
      .slice(0, 10)
      .map(([name, count]) => ({ name, count })),
  };
}


module.exports = {
  exportTrainingData,
  listExports,
  previewExport,
  extractRounds,
  filterQualityRounds,
  toChatML,
  toAlpaca,
  toShareGPT,
  toDPO,
  experiencesToTraining,
  parseSession,
  // Constants for testing
  MIN_TOOL_CALLS,
  MIN_RESPONSE_LENGTH,
  MAX_RESPONSE_LENGTH,
  CORRECTION_PATTERNS,
};
