// PRE Web GUI — Background Agent System
// Fire-and-forget autonomous agents that run in the background while the user
// continues chatting. GPU-serialized via waitForIdle() — same as cron jobs.
//
// Smart Interrupt: Agent state is persisted to disk on every status change.
// Turn-level checkpoints let interrupted agents resume from where they left off
// after sleep, reboot, or crash. Sleep/wake detection pauses agents during
// system hibernation and resumes them when the system wakes.

const fs = require('fs');
const path = require('path');
const { createSession, appendMessage, renameSession, ensureProject, moveSessionToProject } = require('./sessions');
const { runToolLoop } = require('./tools');
const { slotManager } = require('./inference-queue');
const { PRE_DIR } = require('./constants');


// ── Persistence Paths ────────────────────────────────────────────────────────

const STATE_FILE = path.join(PRE_DIR, 'agent-state.json');
const CHECKPOINT_DIR = path.join(PRE_DIR, 'agent-checkpoints');


// ── Agent Registry ───────────────────────────────────────────────────────────

const _agents = new Map();  // id → agent state
let _agentSeq = 0;
let _broadcastWS = null;    // Set by server.js to push events to connected clients

function setBroadcast(fn) {
  _broadcastWS = fn;
}

// ── Disk Persistence ─────────────────────────────────────────────────────────

function _ensureDirs() {
  try { fs.mkdirSync(CHECKPOINT_DIR, { recursive: true }); } catch (_e) {}
}

/**
 * Save the entire agent registry to disk. Called on every state change.
 * Only persists fields needed for recovery — strips runtime objects.
 */
function _persistState() {
  try {
    const agents = [..._agents.values()].map(a => ({
      id: a.id,
      name: a.name,
      task: a.task,
      sessionId: a.sessionId,
      cwd: a.cwd,
      status: a.status,
      createdAt: a.createdAt,
      startedAt: a.startedAt,
      completedAt: a.completedAt,
      result: a.result ? a.result.slice(0, 2000) : null,
      error: a.error,
      toolCalls: a.toolCalls,
      tokens: a.tokens,
      checkpoint: a.checkpoint || null,
    }));
    fs.writeFileSync(STATE_FILE, JSON.stringify({ agents, seq: _agentSeq }, null, 2));
  } catch (err) {
    console.log(`[bg-agent] Failed to persist state: ${err.message}`);
  }
}

/**
 * Load agent state from disk. Called once at startup.
 * Returns agents that were interrupted (running/queued when process died).
 */
function _loadState() {
  try {
    if (!fs.existsSync(STATE_FILE)) return [];
    const data = JSON.parse(fs.readFileSync(STATE_FILE, 'utf-8'));
    if (data.seq) _agentSeq = data.seq;

    const interrupted = [];
    for (const a of (data.agents || [])) {
      // Restore completed/failed/cancelled agents to registry (for history)
      _agents.set(a.id, {
        ...a,
        _abort: null,
      });

      // Identify agents that need recovery
      if (a.status === 'running' || a.status === 'queued') {
        interrupted.push(a);
      }
    }
    return interrupted;
  } catch (err) {
    console.log(`[bg-agent] Failed to load state: ${err.message}`);
    return [];
  }
}

/**
 * Save a turn-level checkpoint for an agent.
 * Stores the turn number so we know where the agent left off.
 */
function _saveCheckpoint(agentId, turn) {
  try {
    _ensureDirs();
    const cpFile = path.join(CHECKPOINT_DIR, `${agentId}.json`);
    fs.writeFileSync(cpFile, JSON.stringify({ agentId, turn, savedAt: Date.now() }));
    // Also update in-memory state
    const agent = _agents.get(agentId);
    if (agent) agent.checkpoint = { turn, savedAt: Date.now() };
  } catch (err) {
    console.log(`[bg-agent] Checkpoint save failed for ${agentId}: ${err.message}`);
  }
}

/**
 * Load a checkpoint for an agent. Returns { turn } or null.
 */
function _loadCheckpoint(agentId) {
  try {
    const cpFile = path.join(CHECKPOINT_DIR, `${agentId}.json`);
    if (!fs.existsSync(cpFile)) return null;
    return JSON.parse(fs.readFileSync(cpFile, 'utf-8'));
  } catch {
    return null;
  }
}

