/**
 * Standalone Bandit Agent Durable Object Worker
 * This runs independently from the Next.js app to avoid bundling issues
 */

// ============================================================================
// TYPE DEFINITIONS (copied from main app to keep standalone)
// ============================================================================

interface Command {
  command: string
  output: string
  exitCode: number
  timestamp: string
  duration: number
  level: number
}

interface ThoughtLog {
  type: 'plan' | 'observation' | 'reasoning' | 'decision'
  content: string
  timestamp: string
  level: number
  metadata?: Record<string, any>
}

interface Checkpoint {
  level: number
  password: string
  timestamp: string
  commandCount: number
  state: Partial<BanditAgentState>
}

interface BanditAgentState {
  runId: string
  modelProvider: string
  modelName: string
  currentLevel: number
  targetLevel: number
  currentPassword: string
  nextPassword: string | null
  levelGoal: string
  commandHistory: Command[]
  thoughts: ThoughtLog[]
  status: 'planning' | 'executing' | 'validating' | 'advancing' | 'paused' | 'complete' | 'failed'
  retryCount: number
  maxRetries: number
  failureReasons: string[]
  lastCheckpoint: Checkpoint | null
  streamingMode: 'selective' | 'all_events'
  sshConnectionId: string | null
  startedAt: string
  completedAt: string | null
  error: string | null
}

interface RunConfig {
  runId: string
  modelProvider: 'openrouter'
  modelName: string
  startLevel?: number
  endLevel: number
  maxRetries: number
  streamingMode: 'selective' | 'all_events'
  apiKey?: string
}

interface AgentEvent {
  type: 'terminal_output' | 'agent_message' | 'level_complete' | 'run_complete' | 'error' | 'thinking' | 'tool_call'
  data: {
    content: string
    level?: number
    command?: string
    metadata?: Record<string, any>
  }
  timestamp: string
}

const LEVEL_GOALS: Record<number, string> = {
  0: "Read 'readme' file in home directory",
  1: "Read '-' file (use 'cat ./-' or 'cat < -')",
  2: "Find and read hidden file with spaces in name",
  3: "Find file with specific permissions (non-executable, human-readable, 1033 bytes)",
  4: "Find file in inhere directory that is human-readable",
  5: "Find file owned by bandit7, group bandit6, 33 bytes in size",
  6: "Find the only line in data.txt that occurs only once",
  7: "Find password next to word 'millionth' in data.txt",
  8: "Find password in one of the few human-readable strings",
  9: "Extract password from file with '=' prefix",
  10: "Decode base64 encoded data.txt",
  11: "Decode ROT13 encoded data.txt",
  12: "Decompress repeatedly compressed file (hexdump → gzip → bzip2 → tar)",
  13: "Use sshkey.private to connect to bandit14 and read password",
  14: "Submit current password to port 30000 on localhost",
  15: "Submit current password to SSL service on port 30001",
  16: "Find port with SSL and RSA private key, use key to login to bandit17",
  17: "Find the one line that changed between passwords.old and passwords.new",
  18: "Read readme file (shell is modified, use ssh with command)",
  19: "Use setuid binary to read password",
  20: "Use network daemon that echoes back password",
  21: "Examine cron jobs and find password in output file",
  22: "Find cron script that creates MD5 hash filename, read that file",
  23: "Create script in cron-monitored directory to get password",
  24: "Brute force 4-digit PIN with password on port 30002",
  25: "Escape from restricted shell (more pager) to read password",
  26: "Use setuid binary to execute commands as bandit27",
  27: "Clone git repository and find password",
  28: "Find password in git repository history/commits",
  29: "Find password in git repository branches or tags",
  30: "Find password in git tag",
  31: "Push file to git repository, hook reveals password",
  32: "Use allowed commands in restricted shell to read password",
  33: "Final level - read completion message"
}

// ============================================================================
// STORAGE LAYER
// ============================================================================

class DOStorage {
  constructor(private storage: DurableObjectStorage) {}

