/**
 * LangGraph tool wrappers for SSH operations
 */

import { tool } from "@langchain/core/tools"
import { z } from "zod"

// SSH Proxy configuration
const SSH_PROXY_URL = process.env.SSH_PROXY_URL || 'http://localhost:3001'
const SSH_TARGET_HOST = 'bandit.labs.overthewire.org'
const SSH_TARGET_PORT = 2220

export interface SSHConnectionResult {
  connectionId: string
  success: boolean
  message: string
}

export interface SSHCommandResult {
  output: string
  exitCode: number
  success: boolean
  duration: number
}

/**
 * Command allowlist based on system prompt
 * These are the only commands the agent is allowed to execute
 */
const ALLOWED_COMMANDS = [
  // File operations
  'ls', 'cat', 'grep', 'find', 'file', 'strings', 'wc', 'head', 'tail',
  // Text processing
  'sort', 'uniq', 'cut', 'tr', 'sed', 'awk',
  // Encoding/Compression
  'base64', 'xxd', 'gunzip', 'bunzip2', 'tar',
  // Network
  'nc', 'nmap', 'ssh', 'openssl',
  // Git
  'git',
  // System
  'chmod', 'pwd', 'cd', 'mkdir', 'mktemp', 'echo', 'printf',
  // Special
  'diff', 'md5sum', 'timeout'
]

/**
 * Validate command against allowlist
 */
function validateCommand(command: string): { valid: boolean; reason?: string } {
  const cmd = command.trim().split(/\s+/)[0]
  
  // Check if command starts with allowed prefix
  const isAllowed = ALLOWED_COMMANDS.some(allowed => cmd.startsWith(allowed))
  
  if (!isAllowed) {
    return { valid: false, reason: `Command '${cmd}' is not in the allowlist` }
  }
  
  // Additional safety checks
  if (command.includes('rm -rf') || command.includes('rm -r')) {
    return { valid: false, reason: 'Destructive rm commands are not allowed' }
  }
  
  if (command.includes('>') && !command.includes('2>')) {
    return { valid: false, reason: 'File write operations are restricted' }
  }
  
  return { valid: true }
}

/**
 * Connect to SSH server
 */
export const sshConnectTool = tool(
  async ({ username, password }) => {
    try {
      const response = await fetch(`${SSH_PROXY_URL}/ssh/connect`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          host: SSH_TARGET_HOST,
          port: SSH_TARGET_PORT,
          username,
          password,
        }),
      })

      const result: SSHConnectionResult = await response.json()
      return JSON.stringify(result)
    } catch (error) {
      return JSON.stringify({
        connectionId: null,
        success: false,
        message: `Connection failed: ${error instanceof Error ? error.message : 'Unknown error'}`,
      })
    }
  },
  {
    name: "ssh_connect",
    description: "Connect to the Bandit SSH server with username and password. Returns a connection ID for subsequent commands.",
    schema: z.object({
      username: z.string().describe("SSH username (e.g., 'bandit0', 'bandit1')"),
      password: z.string().describe("SSH password for the user"),
    }),
  }
)

/**
 * Execute command via SSH
 */
export const sshExecTool = tool(
  async ({ connectionId, command }) => {
    // Validate command
    const validation = validateCommand(command)
    if (!validation.valid) {
      return JSON.stringify({
        output: '',
        exitCode: 127,
        success: false,
        duration: 0,
        error: validation.reason,
      })
    }

    try {
      const startTime = Date.now()
      const response = await fetch(`${SSH_PROXY_URL}/ssh/exec`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          connectionId,
          command,
          timeout: 30000, // 30 second timeout
        }),
      })

      const result: SSHCommandResult = await response.json()
      result.duration = Date.now() - startTime
      return JSON.stringify(result)
    } catch (error) {
      return JSON.stringify({
        output: '',
        exitCode: 1,
        success: false,
        duration: 0,
        error: `Execution failed: ${error instanceof Error ? error.message : 'Unknown error'}`,
      })
    }
  },
  {
    name: "ssh_exec",
    description: "Execute a command on the SSH server. Only commands from the allowlist are permitted.",
    schema: z.object({
      connectionId: z.string().describe("The SSH connection ID from ssh_connect"),
      command: z.string().describe("The command to execute (must be in allowlist)"),
    }),
  }
)

/**
 * Validate password by attempting SSH connection
 */
export const validatePasswordTool = tool(
  async ({ level, password }) => {
    try {
      const username = `bandit${level}`
      const response = await fetch(`${SSH_PROXY_URL}/ssh/connect`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          host: SSH_TARGET_HOST,
          port: SSH_TARGET_PORT,
          username,
          password,
          testOnly: true, // Just test connection, don't keep it open
        }),
      })

      const result: SSHConnectionResult = await response.json()
      
      // Disconnect immediately if successful
      if (result.success && result.connectionId) {
        await fetch(`${SSH_PROXY_URL}/ssh/disconnect`, {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({ connectionId: result.connectionId }),
        })
      }

      return JSON.stringify({
        valid: result.success,
        level,
        message: result.message,
      })
    } catch (error) {
      return JSON.stringify({
        valid: false,
        level,
        message: `Validation failed: ${error instanceof Error ? error.message : 'Unknown error'}`,
      })
    }
  },
  {
    name: "validate_password",
    description: "Validate a password for a specific Bandit level by attempting an SSH connection",
    schema: z.object({
      level: z.number().describe("The Bandit level number"),
      password: z.string().describe("The password to validate"),
    }),
  }
)

/**
 * Disconnect SSH connection
 */
export const sshDisconnectTool = tool(
  async ({ connectionId }) => {
    try {
      const response = await fetch(`${SSH_PROXY_URL}/ssh/disconnect`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ connectionId }),
      })

      const result = await response.json()
      return JSON.stringify(result)
    } catch (error) {
      return JSON.stringify({
        success: false,
        message: `Disconnect failed: ${error instanceof Error ? error.message : 'Unknown error'}`,
      })
    }
  },
  {
    name: "ssh_disconnect",
    description: "Close an SSH connection",
    schema: z.object({
      connectionId: z.string().describe("The SSH connection ID to close"),
    }),
  }
)

/**
 * All available tools for the agent
 */
export const banditTools = [
  sshConnectTool,
  sshExecTool,
  validatePasswordTool,
  sshDisconnectTool,
]