/**
 * Remove checkpoint file after agent completes.
 */
function _clearCheckpoint(agentId) {
  try {
    const cpFile = path.join(CHECKPOINT_DIR, `${agentId}.json`);
    if (fs.existsSync(cpFile)) fs.unlinkSync(cpFile);
  } catch (_e) {}
}


// ── Background Execution Queue ───────────────────────────────────────────────
// Serializes background agent execution so only one runs at a time,
// and all wait for any active user conversation to finish first.

const _queue = [];
let _queueRunning = false;

async function _drainQueue() {
  if (_queueRunning) return;
  _queueRunning = true;
  while (_queue.length > 0) {
    const { agent } = _queue.shift();
    try {
      const slotHandle = await slotManager.acquireSlot('background');
      try {
        await _executeAgent(agent, slotHandle);
      } finally {
        slotHandle.release();
      }
    } catch (err) {
      agent.status = 'failed';
      agent.error = err.message;
      agent.completedAt = Date.now();
      _persistState();
      _notify(agent, 'failed');
      console.log(`[bg-agent] ${agent.id} failed: ${err.message}`);
    }
  }
  _queueRunning = false;
}

// ── Sleep/Wake Detection ─────────────────────────────────────────────────────
// Uses timer drift to detect system sleep. If a 10s interval fires after 30s+,
// the system likely slept. On wake, we health-check Ollama before resuming.

let _sleepDetectorInterval = null;
let _lastTick = Date.now();
const SLEEP_THRESHOLD_MS = 30000; // 30s gap = likely slept
let _systemAwake = true;

function _startSleepDetector() {
  _lastTick = Date.now();
  _sleepDetectorInterval = setInterval(async () => {
    const now = Date.now();
    const drift = now - _lastTick;
    _lastTick = now;

    if (drift > SLEEP_THRESHOLD_MS) {
      console.log(`[bg-agent] Sleep detected (${Math.round(drift / 1000)}s gap). Checking Ollama health...`);
      _systemAwake = false;

      // Pause any running agent — it likely has stale connections
      for (const agent of _agents.values()) {
        if (agent.status === 'running') {
          agent.status = 'paused';
          _notify(agent, 'paused');
          console.log(`[bg-agent] ${agent.id} paused (system slept)`);
        }
      }

      // Wait for Ollama to be healthy again
      const healthy = await _waitForOllama(60000);
      _systemAwake = true;

      if (healthy) {
        console.log('[bg-agent] Ollama healthy after wake. Resuming paused agents...');
        _resumePausedAgents();
      } else {
        console.log('[bg-agent] Ollama unhealthy after wake. Marking paused agents as interrupted.');
        for (const agent of _agents.values()) {
          if (agent.status === 'paused') {
            agent.status = 'interrupted';
            agent.error = 'System slept and Ollama did not recover';
            agent.completedAt = Date.now();
            _persistState();
            _notify(agent, 'failed');
          }
        }
      }
    }
  }, 10000);
}

function _stopSleepDetector() {
  if (_sleepDetectorInterval) {
    clearInterval(_sleepDetectorInterval);
    _sleepDetectorInterval = null;
  }
}

/**
 * Poll Ollama health with retries until it responds or timeout.
 */
async function _waitForOllama(timeoutMs) {
  const { healthCheck } = require('./ollama');
  const deadline = Date.now() + timeoutMs;
  let attempt = 0;
  while (Date.now() < deadline) {
    attempt++;
    try {
      const result = await healthCheck();
      if (result.ok) return true;
    } catch (_e) {}
    // Exponential backoff: 1s, 2s, 4s, 8s, capped at 10s
    const delay = Math.min(1000 * Math.pow(2, attempt - 1), 10000);
    await new Promise(r => setTimeout(r, delay));
  }
  return false;
}

/**
 * Resume agents that were paused due to system sleep.
 * Re-queues them from their last checkpoint.
 */
function _resumePausedAgents() {
  for (const agent of _agents.values()) {
    if (agent.status === 'paused') {
      agent.status = 'queued';
      agent._resumeFromCheckpoint = true;
      _persistState();
      _notify(agent, 'queued');
      console.log(`[bg-agent] ${agent.id} re-queued for resume`);
      _queue.push({ agent });
    }
  }
  _drainQueue();
}


