/**
 * Bandit Agent Durable Object
 * Runs LangGraph.js state machine and manages WebSocket connections
 */

import type { DurableObject, DurableObjectState } from "@cloudflare/workers-types"
import type { BanditAgentState, RunConfig, AgentEvent } from "../agents/bandit-state"
import { LEVEL_GOALS } from "../agents/bandit-state"
import { DOStorage } from "../storage/run-storage"

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

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

  /**
   * Handle HTTP requests and WebSocket upgrades
   */
  async fetch(request: Request): Promise<Response> {
    const url = new URL(request.url)

    // Handle WebSocket upgrade
    if (request.headers.get("Upgrade") === "websocket") {
      return this.handleWebSocket(request)
    }

    // 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 })
    }
  }

  /**
   * Handle WebSocket connection using Hibernatable WebSockets API
   */
  private handleWebSocket(request: Request): Response {
    const pair = new WebSocketPair()
    const [client, server] = Object.values(pair)

    // Use modern Hibernatable WebSockets API
    // This allows the DO to be evicted from memory during inactivity
    this.ctx.acceptWebSocket(server)

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

  /**
   * Handle incoming WebSocket messages (Hibernatable API handler)
   */
  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)
    }
  }

  /**
   * Handle WebSocket close (Hibernatable API handler)
   */
  async webSocketClose(ws: WebSocket, code: number, reason: string, wasClean: boolean): Promise<void> {
    console.log(`WebSocket closed: Code ${code}, Reason: ${reason}, Clean: ${wasClean}`)
    // Cleanup is automatic with Hibernatable WebSockets
  }

  /**
   * Handle WebSocket errors (Hibernatable API handler)
   */
  async webSocketError(ws: WebSocket, error: unknown): Promise<void> {
    console.error("WebSocket error:", error)
  }

  /**
   * Handle POST requests
   */
  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 })
  }

  /**
   * Handle GET requests
   */
  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 })
  }

  /**
   * Start a new agent run - delegate to SSH proxy
   */
  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" },
      })
    }

    // Initialize state
    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,
    }

    // Save initial state
    await this.storage.saveState(this.state)
    this.isRunning = true

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

    // Start agent run in SSH proxy (in background)
    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" },
    })
  }

  /**
   * Run agent via SSH proxy - streams JSONL events back
   */
  private async runAgentViaProxy(config: RunConfig) {
    try {
      const sshProxyUrl = this.env.SSH_PROXY_URL || 'https://bandit-ssh-proxy.fly.dev'
      
      // Call SSH proxy /agent/run endpoint
      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()}`)
      }

      // Stream JSONL events from SSH proxy
      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
        }

        // Decode chunk and add to buffer
        buffer += decoder.decode(value, { stream: true })
        
        // Process complete JSON lines
        const lines = buffer.split('\n')
        buffer = lines.pop() || '' // Keep incomplete line in buffer

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

          try {
            const event = JSON.parse(line)
            
            // Broadcast event to all WebSocket clients
            this.broadcast(event)

            // Update local state based on events
            this.updateStateFromEvent(event)
          } catch (parseError) {
            console.error('Failed to parse JSONL event:', line, parseError)
          }
        }
      }

      // Mark run as complete
      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)}`)
    }
  }

  /**
   * Update DO state based on events from SSH proxy
   */
  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
    }

    // Save updated state
    this.storage.saveState(this.state)
  }

  /**
   * Pause the current run
   */
  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" },
    })
  }

  /**
   * Resume a paused run
   */
  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(),
    })

    // Continue graph execution
    this.runGraph().catch(error => {
      console.error("Graph execution error:", error)
      this.handleError(error)
    })

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

  /**
   * Execute a manual command (human intervention)
   */
  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" },
      })
    }

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

    // TODO: Actually execute command via SSH proxy
    // For now, simulate execution
    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" },
    })
  }

  /**
   * Retry current level
   */
  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" },
    })
  }

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

  /**
   * Handle errors
   */
  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
  }

  /**
   * Broadcast event to all connected WebSocket clients
   * Uses Hibernatable WebSockets API
   */
  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)
        // With Hibernatable API, no need to manually delete from a Set
      }
    }
  }

  /**
   * Alarm handler for cleanup
   */
  async alarm() {
    // Auto-cleanup after 2 hours of inactivity
    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
      }
    }

    // Schedule next alarm in 1 hour
    await this.ctx.storage.setAlarm(Date.now() + 60 * 60 * 1000)
  }
}