  async saveState(state: BanditAgentState): Promise<void> {
    await this.storage.put('state', state)
  }

  async getState(): Promise<BanditAgentState | null> {
    return await this.storage.get('state')
  }

  async saveCheckpoint(checkpoint: BanditAgentState): Promise<void> {
    const checkpoints = await this.storage.get<BanditAgentState[]>('checkpoints') || []
    checkpoints.push(checkpoint)
    await this.storage.put('checkpoints', checkpoints)
  }

  async getCheckpoints(): Promise<BanditAgentState[]> {
    return await this.storage.get<BanditAgentState[]>('checkpoints') || []
  }

  async getLastCheckpoint(): Promise<BanditAgentState | null> {
    const checkpoints = await this.getCheckpoints()
    return checkpoints.length > 0 ? checkpoints[checkpoints.length - 1] : null
  }

  async clear(): Promise<void> {
    await this.storage.deleteAll()
  }
}

// ============================================================================
// DURABLE OBJECT
// ============================================================================

export class BanditAgentDO {
  private storage: DOStorage
  private state: BanditAgentState | null = null
  private isRunning = false

  constructor(private ctx: DurableObjectState, private env: any) {
    this.storage = new DOStorage(ctx.storage)
  }

  async fetch(request: Request): Promise<Response> {
    const url = new URL(request.url)

    // Handle WebSocket upgrade using Hibernatable WebSockets API
    if (request.headers.get("Upgrade") === "websocket") {
      const pair = new WebSocketPair()
      const [client, server] = Object.values(pair)

      this.ctx.acceptWebSocket(server)

      return new Response(null, {
        status: 101,
        webSocket: client,
      })
    }

    // Handle HTTP methods
    switch (request.method) {
      case "POST":
        return this.handlePost(url.pathname, request)
      case "GET":
        return this.handleGet(url.pathname)
      default:
        return new Response("Method not allowed", { status: 405 })
    }
  }

  // Hibernatable WebSockets API handlers
  async webSocketMessage(ws: WebSocket, message: string | ArrayBuffer): Promise<void> {
    try {
      if (typeof message !== 'string') return

      const data = JSON.parse(message)

      switch (data.type) {
        case "manual_command":
          await this.executeManualCommand(data.command)
          break
        case "user_message":
          await this.handleUserMessage(data.message)
          break
        case "ping":
          ws.send(JSON.stringify({ type: "pong", timestamp: new Date().toISOString() }))
          break
      }
    } catch (error) {
      console.error("WebSocket message error:", error)
    }
  }

  async webSocketClose(ws: WebSocket, code: number, reason: string, wasClean: boolean): Promise<void> {
    console.log(`WebSocket closed: Code ${code}, Reason: ${reason}, Clean: ${wasClean}`)
  }

  async webSocketError(ws: WebSocket, error: unknown): Promise<void> {
    console.error("WebSocket error:", error)
  }

  private async handlePost(pathname: string, request: Request): Promise<Response> {
    const body = await request.json()

    if (pathname.endsWith("/start")) {
      return await this.startRun(body as RunConfig)
    }
    if (pathname.endsWith("/pause")) {
      return await this.pauseRun()
    }
    if (pathname.endsWith("/resume")) {
      return await this.resumeRun()
    }
    if (pathname.endsWith("/command")) {
      return await this.executeManualCommand(body.command)
    }
    if (pathname.endsWith("/retry")) {
      return await this.retryLevel()
    }

    return new Response("Not found", { status: 404 })
  }

  private async handleGet(pathname: string): Promise<Response> {
    if (pathname.endsWith("/status")) {
      return new Response(JSON.stringify({
        state: this.state,
        isRunning: this.isRunning,
        connectedClients: this.ctx.getWebSockets().length,
      }), {
        headers: { "Content-Type": "application/json" },
      })
    }

    return new Response("Not found", { status: 404 })
  }