// ── Agent Lifecycle ──────────────────────────────────────────────────────────

const AGENT_TIMEOUT_MS = 10 * 60 * 1000; // 10 minutes max per agent

/**
 * Spawn a background agent. Returns immediately with the agent ID.
 * The agent runs in the background and notifies on completion.
 *
 * @param {Object} opts
 * @param {string} opts.task - What the agent should do
 * @param {string} [opts.cwd] - Working directory
 * @param {string} [opts.name] - Display name for the agent
 * @returns {{ id: string, status: string }}
 */
function spawn({ task, cwd, name }) {
  if (!task || !task.trim()) {
    throw new Error('task is required');
  }

  const id = `bg_${++_agentSeq}`;
  const displayName = name || task.slice(0, 60);

  // Create a dedicated session for this agent
  const ts = Date.now().toString(36);
  const sessionId = createSession('web', `bg-agent-${id}-${ts}`, true);
  renameSession(sessionId, `Agent: ${displayName}`);
  ensureProject('background-agents', 'Background Agents');
  moveSessionToProject(sessionId, 'background-agents');
  appendMessage(sessionId, { role: 'user', content: task });

  const agent = {
    id,
    name: displayName,
    task,
    sessionId,
    cwd: cwd || process.env.HOME || '/tmp',
    status: 'queued',        // queued → running → completed/failed/cancelled/paused/interrupted
    createdAt: Date.now(),
    startedAt: null,
    completedAt: null,
    result: null,
    error: null,
    toolCalls: 0,            // Count of tools executed
    currentTool: null,       // Currently executing tool name
    tokens: { in: 0, out: 0 },
    checkpoint: null,        // { turn, savedAt } — last completed turn
  };

  _agents.set(id, agent);
  _persistState();
  _notify(agent, 'queued');
  console.log(`[bg-agent] ${id} queued: "${displayName}"`);

  // Enqueue for execution
  _queue.push({ agent });
  _drainQueue();

  return { id, status: 'queued', sessionId };
}

/**
 * Execute an agent through the full tool loop.
 */
async function _executeAgent(agent, slotHandle) {
  agent.status = 'running';
  agent.startedAt = agent.startedAt || Date.now();
  _persistState();
  _notify(agent, 'running');
  console.log(`[bg-agent] ${agent.id} running`);

  const abortController = new AbortController();
  const timeout = setTimeout(() => abortController.abort(), AGENT_TIMEOUT_MS);

  // Store abort controller on agent for cancel support
  agent._abort = abortController;

  let collectedTokens = [];

  // Determine if this is a resume — skip turns already completed
  const checkpoint = agent._resumeFromCheckpoint ? _loadCheckpoint(agent.id) : null;
  const skipTurns = checkpoint ? checkpoint.turn : 0;
  if (skipTurns > 0) {
    console.log(`[bg-agent] ${agent.id} resuming from turn ${skipTurns}`);
    delete agent._resumeFromCheckpoint;
  }

  try {
    await runToolLoop({
      sessionId: agent.sessionId,
      cwd: agent.cwd,
      signal: abortController.signal,
      userMessage: skipTurns > 0 ? null : agent.task, // Don't re-send task on resume
      needsTitle: false,
      skipTurns,
      send: (event) => {
        // Track progress
        if (event.type === 'token' && event.content) {
          collectedTokens.push(event.content);
        }
        if (event.type === 'tool_call') {
          agent.toolCalls++;
          agent.currentTool = event.name || null;
          _notify(agent, 'progress');
        }
        if (event.type === 'tool_result') {
          agent.currentTool = null;
        }
        if (event.type === 'done') {
          if (event.stats) {
            agent.tokens.in += event.stats.prompt_eval_count || 0;
            agent.tokens.out += event.stats.eval_count || 0;
          }
        }
        // Forward to connected clients (prefixed so UI can filter)
        if (_broadcastWS) {
          _broadcastWS({
            ...event,
            bgAgentId: agent.id,
            bgAgentName: agent.name,
          });
        }
      },
      onConfirmRequest: async () => {
        // Auto-deny confirmation-required tools in background mode
        console.log(`[bg-agent] ${agent.id} auto-denied confirmation — background mode`);
        return false;
      },
      onCheckpoint: (turn) => {
        // Called after each tool turn completes — save progress
        _saveCheckpoint(agent.id, turn);
        _persistState();
      },
      _slotHandle: slotHandle,
    });

    agent.status = 'completed';
    agent.result = collectedTokens.join('') || '(Agent completed with no text response)';
    agent.completedAt = Date.now();
    agent.currentTool = null;
    delete agent._abort;

    // Save result to session
    appendMessage(agent.sessionId, {
      role: 'assistant',
      content: `[Background agent result]\n${agent.result}`,
    });

    _clearCheckpoint(agent.id);
    _persistState();
    _notify(agent, 'completed');
    console.log(`[bg-agent] ${agent.id} completed in ${((agent.completedAt - agent.startedAt) / 1000).toFixed(1)}s, ${agent.toolCalls} tool calls`);

  } catch (err) {
    const partial = collectedTokens.join('');
    if (abortController.signal.aborted && agent.status === 'cancelled') {
      // User cancelled
      agent.result = partial || '(Cancelled by user)';
      agent.completedAt = Date.now();
      _clearCheckpoint(agent.id);
      _persistState();
      _notify(agent, 'cancelled');
      console.log(`[bg-agent] ${agent.id} cancelled`);
    } else if (abortController.signal.aborted && partial) {
      // Timeout with partial results
      agent.status = 'completed';
      agent.result = partial + '\n\n*(Agent timed out — partial results shown)*';
      agent.completedAt = Date.now();
      _clearCheckpoint(agent.id);
      _persistState();
      _notify(agent, 'completed');
      console.log(`[bg-agent] ${agent.id} timed out with partial results`);
    } else {
      agent.status = 'failed';
      agent.error = err.message;
      agent.result = partial || null;
      agent.completedAt = Date.now();
      // Keep checkpoint for potential manual retry
      _persistState();
      _notify(agent, 'failed');
      console.log(`[bg-agent] ${agent.id} failed: ${err.message}`);
    }
    delete agent._abort;
  } finally {
    clearTimeout(timeout);
  }
}

