/**
 * LangGraph state machine for Bandit Runner agent
 */

import { StateGraph, END, START } from "@langchain/langgraph"
import { HumanMessage, SystemMessage, AIMessage } from "@langchain/core/messages"
import type { BanditAgentState, Command, ThoughtLog } from "./bandit-state"
import { LEVEL_GOALS } from "./bandit-state"
import { createLLMProvider } from "./llm-provider"
import { banditTools } from "./tools"
import { ToolNode } from "@langchain/langgraph/prebuilt"

/**
 * System prompt for the Bandit Runner agent
 */
const SYSTEM_PROMPT = `You are BanditRunner, an autonomous operator tasked with solving the OverTheWire Bandit wargame.

IMPORTANT RULES:
1. You have access to SSH tools: ssh_connect, ssh_exec, validate_password, ssh_disconnect
2. Only commands from the allowlist are permitted (ls, cat, grep, find, base64, etc.)
3. Never use destructive commands (rm -rf, format, etc.)
4. Always validate passwords before advancing to the next level
5. Think step-by-step and explain your reasoning
6. If a command fails, analyze the error and try a different approach
7. Keep track of the current level and goal

WORKFLOW:
1. Plan: Analyze the current level goal and decide which command(s) to run
2. Execute: Run the command via ssh_exec
3. Validate: Check if the output contains the password
4. Advance: Validate the password and move to the next level

Remember: Passwords are typically 32-character alphanumeric strings.`

/**
 * Planning node - LLM decides next action
 */
async function planLevel(state: BanditAgentState): Promise<Partial<BanditAgentState>> {
  const { currentLevel, levelGoal, commandHistory, modelProvider, modelName, thoughts } = state

  // Create LLM provider
  const apiKey = process.env.OPENROUTER_API_KEY || ''
  const llm = createLLMProvider(modelProvider as 'openrouter', modelName, apiKey)

  // Build context from recent commands
  const recentCommands = commandHistory.slice(-5).map(cmd => 
    `Command: ${cmd.command}\nOutput: ${cmd.output.slice(0, 500)}\nExit Code: ${cmd.exitCode}`
  ).join('\n\n')

  const messages = [
    new SystemMessage(SYSTEM_PROMPT),
    new HumanMessage(`Current Level: ${currentLevel}
Goal: ${levelGoal}

Recent Commands:
${recentCommands || 'No commands executed yet.'}

What is your next step? Think through this carefully and decide on a command to execute.
Respond with your reasoning and then the exact command to run.`),
  ]

  const response = await llm.generateResponse(messages)

  // Add thought to log
  const thought: ThoughtLog = {
    type: 'plan',
    content: response,
    timestamp: new Date().toISOString(),
    level: currentLevel,
  }

  return {
    thoughts: [...thoughts, thought],
    status: 'executing',
  }
}

/**
 * Command execution node - Run SSH command
 */
async function executeCommand(state: BanditAgentState): Promise<Partial<BanditAgentState>> {
  const { thoughts, currentLevel, sshConnectionId } = state

  // Extract command from latest thought
  const latestThought = thoughts[thoughts.length - 1]
  const commandMatch = latestThought.content.match(/```(?:bash|sh)?\n(.+?)\n```/s) ||
                       latestThought.content.match(/(?:command|execute|run):\s*`?([^`\n]+)`?/i)
  
  if (!commandMatch) {
    return {
      status: 'failed',
      error: 'Could not extract command from LLM response',
      failureReasons: [...state.failureReasons, 'Command extraction failed'],
    }
  }

  const command = commandMatch[1].trim()

  // TODO: Actually call ssh_exec tool
  // For now, this is a placeholder
  const mockResult: Command = {
    command,
    output: `[SSH Proxy not yet implemented - command would execute: ${command}]`,
    exitCode: 0,
    timestamp: new Date().toISOString(),
    duration: 100,
    level: currentLevel,
  }

  return {
    commandHistory: [...state.commandHistory, mockResult],
    status: 'validating',
  }
}

/**
 * Validation node - Check if password was found
 */
async function validateResult(state: BanditAgentState): Promise<Partial<BanditAgentState>> {
  const { commandHistory, currentLevel, modelProvider, modelName, thoughts } = state

  const lastCommand = commandHistory[commandHistory.length - 1]
  
  // Use LLM to analyze output and extract password
  const apiKey = process.env.OPENROUTER_API_KEY || ''
  const llm = createLLMProvider(modelProvider as 'openrouter', modelName, apiKey)

  const messages = [
    new SystemMessage(SYSTEM_PROMPT),
    new HumanMessage(`Command executed: ${lastCommand.command}