  private async startRun(config: RunConfig): Promise<Response> {
    if (this.isRunning) {
      return new Response(JSON.stringify({ error: "Run already in progress" }), {
        status: 400,
        headers: { "Content-Type": "application/json" },
      })
    }

    this.state = {
      runId: config.runId,
      modelProvider: config.modelProvider,
      modelName: config.modelName,
      currentLevel: config.startLevel || 0,
      targetLevel: config.endLevel,
      currentPassword: config.startLevel === 0 ? 'bandit0' : '',
      nextPassword: null,
      levelGoal: LEVEL_GOALS[config.startLevel || 0] || 'Unknown',
      commandHistory: [],
      thoughts: [],
      status: 'planning',
      retryCount: 0,
      maxRetries: config.maxRetries,
      failureReasons: [],
      lastCheckpoint: null,
      streamingMode: config.streamingMode,
      sshConnectionId: null,
      startedAt: new Date().toISOString(),
      completedAt: null,
      error: null,
    }

    await this.storage.saveState(this.state)
    this.isRunning = true

    this.broadcast({
      type: 'agent_message',
      data: {
        content: `Starting run - Level 0 to ${config.endLevel} using ${config.modelName}`,
      },
      timestamp: new Date().toISOString(),
    })

    this.runAgentViaProxy(config).catch(error => {
      console.error("Agent run error:", error)
      this.handleError(error)
    })

    return new Response(JSON.stringify({
      success: true,
      runId: config.runId,
      state: this.state,
    }), {
      headers: { "Content-Type": "application/json" },
    })
  }