/**
 * Push a notification about an agent's state change to all connected clients.
 */
function _notify(agent, event) {
  if (_broadcastWS) {
    _broadcastWS({
      type: 'bg_agent',
      event,
      id: agent.id,
      name: agent.name,
      status: agent.status,
      sessionId: agent.sessionId,
      task: agent.task.slice(0, 200),
      toolCalls: agent.toolCalls,
      currentTool: agent.currentTool,
      duration: agent.completedAt
        ? agent.completedAt - (agent.startedAt || agent.createdAt)
        : agent.startedAt
          ? Date.now() - agent.startedAt
          : 0,
      result: event === 'completed' ? (agent.result || '').slice(0, 500) : undefined,
      error: event === 'failed' ? agent.error : undefined,
    });
  }

  // Forward completion/failure to Telegram if configured
  if (event === 'completed' || event === 'failed') {
    _notifyTelegram(agent, event).catch(() => {});
  }
}

async function _notifyTelegram(agent, event) {
  try {
    const { CONNECTIONS_FILE } = require('./constants');
    const conns = JSON.parse(fs.readFileSync(CONNECTIONS_FILE, 'utf-8'));
    if (!conns.telegram_key || !conns.telegram_chat_id) return;

    const telegramTool = require('./tools/telegram');
    const icon = event === 'completed' ? '✅' : '❌';
    const status = event === 'completed' ? 'completed' : 'failed';
    const preview = event === 'completed'
      ? (agent.result || '').slice(0, 300)
      : (agent.error || 'Unknown error');
    const dur = agent.completedAt && agent.startedAt
      ? `${((agent.completedAt - agent.startedAt) / 1000).toFixed(0)}s`
      : 'unknown';

    const text = `${icon} *Agent ${status}:* ${agent.name || agent.id}\n` +
      `Task: ${agent.task.slice(0, 100)}\n` +
      `Duration: ${dur} | Tools: ${agent.toolCalls}\n\n${preview}`;

    await telegramTool.telegram({ action: 'send', text, parse_mode: 'Markdown' });
  } catch {}
}

// ── Public API ───────────────────────────────────────────────────────────────

/**
 * Get agent state by ID.
 */
function getAgent(id) {
  const agent = _agents.get(id);
  if (!agent) return null;
  return _serialize(agent);
}

