/**
 * Storage layer abstraction for run data
 * Handles DO → D1 → R2 data lifecycle
 */

import type { BanditAgentState, Command } from "../agents/bandit-state"

export interface RunMetadata {
  runId: string
  modelProvider: string
  modelName: string
  startLevel: number
  endLevel: number
  currentLevel: number
  status: string
  startedAt: string
  completedAt: string | null
  commandCount: number
  totalCost: number
  error: string | null
}

/**
 * Durable Object Storage Interface
 */
export 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()
  }
}

/**
 * D1 Database Interface (for historical data)
 */
export class D1Storage {
  constructor(private db: D1Database) {}

  async saveRunMetadata(metadata: RunMetadata): Promise<void> {
    await this.db
      .prepare(
        `INSERT OR REPLACE INTO runs 
        (run_id, model_provider, model_name, start_level, end_level, current_level, 
         status, started_at, completed_at, command_count, total_cost, error)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
      )
      .bind(
        metadata.runId,
        metadata.modelProvider,
        metadata.modelName,
        metadata.startLevel,
        metadata.endLevel,
        metadata.currentLevel,
        metadata.status,
        metadata.startedAt,
        metadata.completedAt,
        metadata.commandCount,
        metadata.totalCost,
        metadata.error
      )
      .run()
  }

  async getRunMetadata(runId: string): Promise<RunMetadata | null> {
    const result = await this.db
      .prepare('SELECT * FROM runs WHERE run_id = ?')
      .bind(runId)
      .first<RunMetadata>()
    return result
  }

  async listRuns(limit = 50, offset = 0): Promise<RunMetadata[]> {
    const { results } = await this.db
      .prepare('SELECT * FROM runs ORDER BY started_at DESC LIMIT ? OFFSET ?')
      .bind(limit, offset)
      .all<RunMetadata>()
    return results
  }

  async saveCommand(runId: string, command: Command): Promise<void> {
    await this.db
      .prepare(
        `INSERT INTO commands 
        (run_id, command, output, exit_code, timestamp, duration, level)
        VALUES (?, ?, ?, ?, ?, ?, ?)`
      )
      .bind(
        runId,
        command.command,
        command.output,
        command.exitCode,
        command.timestamp,
        command.duration,
        command.level
      )
      .run()
  }
}

/**
 * R2 Storage Interface (for logs and artifacts)
 */
export class R2Storage {
  constructor(private bucket: R2Bucket) {}

  async saveJSONLLog(runId: string, lines: string[]): Promise<void> {
    const content = lines.join('\n')
    await this.bucket.put(`logs/${runId}.jsonl`, content, {
      httpMetadata: {
        contentType: 'application/x-ndjson',
      },
    })
  }

  async appendJSONLLine(runId: string, line: string): Promise<void> {
    // Read existing log
    const existing = await this.bucket.get(`logs/${runId}.jsonl`)
    const content = existing ? await existing.text() : ''
    
    // Append new line
    const updated = content ? `${content}\n${line}` : line
    await this.bucket.put(`logs/${runId}.jsonl`, updated, {
      httpMetadata: {
        contentType: 'application/x-ndjson',
      },
    })
  }

  async getJSONLLog(runId: string): Promise<string | null> {
    const object = await this.bucket.get(`logs/${runId}.jsonl`)
    return object ? await object.text() : null
  }

  async savePasswordVault(runId: string, passwords: Record<number, string>, encryptionKey: string): Promise<void> {
    // Simple encryption (in production, use proper encryption)
    const encrypted = await this.encrypt(JSON.stringify(passwords), encryptionKey)
    
    await this.bucket.put(`passwords/${runId}.enc`, encrypted, {
      httpMetadata: {
        contentType: 'application/octet-stream',
      },
      customMetadata: {
        'ttl': String(Date.now() + 2 * 60 * 60 * 1000), // 2 hour TTL
      },
    })
  }

  private async encrypt(data: string, key: string): Promise<string> {
    // TODO: Implement proper encryption using Web Crypto API
    // For now, just base64 encode
    return btoa(data)
  }

  private async decrypt(data: string, key: string): Promise<string> {
    // TODO: Implement proper decryption
    return atob(data)
  }
}

/**
 * D1 Schema migrations
 */
export const D1_SCHEMA = `
CREATE TABLE IF NOT EXISTS runs (
  run_id TEXT PRIMARY KEY,
  model_provider TEXT NOT NULL,
  model_name TEXT NOT NULL,
  start_level INTEGER NOT NULL,
  end_level INTEGER NOT NULL,
  current_level INTEGER NOT NULL,
  status TEXT NOT NULL,
  started_at TEXT NOT NULL,
  completed_at TEXT,
  command_count INTEGER DEFAULT 0,
  total_cost REAL DEFAULT 0.0,
  error TEXT
);

CREATE TABLE IF NOT EXISTS commands (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  run_id TEXT NOT NULL,
  command TEXT NOT NULL,
  output TEXT,
  exit_code INTEGER,
  timestamp TEXT NOT NULL,
  duration INTEGER,
  level INTEGER,
  FOREIGN KEY (run_id) REFERENCES runs(run_id)
);

CREATE INDEX IF NOT EXISTS idx_runs_started_at ON runs(started_at DESC);
CREATE INDEX IF NOT EXISTS idx_commands_run_id ON commands(run_id);
`