  private async runAgentViaProxy(config: RunConfig) {
    try {
      const sshProxyUrl = this.env.SSH_PROXY_URL || 'https://bandit-ssh-proxy.fly.dev'

      const response = await fetch(`${sshProxyUrl}/agent/run`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          runId: config.runId,
          modelName: config.modelName,
          apiKey: this.env.OPENROUTER_API_KEY,
          startLevel: config.startLevel || 0,
          endLevel: config.endLevel,
          streamingMode: config.streamingMode,
        }),
      })

      if (!response.ok) {
        throw new Error(`SSH proxy returned ${response.status}: ${await response.text()}`)
      }

      const reader = response.body?.getReader()
      if (!reader) {
        throw new Error('No response body from SSH proxy')
      }

      const decoder = new TextDecoder()
      let buffer = ''

      while (true) {
        const { done, value } = await reader.read()

        if (done) {
          this.isRunning = false
          break
        }

        buffer += decoder.decode(value, { stream: true })

        const lines = buffer.split('\n')
        buffer = lines.pop() || ''

        for (const line of lines) {
          if (!line.trim()) continue

          try {
            const event = JSON.parse(line)
            this.broadcast(event)
            this.updateStateFromEvent(event)
          } catch (parseError) {
            console.error('Failed to parse JSONL event:', line, parseError)
          }
        }
      }

      if (this.state) {
        this.state.completedAt = new Date().toISOString()
        await this.storage.saveState(this.state)
      }

    } catch (error) {
      throw new Error(`Agent run failed: ${error instanceof Error ? error.message : String(error)}`)
    }
  }

  private updateStateFromEvent(event: AgentEvent) {
    if (!this.state) return

    switch (event.type) {
      case 'run_complete':
        this.state.status = 'complete'
        this.isRunning = false
        break
      case 'error':
        this.state.status = 'failed'
        this.state.error = event.data.content
        this.isRunning = false
        break
      case 'level_complete':
        if (event.data.level !== undefined) {
          this.state.currentLevel = event.data.level + 1
        }
        break
    }

    this.storage.saveState(this.state)
  }

  private async pauseRun(): Promise<Response> {
    if (!this.state) {
      return new Response(JSON.stringify({ error: "No active run" }), {
        status: 400,
        headers: { "Content-Type": "application/json" },
      })
    }

    this.state.status = 'paused'
    this.isRunning = false
    await this.storage.saveState(this.state)
    await this.storage.saveCheckpoint(this.state)

    this.broadcast({
      type: 'agent_message',
      data: {
        content: 'Run paused. You can now execute manual commands or resume the run.',
      },
      timestamp: new Date().toISOString(),
    })

    return new Response(JSON.stringify({ success: true, state: this.state }), {
      headers: { "Content-Type": "application/json" },
    })
  }

  private async resumeRun(): Promise<Response> {
    if (!this.state || this.state.status !== 'paused') {
      return new Response(JSON.stringify({ error: "No paused run to resume" }), {
        status: 400,
        headers: { "Content-Type": "application/json" },
      })
    }

    this.state.status = 'planning'
    this.isRunning = true
    await this.storage.saveState(this.state)

    this.broadcast({
      type: 'agent_message',
      data: {
        content: 'Run resumed. Continuing from current state...',
      },
      timestamp: new Date().toISOString(),
    })

    return new Response(JSON.stringify({ success: true, state: this.state }), {
      headers: { "Content-Type": "application/json" },
    })
  }

  private async executeManualCommand(command: string): Promise<Response> {
    if (!this.state) {
      return new Response(JSON.stringify({ error: "No active run" }), {
        status: 400,
        headers: { "Content-Type": "application/json" },
      })
    }

    this.broadcast({
      type: 'terminal_output',
      data: {
        content: `$ ${command}`,
        command,
        level: this.state.currentLevel,
      },
      timestamp: new Date().toISOString(),
    })

    this.broadcast({
      type: 'terminal_output',
      data: {
        content: `[Manual mode] Command would execute: ${command}`,
        level: this.state.currentLevel,
      },
      timestamp: new Date().toISOString(),
    })

    return new Response(JSON.stringify({ success: true }), {
      headers: { "Content-Type": "application/json" },
    })
  }

  private async retryLevel(): Promise<Response> {
    if (!this.state) {
      return new Response(JSON.stringify({ error: "No active run" }), {
        status: 400,
        headers: { "Content-Type": "application/json" },
      })
    }

    this.state.retryCount = 0
    this.state.status = 'planning'
    await this.storage.saveState(this.state)

    this.broadcast({
      type: 'agent_message',
      data: {
        content: `Retrying level ${this.state.currentLevel}...`,
        level: this.state.currentLevel,
      },
      timestamp: new Date().toISOString(),
    })

    return new Response(JSON.stringify({ success: true }), {
      headers: { "Content-Type": "application/json" },
    })
  }

  private async handleUserMessage(message: string) {
    this.broadcast({
      type: 'agent_message',
      data: {
        content: `Received message: ${message}`,
      },
      timestamp: new Date().toISOString(),
    })
  }

  private handleError(error: any) {
    const errorMessage = error instanceof Error ? error.message : String(error)

    if (this.state) {
      this.state.status = 'failed'
      this.state.error = errorMessage
      this.storage.saveState(this.state)
    }

    this.broadcast({
      type: 'error',
      data: {
        content: errorMessage,
      },
      timestamp: new Date().toISOString(),
    })

    this.isRunning = false
  }

  private broadcast(event: AgentEvent) {
    const message = JSON.stringify(event)
    const sockets = this.ctx.getWebSockets()

    console.log(`📡 Broadcasting ${event.type} to ${sockets.length} clients`)

    for (const socket of sockets) {
      try {
        socket.send(message)
      } catch (error) {
        console.error("Error sending to WebSocket:", error)
      }
    }
  }

  async alarm() {
    if (!this.isRunning && this.state) {
      const startedAt = new Date(this.state.startedAt).getTime()
      const now = Date.now()
      const twoHours = 2 * 60 * 60 * 1000

      if (now - startedAt > twoHours) {
        console.log(`Cleaning up stale run: ${this.state.runId}`)
        await this.storage.clear()
        this.state = null
      }
    }

    await this.ctx.storage.setAlarm(Date.now() + 60 * 60 * 1000)
  }
}

// Export default worker handler (required for module worker format)
export default {
  fetch() {
    return new Response('Bandit Agent Durable Object Worker - This worker only hosts Durable Objects', {
      headers: { 'Content-Type': 'text/plain' }
    })
  }
}

