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

/**
 * NERF Benchmark — Security agent for auto solving.
 *
 * Runs a multi-turn tool-use conversation loop with an LLM,
 * dispatching tool calls to the sandbox. Supports file injection
 * from benchmark configs and multi-service targets.
 */

import { getToolsForWinCondition, dispatchTool } from './tools.js';
import { ConfidenceTracker, createCognitiveDispatcher } from './cognitive.js';
import { generateSystemPrompt } from './prompts.js';
import { join } from 'node:path';
import { existsSync } from 'node:fs';

export class AgentResult {
  constructor(opts = {}) {
    this.flagFound = opts.flagFound ?? false;
    this.flagValue = opts.flagValue ?? '';
    this.answerFound = opts.answerFound ?? false;
    this.answerValue = opts.answerValue ?? '';
    this.tokensIn = opts.tokensIn ?? 0;
    this.tokensOut = opts.tokensOut ?? 0;
    this.toolCalls = opts.toolCalls ?? 0;
    this.steps = opts.steps ?? [];
    this.error = opts.error ?? null;
  }
}

export class SecurityAgent {
  constructor({ client, model, sandbox, maxTurns = 30, preToolHook = null }) {
    this._client = client;
    this._model = model;
    this._sandbox = sandbox;
    this._maxTurns = maxTurns;
    this._preToolHook = preToolHook;
  }

  /**
   * Inject benchmark files into the sandbox before solving.
   *
   * @param {Array<{ name: string, path: string }>} files - Files from BenchmarkConfig
   * @param {string} benchmarkPath - Base path of the benchmark
   */
  injectFiles(files, benchmarkPath) {
    if (!files || files.length === 0) return;

    for (const file of files) {
      const hostPath = join(benchmarkPath, file.path || file.name || file);
      const containerPath = `/tmp/challenge/${file.name || file.path || file}`;
      if (existsSync(hostPath)) {
        try {
          // Ensure container directory exists
          const dir = containerPath.replace(/\/[^/]+$/, '');
          this._sandbox.execTool(`mkdir -p '${dir}'`);
          this._sandbox.copyFileIn(hostPath, containerPath);
        } catch (e) {
          // Non-fatal — log and continue
        }
      }
    }
  }

  /**
   * Run the agent against a target.
   *
   * @param {object} opts
   * @param {string} opts.targetUrl - Primary target URL
   * @param {string} opts.challengeDescription - Challenge description
   * @param {string} [opts.winCondition='flag'] - 'flag' or 'question'
   * @param {Array<{ name: string, url: string }>} [opts.serviceUrls] - All service URLs
   * @param {Array} [opts.files] - Benchmark files to inject
   * @param {string} [opts.benchmarkPath] - Benchmark base path for file injection
   * @returns {Promise<AgentResult>}
   */
  async run({ targetUrl, challengeDescription, winCondition = 'flag', serviceUrls = [], files = [], benchmarkPath = '', tags = [] }) {
    // Inject files if provided
    if (files.length > 0 && benchmarkPath) {
      this.injectFiles(files, benchmarkPath);
    }

    // Initialize cognitive architecture
    const tracker = new ConfidenceTracker();
    const cognitiveDispatch = createCognitiveDispatcher(tracker);

    // Generate tag-aware system prompt
    const systemPrompt = generateSystemPrompt({
      targetUrl,
      challengeDescription,
      tags,
      winCondition,
      serviceUrls,
      hasFiles: files.length > 0,
      phase: tracker.phase,
    });
    const tools = getToolsForWinCondition(winCondition);
    const messages = [{ role: 'user', content: systemPrompt }];
    let totalIn = 0;
    let totalOut = 0;
    let toolCallCount = 0;
    const steps = [];

    for (let turn = 0; turn < this._maxTurns; turn++) {
      let response;
      try {
        response = await this._client.messages.create({
          model: this._model,
          max_tokens: 4096,
          tools,
          messages,
        });
      } catch (err) {
        return new AgentResult({ error: `LLM error: ${err.message}`, steps });
      }

      totalIn += response.usage?.input_tokens || 0;
      totalOut += response.usage?.output_tokens || 0;

      const assistantContent = response.content || [];
      messages.push({ role: 'assistant', content: assistantContent });

      const toolUseBlocks = assistantContent.filter((b) => b.type === 'tool_use');

      if (toolUseBlocks.length === 0) {
        const text = assistantContent.find((b) => b.type === 'text')?.text || '';
        steps.push(`[text] ${text.slice(0, 200)}`);
        if (response.stop_reason === 'end_turn') {
          // Nudge the model to keep going if it hasn't found the flag yet
          if (turn < this._maxTurns - 1) {
            messages.push({
              role: 'user',
              content: 'You have not found the flag yet. Continue investigating — use your tools to take the next action. Do not just describe what you would do; actually do it by calling a tool.',
            });
            continue;
          }
          break;
        }
        continue;
      }

      const toolResults = [];
      for (const block of toolUseBlocks) {
        toolCallCount++;
        steps.push(`[tool] ${block.name}: ${JSON.stringify(block.input).slice(0, 150)}`);

        if (this._preToolHook) {
          const approved = await this._preToolHook(block.name, block.input);
          if (!approved) {
            toolResults.push({ type: 'tool_result', tool_use_id: block.id, content: 'Tool execution denied by supervisor.' });
            continue;
          }
        }

        const result = cognitiveDispatch(block.name, block.input) || dispatchTool(block.name, block.input, this._sandbox);
        steps.push(`[result] ${result.output.slice(0, 200)}`);

        if (result.flagSubmitted) {
          return new AgentResult({
            flagFound: true, flagValue: result.flagSubmitted,
            tokensIn: totalIn, tokensOut: totalOut, toolCalls: toolCallCount, steps,
          });
        }

        if (result.answerSubmitted) {
          return new AgentResult({
            answerFound: true, answerValue: result.answerSubmitted,
            tokensIn: totalIn, tokensOut: totalOut, toolCalls: toolCallCount, steps,
          });
        }

        toolResults.push({ type: 'tool_result', tool_use_id: block.id, content: result.output });
      }

      messages.push({ role: 'user', content: toolResults });
    }

    return new AgentResult({ tokensIn: totalIn, tokensOut: totalOut, toolCalls: toolCallCount, steps });
  }
}
