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

/**
 * Claude Code Solver — Delegates CTF challenges to a Claude Code subagent.
 *
 * Instead of using the Anthropic API directly (expensive) or Ollama (inconsistent),
 * this solver spawns a Claude Code subagent that uses bash/curl/python to interact
 * with the Docker-hosted challenge target.
 *
 * The subagent has full access to the host's tools — it can curl the Docker network,
 * run Python exploit scripts, use netcat, etc. This is the most capable solver
 * because it leverages Claude Code's excellent tool use.
 *
 * @module benchmark/claude-code-solver
 */

import { SolverResult } from './models.js';
import { SolverAdapter } from './solver.js';
import { analyzeDockerEnvironment, formatEnvironmentBrief } from './docker-analyzer.js';
import { selectStrategy, formatStrategy } from './strategies/loader.js';
import { analyzeOutput, buildPivotPrompt } from './progress-evaluator.js';
import { execMultiTurn, isMultiTurnAvailable } from './multi-turn-solver.js';
import { execSync } from 'node:child_process';
import { existsSync, readFileSync, readdirSync } from 'node:fs';
import { join, dirname, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';

const __dirname = dirname(fileURLToPath(import.meta.url));

/**
 * Map challenge tags to relevant NERF knowledge docs and skill domains.
 */
const TAG_TO_KNOWLEDGE = {
  sqli:       ['api-exploitation-deep.md', 'pentest-cheatsheet-ultimate.md'],
  xss:        ['api-exploitation-deep.md', 'pentest-cheatsheet-ultimate.md'],
  ssrf:       ['api-exploitation-deep.md', 'pentest-cheatsheet-ultimate.md'],
  lfi:        ['pentest-cheatsheet-ultimate.md', 'linux-exploitation-deep.md'],
  rce:        ['linux-exploitation-deep.md', 'pentest-cheatsheet-ultimate.md'],
  web:        ['api-exploitation-deep.md', 'pentest-cheatsheet-ultimate.md'],
  crypto:     ['crypto-pki-tls-deep.md'],
  forensics:  ['ctf-methodology-deep.md'],
  pwn:        ['binary-exploitation-deep.md', 'linux-exploitation-deep.md'],
  reversing:  ['binary-exploitation-deep.md'],
  default:    ['ctf-methodology-deep.md', 'pentest-cheatsheet-ultimate.md'],
};

const TAG_TO_SKILLS = {
  sqli:       ['application-security', 'api-security'],
  xss:        ['application-security', 'api-security'],
  ssrf:       ['application-security', 'cloud-security'],
  lfi:        ['application-security'],
  rce:        ['exploit-development'],
  web:        ['application-security', 'api-security', 'bug-bounty'],
  crypto:     ['cryptography'],
  forensics:  ['digital-forensics', 'nerf-forensics'],
  pwn:        ['binary-exploitation', 'exploit-development'],
  reversing:  ['reverse-engineering'],
};

/**
 * Load relevant knowledge snippets for challenge tags.
 * Returns up to maxChars of combined knowledge content.
 */
function loadKnowledge(tags, maxChars = 8000) {
  const root = resolve(__dirname, '..', '..', '..');
  const knowledgeDir = join(root, 'knowledge');
  if (!existsSync(knowledgeDir)) return '';

  const fileSet = new Set();
  for (const tag of tags) {
    const files = TAG_TO_KNOWLEDGE[tag.toLowerCase()] || TAG_TO_KNOWLEDGE.default;
    files.forEach(f => fileSet.add(f));
  }
  if (fileSet.size === 0) {
    TAG_TO_KNOWLEDGE.default.forEach(f => fileSet.add(f));
  }

  const chunks = [];
  let totalLen = 0;
  for (const filename of fileSet) {
    const filepath = join(knowledgeDir, filename);
    if (!existsSync(filepath)) continue;
    try {
      const content = readFileSync(filepath, 'utf8');
      // Take first N chars to stay within budget
      const budget = Math.min(maxChars - totalLen, 4000);
      if (budget <= 0) break;
      chunks.push(`### ${filename}\n${content.slice(0, budget)}`);
      totalLen += budget;
    } catch { /* skip unreadable */ }
  }
  return chunks.join('\n\n');
}

/**
 * Load the domain-level SKILL.md for relevant skill domains.
 * Returns up to maxChars of combined skill summaries.
 */
function loadSkillSummaries(tags, maxChars = 4000) {
  const root = resolve(__dirname, '..', '..', '..');
  const skillsDir = existsSync(join(root, 'skills'))
    ? join(root, 'skills')
    : join(root, 'src', 'resources', 'skills');
  if (!existsSync(skillsDir)) return '';

  const domainSet = new Set();
  for (const tag of tags) {
    const domains = TAG_TO_SKILLS[tag.toLowerCase()] || [];
    domains.forEach(d => domainSet.add(d));
  }

  const chunks = [];
  let totalLen = 0;
  for (const domain of domainSet) {
    const skillPath = join(skillsDir, domain, 'SKILL.md');
    if (!existsSync(skillPath)) continue;
    try {
      const content = readFileSync(skillPath, 'utf8');
      const budget = Math.min(maxChars - totalLen, 2000);
      if (budget <= 0) break;
      chunks.push(`### ${domain}\n${content.slice(0, budget)}`);
      totalLen += budget;
    } catch { /* skip */ }
  }
  return chunks.join('\n\n');
}

// ---------------------------------------------------------------------------
// ClaudeCodeSolver
// ---------------------------------------------------------------------------

/**
 * Solver that delegates challenges to a Claude Code subagent.
 *
 * The subagent uses bash, curl, python3, and other host tools to
 * interact with the Docker-hosted challenge target. This avoids
 * the API cost of direct Claude API calls and the inconsistency
 * of smaller models like Ollama.
 */
export class ClaudeCodeSolver extends SolverAdapter {
  /**
   * @param {object} [opts]
   * @param {number} [opts.timeoutS=300] - Max time for the subagent
   * @param {string} [opts.agent='ctf-solver'] - Subagent name
   * @param {Function} [opts.subagentFn] - Direct subagent function (when running inside pi)
   */
  constructor({ timeoutS = 300, agent = 'ctf-solver', subagentFn } = {}) {
    super();
    this._timeoutS = timeoutS;
    this._agent = agent;
    this._subagentFn = subagentFn || null;
  }

  get name() { return 'claude-code'; }

  /**
   * Solve a benchmark challenge using a Claude Code subagent.
   *
   * @param {import('./models.js').BenchmarkConfig} config
   * @param {string} targetUrl - Target URL (from Docker compose)
   * @param {string} expectedFlag - Expected flag for validation
   * @returns {Promise<SolverResult>}
   */
  async solve(config, targetUrl, expectedFlag) {
    const start = performance.now();
    const steps = [];

    // ── Try multi-turn RPC solver first (preferred path) ──────────────
    // Skip for subclasses (e.g., test mocks) and when subagentFn is set
    if (isMultiTurnAvailable() && !this._subagentFn && this.constructor === ClaudeCodeSolver) {
      try {
        const task = this._buildTask(config, targetUrl);
        steps.push('[solver] Using multi-turn RPC solver');

        const result = await execMultiTurn(task, {
          timeoutS: this._timeoutS,
          tags: config.tags,
        });

        steps.push(...result.steps);

        if (result.flag) {
          steps.push(`[flag] Found: ${result.flag}`);
          return new SolverResult({
            flagFound: true,
            flagValue: result.flag,
            durationS: (performance.now() - start) / 1000,
            steps,
          });
        }

        // Check output for flags the multi-turn engine may have missed
        const flagResult = this._extractFlag(result.text, start, steps);
        if (flagResult) return flagResult;

        steps.push(`[result] No flag found in multi-turn output (${result.text.length} chars)`);
        steps.push(`[output-tail] ${result.text.slice(-800)}`);

        return new SolverResult({
          durationS: (performance.now() - start) / 1000,
          steps,
        });
      } catch (e) {
        steps.push(`[multi-turn] Failed: ${e.message} — falling back to single-shot`);
        // Fall through to single-shot loop
      }
    }

    // ── Fallback: single-shot solver loop with retry ────────────────
    const maxAttempts = 2;
    const attempt1Budget = Math.floor(this._timeoutS * 0.65);
    const attempt2Budget = this._timeoutS - attempt1Budget;
    let lastError = null;

    for (let attempt = 1; attempt <= maxAttempts; attempt++) {
      const attemptStart = performance.now();
      const budget = attempt === 1 ? attempt1Budget : attempt2Budget;

      try {
        // Build prompt — attempt 2 gets pivot instructions prepended
        let task;
        if (attempt === 1) {
          task = this._buildTask(config, targetUrl);
          steps.push(`[solver-loop] Attempt ${attempt}/${maxAttempts} (budget: ${budget}s)`);
        } else {
          // Analyze previous attempt output and build pivot prompt
          const prevOutput = steps.filter(s => s.startsWith('[output-tail]')).map(s => s.slice(14)).join('\n');
          const allOutput = steps.join('\n');
          const findings = analyzeOutput(allOutput);
          const pivotPrompt = buildPivotPrompt(findings, config.tags);
          steps.push(`[solver-loop] Attempt ${attempt}/${maxAttempts} — pivot (budget: ${budget}s)`);
          steps.push(`[evaluator] Stuck: ${findings.stuckReason || 'unknown'}`);
          steps.push(`[evaluator] Techniques tried: ${findings.techniques.join(', ') || 'none detected'}`);
          steps.push(`[evaluator] Endpoints found: ${findings.endpoints.length}`);
          if (findings.credentials.length > 0) steps.push(`[evaluator] Credentials found: ${findings.credentials.length}`);

          task = pivotPrompt + '\n\n' + this._buildTask(config, targetUrl);
        }

        steps.push(`[subagent] Delegating to ${this._agent}: ${config.name}`);
        steps.push(`[target] ${targetUrl}`);

        // Execute with attempt-specific timeout
        const origTimeout = this._timeoutS;
        this._timeoutS = budget;
        let result;
        try {
          result = await this._execSubagent(task);
        } finally {
          this._timeoutS = origTimeout;
        }

        steps.push(`[subagent] Attempt ${attempt} completed in ${((performance.now() - attemptStart) / 1000).toFixed(1)}s`);

        // Try to extract flag from output
        const flagResult = this._extractFlag(result, start, steps);
        if (flagResult) return flagResult;

        // No flag found — record output tail for potential pivot
        steps.push(`[result] No flag found in attempt ${attempt} output (${result.length} chars)`);
        steps.push(`[output-tail] ${result.slice(-800)}`);

        // Check remaining time budget
        const elapsed = (performance.now() - start) / 1000;
        if (elapsed > this._timeoutS * 0.9) {
          steps.push(`[solver-loop] Time budget exhausted after attempt ${attempt}`);
          break;
        }

      } catch (e) {
        lastError = String(e);
        steps.push(`[error] Attempt ${attempt}: ${e}`);
        // If timeout on attempt 1, still try attempt 2 with remaining budget
        if (attempt === 1 && e.message?.includes('timeout')) {
          const elapsed = (performance.now() - start) / 1000;
          if (elapsed < this._timeoutS * 0.7) continue;
        }
        break;
      }
    }

    return new SolverResult({
      durationS: (performance.now() - start) / 1000,
      error: lastError,
      steps,
    });
  }

  /**
   * Extract flag from solver output. Returns SolverResult if found, null otherwise.
   * @private
   */
  _extractFlag(result, start, steps) {
    // Check for explicit FLAG_FOUND marker
    const flagMatch = result.match(/FLAG_FOUND:\s*(.+)/);
    if (flagMatch) {
      const flag = flagMatch[1].trim().replace(/[\s*`_~]+$/, '');
      steps.push(`[flag] Found: ${flag}`);
      return new SolverResult({
        flagFound: true,
        flagValue: flag,
        durationS: (performance.now() - start) / 1000,
        steps,
      });
    }

    // Check for common flag patterns in output
    const patterns = [
      /FLAG\{[^}]+\}/,
      /flag\{[^}]+\}/,
      /picoCTF\{[^}]+\}/,
      /CTF\{[^}]+\}/,
    ];

    for (const pattern of patterns) {
      const match = result.match(pattern);
      if (match) {
        steps.push(`[flag] Extracted from output: ${match[0]}`);
        return new SolverResult({
          flagFound: true,
          flagValue: match[0],
          durationS: (performance.now() - start) / 1000,
          steps,
        });
      }
    }

    return null;
  }

  /**
   * Build the task prompt for the subagent.
   */
  _buildTask(config, targetUrl) {
    const parts = [
      `You are NERF — an expert CTF solver with deep security knowledge.`,
      ``,
      `## Challenge`,
      `**Name:** ${config.name}`,
      `**Description:** ${config.description}`,
      `**Level:** ${config.level}/5`,
      `**Tags:** ${config.tags.join(', ')}`,
      `**Target:** ${targetUrl}`,
      ``,
      `The target is a Docker-hosted web application at ${targetUrl}.`,
      ``,
    ];

    // Analyze Docker environment once (used for both Target Environment and strategy selection)
    let envBrief = null;
    if (config.path) {
      try { envBrief = analyzeDockerEnvironment(config.path); } catch { /* graceful */ }
    }

    // Inject Docker environment brief if substantive
    if (envBrief) {
      const formatted = formatEnvironmentBrief(envBrief);
      if (formatted && (envBrief.framework || envBrief.database || envBrief.ports?.length > 0)) {
        parts.push('## Target Environment');
        parts.push(formatted);
        parts.push('');
      }
    }

    // Select and inject vuln-class-specific strategy template
    const strategy = selectStrategy({
      mode: 'red',
      tags: config.tags,
      envBrief,
      override: config.strategyOverride || null,
    });

    if (strategy) {
      parts.push(`## Attack Strategy (${strategy.name})`);
      parts.push(`*Selected via: ${strategy.source}*`);
      parts.push('');
      parts.push(formatStrategy(strategy.content));
      parts.push('');
    }

    // General methodology (supplements strategy template)
    parts.push('## Methodology');
    parts.push('1. **Recon:** `curl -s TARGET` to see the app. Check /robots.txt, /sitemap.xml, /.git/, /admin, view source.');
    parts.push('2. **Identify attack surface:** forms, API endpoints, file uploads, auth mechanisms, URL parameters.');
    parts.push('3. **Exploit:** based on tags and discovered attack surface. Try the obvious first, escalate complexity.');
    parts.push('4. **Extract flag:** check response bodies, headers, cookies, database entries, file contents.');
    parts.push('5. **If stuck:** try a completely different attack vector. Check source code if accessible. Look for backup files (.bak, .old, ~).');
    parts.push('');
    parts.push('## Auth Patterns');
    parts.push('- COOKIE: if server returns set-cookie, send it back: `Cookie: name=value`');
    parts.push('- BEARER: if response has a token, use `Authorization: Bearer TOKEN`');
    parts.push('- SESSION: some apps use session cookies — maintain them across requests');
    parts.push('');
    parts.push('## Tools Available');
    parts.push('You have full access to: curl, python3, netcat (nc), nmap, bash, and any installed CLI tools.');
    parts.push('Write Python exploit scripts when curl alone is insufficient.');
    parts.push('');
    parts.push('## Rules');
    parts.push('- DO NOT stop until you find the flag. Keep trying different approaches.');
    parts.push('- After each action, decide your NEXT action immediately — don\'t just summarize.');
    parts.push('- If you find a promising lead (login, IDOR, endpoint), EXPLOIT IT immediately.');
    parts.push('- The flag format is usually FLAG{...} or similar.');
    parts.push('');
    parts.push('Find the flag and output it as: FLAG_FOUND: <flag>');

    // Inject NERF knowledge and skills relevant to challenge tags
    const knowledge = loadKnowledge(config.tags);
    if (knowledge) {
      parts.push('');
      parts.push('## Reference Knowledge');
      parts.push(knowledge);
    }

    const skills = loadSkillSummaries(config.tags);
    if (skills) {
      parts.push('');
      parts.push('## Relevant NERF Skills');
      parts.push(skills);
    }

    return parts.join('\n');
  }

  /**
   * Execute the subagent. Tries the subagent tool first,
   * falls back to direct CLI execution.
   */
  async _execSubagent(task) {
    // Option 1: Direct subagent function (when called from within pi runtime)
    if (this._subagentFn) {
      const result = await this._subagentFn({
        agent: this._agent,
        task,
      });
      return typeof result === 'string' ? result : JSON.stringify(result);
    }

    // Option 2: Pi SDK (when pi-coding-agent is available)
    try {
      const { importPiCodingAgent } = await import('../gateway/pi-imports.js');
      const { createAgentSession, SessionManager, AuthStorage, ModelRegistry } =
        await importPiCodingAgent();

      const authStorage = AuthStorage.create();
      const modelRegistry = new ModelRegistry(authStorage);

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

      const textParts = [];
      session.subscribe((event) => {
        if (event.type === 'message_update' && event.assistantMessageEvent?.type === 'text_delta') {
          textParts.push(event.assistantMessageEvent.delta);
        }
      });

      // Set up timeout with cleanup
      let timeoutId;
      const timeoutPromise = new Promise((_, reject) => {
        timeoutId = setTimeout(() => reject(new Error(`Pi SDK solver timeout after ${this._timeoutS}s`)), this._timeoutS * 1000);
      });

      try {
        await Promise.race([session.prompt(task), timeoutPromise]);
      } finally {
        clearTimeout(timeoutId);
      }
      session.dispose();
      return textParts.join('');
    } catch (e) {
      if (e.message?.includes('timeout')) throw e;
      // Pi SDK not available — fall through to other options
    }

    // Option 3: Fallback — pi-mono not available, no subagent
    throw new Error(
      'No agent execution backend available. ' +
      'Ensure @nerf/pi-coding-agent is installed (npm install @nerf/pi-coding-agent).'
    );
  }
}