/**
 * List all agents (newest first).
 */
function listAgents() {
  return [..._agents.values()]
    .sort((a, b) => b.createdAt - a.createdAt)
    .map(_serialize);
}

/**
 * Cancel a running or queued agent.
 */
function cancelAgent(id) {
  const agent = _agents.get(id);
  if (!agent) return { success: false, error: 'Agent not found' };
  if (agent.status === 'completed' || agent.status === 'failed' || agent.status === 'cancelled') {
    return { success: false, error: `Agent already ${agent.status}` };
  }

  // If queued, remove from queue
  const qIdx = _queue.findIndex(q => q.agent.id === id);
  if (qIdx >= 0) {
    _queue.splice(qIdx, 1);
    agent.status = 'cancelled';
    agent.completedAt = Date.now();
    _clearCheckpoint(agent.id);
    _persistState();
    _notify(agent, 'cancelled');
    return { success: true };
  }

  // If running or paused, abort
  if (agent._abort) {
    agent.status = 'cancelled';
    agent._abort.abort();
    _clearCheckpoint(agent.id);
    _persistState();
    return { success: true };
  }

  agent.status = 'cancelled';
  agent.completedAt = Date.now();
  _clearCheckpoint(agent.id);
  _persistState();
  _notify(agent, 'cancelled');
  return { success: true };
}

/**
 * Get the full result text of a completed agent.
 */
function getResult(id) {
  const agent = _agents.get(id);
  if (!agent) return null;
  return agent.result;
}

/**
 * Serialize agent state for API responses (strip internal fields).
 */
function _serialize(agent) {
  return {
    id: agent.id,
    name: agent.name,
    task: agent.task,
    sessionId: agent.sessionId,
    status: agent.status,
    createdAt: agent.createdAt,
    startedAt: agent.startedAt,
    completedAt: agent.completedAt,
    toolCalls: agent.toolCalls,
    currentTool: agent.currentTool,
    tokens: agent.tokens,
    checkpoint: agent.checkpoint || null,
    duration: agent.completedAt
      ? agent.completedAt - (agent.startedAt || agent.createdAt)
      : agent.startedAt
        ? Date.now() - agent.startedAt
        : 0,
    resultPreview: agent.result ? agent.result.slice(0, 300) : null,
    error: agent.error,
  };
}

/**
 * Get queue status summary.
 */
function getQueueStatus() {
  const agents = [..._agents.values()];
  return {
    queued: agents.filter(a => a.status === 'queued').length,
    running: agents.filter(a => a.status === 'running').length,
    completed: agents.filter(a => a.status === 'completed').length,
    failed: agents.filter(a => a.status === 'failed').length,
    cancelled: agents.filter(a => a.status === 'cancelled').length,
    paused: agents.filter(a => a.status === 'paused').length,
    interrupted: agents.filter(a => a.status === 'interrupted').length,
    total: agents.length,
  };
}

/**
 * Initialize the agent system: load persisted state, recover interrupted agents,
 * and start sleep detection. Called once from server.js at startup.
 *
 * @returns {{ recovered: number }} count of agents re-queued for recovery
 */
function init() {
  _ensureDirs();
  const interrupted = _loadState();
  _startSleepDetector();

  let recovered = 0;
  for (const agent of interrupted) {
    const cp = _loadCheckpoint(agent.id);
    const turns = cp ? cp.turn : 0;
    console.log(`[bg-agent] Recovering interrupted agent ${agent.id}: "${agent.name}" (${turns} turns completed)`);

    // Mark as queued for re-execution with checkpoint resume
    agent.status = 'queued';
    agent._resumeFromCheckpoint = true;
    _agents.set(agent.id, agent);
    _persistState();
    _notify(agent, 'queued');

    _queue.push({ agent });
    recovered++;
  }

  if (recovered > 0) {
    console.log(`[bg-agent] ${recovered} interrupted agent(s) queued for recovery`);
    _drainQueue();
  }

  return { recovered };
}

/**
 * Shutdown: stop sleep detector and persist final state.
 */
function shutdown() {
  _stopSleepDetector();
  _persistState();
}

module.exports = {
  spawn,
  getAgent,
  listAgents,
  cancelAgent,
  getResult,
  getQueueStatus,
  setBroadcast,
  init,
  shutdown,
};