Output: ${lastCommand.output}

Does this output contain the password for the next level?
Passwords are typically 32-character alphanumeric strings.

If you found a password, respond with: PASSWORD: <the_password>
If not found, respond with: NOT_FOUND and explain what to try next.`),
  ]

  const response = await llm.generateResponse(messages)

  const thought: ThoughtLog = {
    type: 'observation',
    content: response,
    timestamp: new Date().toISOString(),
    level: currentLevel,
  }

  // Check if password was found
  const passwordMatch = response.match(/PASSWORD:\s*([A-Za-z0-9+/=]{16,})/i)
  
  if (passwordMatch) {
    const password = passwordMatch[1]
    return {
      thoughts: [...thoughts, thought],
      nextPassword: password,
      status: 'advancing',
    }
  }

  // Password not found, retry if under limit
  const newRetryCount = state.retryCount + 1
  if (newRetryCount >= state.maxRetries) {
    return {
      thoughts: [...thoughts, thought],
      status: 'failed',
      error: `Max retries (${state.maxRetries}) reached for level ${currentLevel}`,
      failureReasons: [...state.failureReasons, `Level ${currentLevel} max retries exceeded`],
    }
  }

  return {
    thoughts: [...thoughts, thought],
    retryCount: newRetryCount,
    status: 'planning',
  }
}

/**
 * Advance level node - Validate password and move to next level
 */
async function advanceLevel(state: BanditAgentState): Promise<Partial<BanditAgentState>> {
  const { currentLevel, nextPassword, targetLevel } = state

  if (!nextPassword) {
    return {
      status: 'failed',
      error: 'No password to validate',
    }
  }

  // TODO: Actually validate password via SSH
  // For now, assume it's valid
  const nextLevel = currentLevel + 1

  // Check if we've reached target
  if (nextLevel > targetLevel) {
    return {
      status: 'complete',
      completedAt: new Date().toISOString(),
      currentLevel: nextLevel,
      currentPassword: nextPassword,
    }
  }

  // Move to next level
  const newGoal = LEVEL_GOALS[nextLevel] || 'Unknown level goal'

  return {
    currentLevel: nextLevel,
    currentPassword: nextPassword,
    nextPassword: null,
    levelGoal: newGoal,
    retryCount: 0,
    status: 'planning',
    lastCheckpoint: {
      level: currentLevel,
      password: nextPassword,
      timestamp: new Date().toISOString(),
      commandCount: state.commandHistory.length,
      state: { currentLevel: nextLevel, currentPassword: nextPassword },
    },
  }
}

/**
 * Conditional edge function - determines next node based on state
 */
function shouldContinue(state: BanditAgentState): string {
  if (state.status === 'complete' || state.status === 'failed') {
    return END
  }
  if (state.status === 'paused') {
    return 'paused'
  }
  if (state.status === 'planning') {
    return 'plan_level'
  }
  if (state.status === 'executing') {
    return 'execute_command'
  }
  if (state.status === 'validating') {
    return 'validate_result'
  }
  if (state.status === 'advancing') {
    return 'advance_level'
  }
  return END
}

/**
 * Create the Bandit Runner state graph
 */
export function createBanditGraph() {
  const workflow = new StateGraph<BanditAgentState>({
    channels: {
      runId: null,
      modelProvider: null,
      modelName: null,
      currentLevel: null,
      targetLevel: null,
      currentPassword: null,
      nextPassword: null,
      levelGoal: null,
      commandHistory: null,
      thoughts: null,
      status: null,
      retryCount: null,
      maxRetries: null,
      failureReasons: null,
      lastCheckpoint: null,
      streamingMode: null,
      sshConnectionId: null,
      startedAt: null,
      completedAt: null,
      error: null,
    },
  })

  // Add nodes
  workflow.addNode('plan_level', planLevel)
  workflow.addNode('execute_command', executeCommand)
  workflow.addNode('validate_result', validateResult)
  workflow.addNode('advance_level', advanceLevel)

  // Add edges
  workflow.addEdge(START, 'plan_level')
  workflow.addConditionalEdges('plan_level', shouldContinue)
  workflow.addConditionalEdges('execute_command', shouldContinue)
  workflow.addConditionalEdges('validate_result', shouldContinue)
  workflow.addConditionalEdges('advance_level', shouldContinue)

  return workflow.compile({
    // Enable checkpointing for pause/resume
    checkpointer: undefined, // Will be set in Durable Object
  })
}

