// Copyright (c) 2026 defconxt. All rights reserved.
// Licensed under AGPL-3.0 — see LICENSE file for details.

/**
 * Multi-Turn Solver Engine — uses pi-mono's Agent API for iterative
 * solver sessions with mid-solve progress evaluation and pivoting.
 *
 * Uses pi-agent-core's Agent with steer() for mid-solve pivots,
 * replacing the previous nerf CLI RPC approach.
 *
 * @module benchmark/multi-turn-solver
 */

import { importPiCodingAgent, importPiAi, isPiCodingAgentAvailable } from '../gateway/pi-imports.js';
import { analyzeOutput, buildPivotPrompt } from './progress-evaluator.js';

// ---------------------------------------------------------------------------
// Multi-Turn Solver via pi-mono AgentSession
// ---------------------------------------------------------------------------

/**
 * Execute a solver task using multi-turn pi-mono AgentSession.
 *
 * Phase 1: Initial prompt with full strategy template
 * Phase 2: Evaluate progress, inject pivot if stuck
 * Phase 3: Continue or provide extraction guidance
 *
 * @param {string} task - Initial solver prompt
 * @param {object} opts
 * @param {number} opts.timeoutS - Total timeout in seconds
 * @param {string[]} [opts.tags] - Challenge tags for pivot construction
 * @param {string} [opts.cwd] - Working directory for the agent
 * @returns {Promise<{ text: string, steps: string[], toolCalls: number, flag?: string }>}
 */
export async function execMultiTurn(task, { timeoutS = 300, tags = [], cwd } = {}) {
  const piCoding = await importPiCodingAgent();
  const piAi = await importPiAi();

  const {
    createAgentSession,
    SessionManager,
    AuthStorage,
    ModelRegistry,
  } = piCoding;

  const steps = [];
  const totalMs = timeoutS * 1000;
  const phase1Budget = Math.floor(totalMs * 0.60);
  const phase2Budget = Math.floor(totalMs * 0.35);
  const startTime = Date.now();

  const authStorage = AuthStorage.create();
  const modelRegistry = new ModelRegistry(authStorage);
  const model = piAi.getModel('anthropic', 'claude-sonnet-4-6');

  const { session } = await createAgentSession({
    model,
    sessionManager: SessionManager.inMemory(),
    authStorage,
    modelRegistry,
    cwd: cwd || process.cwd(),
  });

  try {
    // ── Phase 1: Initial solve attempt ──────────────────────────────────
    steps.push(`[multi-turn] Phase 1: Initial solve (budget: ${Math.round(phase1Budget / 1000)}s)`);

    const phase1 = await _promptWithTimeout(session, task, phase1Budget);
    steps.push(`[multi-turn] Phase 1 complete: ${phase1.text.length} chars, ${phase1.toolCalls} tool calls, timedOut=${phase1.timedOut}`);

    // Check for flag in phase 1 output
    const flag1 = extractFlagFromText(phase1.text);
    if (flag1) {
      steps.push(`[multi-turn] Flag found in phase 1: ${flag1}`);
      return { text: phase1.text, steps, toolCalls: phase1.toolCalls, flag: flag1 };
    }

    // ── Progress evaluation ────────────────────────────────────────────
    const elapsed = Date.now() - startTime;
    const remaining = totalMs - elapsed;
    if (remaining < 15000) {
      steps.push('[multi-turn] Insufficient time for phase 2 — returning phase 1 output');
      return { text: phase1.text, steps, toolCalls: phase1.toolCalls };
    }

    const findings = analyzeOutput(phase1.text);
    steps.push(`[multi-turn] Evaluator: techniques=${findings.techniques.join(',') || 'none'}, endpoints=${findings.endpoints.length}, stuck=${findings.stuckReason || 'no'}`);

    // ── Phase 2: Pivot via follow-up prompt ─────────────────────────────
    const pivotPrompt = buildPivotPrompt(findings, tags);
    const phase2Msg = pivotPrompt + '\n\nContinue working on the challenge. The target is still running. Find the flag.';

    steps.push(`[multi-turn] Phase 2: Pivot (budget: ${Math.round(Math.min(phase2Budget, remaining) / 1000)}s)`);
    const phase2 = await _promptWithTimeout(session, phase2Msg, Math.min(phase2Budget, remaining));
    steps.push(`[multi-turn] Phase 2 complete: ${phase2.text.length} chars, ${phase2.toolCalls} tool calls`);

    const flag2 = extractFlagFromText(phase2.text);
    if (flag2) {
      steps.push(`[multi-turn] Flag found in phase 2: ${flag2}`);
      return {
        text: phase1.text + '\n\n--- PIVOT ---\n\n' + phase2.text,
        steps,
        toolCalls: phase1.toolCalls + phase2.toolCalls,
        flag: flag2,
      };
    }

    return {
      text: phase1.text + '\n\n--- PIVOT ---\n\n' + phase2.text,
      steps,
      toolCalls: phase1.toolCalls + phase2.toolCalls,
    };

  } finally {
    session.dispose();
  }
}

/**
 * Check if multi-turn execution is available (pi-mono installed).
 */
export function isMultiTurnAvailable() {
  // Sync check — we cache the result after first async check
  return _multiTurnAvailable;
}

let _multiTurnAvailable = false;

// Async initialization — resolve at first import
isPiCodingAgentAvailable().then(avail => { _multiTurnAvailable = avail; }).catch(() => {});

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/**
 * Send a prompt to the session and collect output with a timeout.
 */
async function _promptWithTimeout(session, message, timeoutMs) {
  const textParts = [];
  let toolCalls = 0;
  let timedOut = false;

  const unsub = session.subscribe((event) => {
    if (event.type === 'message_update' && event.assistantMessageEvent?.type === 'text_delta') {
      textParts.push(event.assistantMessageEvent.delta);
    }
    if (event.type === 'message_update' && event.assistantMessageEvent?.type === 'tool_use_start') {
      toolCalls++;
    }
  });

  const timer = setTimeout(() => {
    timedOut = true;
    try { session.abort(); } catch { /* may not support abort */ }
  }, timeoutMs);

  try {
    await session.prompt(message);
  } catch {
    // Collect whatever output we got (timeout or other error)
  }

  clearTimeout(timer);
  if (typeof unsub === 'function') unsub();

  return { text: textParts.join(''), toolCalls, timedOut };
}

function extractFlagFromText(text) {
  if (!text) return null;
  const marker = text.match(/FLAG_FOUND:\s*(.+)/);
  if (marker) return marker[1].trim().replace(/[\s*`_~]+$/, '');

  const patterns = [/FLAG\{[^}]+\}/, /flag\{[^}]+\}/, /picoCTF\{[^}]+\}/, /CTF\{[^}]+\}/];
  for (const p of patterns) {
    const m = text.match(p);
    if (m) return m[0];
  }
  return null;
}
