import { randomUUID } from 'crypto';
import type { SessionManager } from '../session/manager.js';
import type { ModelClient } from '../models/types.js';
import type { ModelRouter, ModelTier } from '../models/router.js';
import type { ToolRegistry } from '../tools/registry.js';
import type { ToolExecutor } from '../tools/executor.js';
import { AgentOrchestrator, type DelegationConfig } from '../backends/native/orchestrator.js';
import type { ToolUseEvent } from '../backends/native/agent.js';
import type { MemoryStore } from '../memory/store.js';
import type { Config } from '../config/index.js';

export interface SessionBridgeConfig {
  sessionManager: SessionManager;
  modelClient: ModelClient | ModelRouter;
  systemPrompt: string;
  toolRegistry: ToolRegistry;
  toolExecutor: ToolExecutor;
  config?: Config;
  memoryStore?: MemoryStore;
}

interface ClientEntry {
  connectionId: string;
  sessionId: string;
  agent: AgentOrchestrator;
  busy: boolean;
}

export class SessionBridge {
  private clients: Map<string, ClientEntry> = new Map();
  private agents: Map<string, AgentOrchestrator> = new Map();
  private config: SessionBridgeConfig;

  constructor(config: SessionBridgeConfig) {
    this.config = config;
  }

  /** Register a new WS connection. Returns the assigned connection ID. */
  connect(connectionId?: string): string {
    const id = connectionId ?? randomUUID();
    const sessionId = `ws:${id}`;
    const agent = this.getOrCreateAgent(sessionId);

    this.clients.set(id, {
      connectionId: id,
      sessionId,
      agent,
      busy: false,
    });

    return id;
  }

  /** Remove a WS connection. Does NOT destroy the session (persists in SQLite). */
  disconnect(connectionId: string): void {
    const client = this.clients.get(connectionId);
    if (client) {
      // Only remove the agent if no other clients share the session
      const otherClients = Array.from(this.clients.values())
        .filter(c => c.sessionId === client.sessionId && c.connectionId !== connectionId);
      if (otherClients.length === 0) {
        this.agents.delete(client.sessionId);
      }
      this.clients.delete(connectionId);
    }
  }

  /** Switch a connection to a different session (e.g. resuming an old session). */
  switchSession(connectionId: string, sessionId: string): void {
    const client = this.clients.get(connectionId);
    if (!client) {throw new Error(`Unknown connection: ${connectionId}`);}
    if (client.busy) {throw new Error('Cannot switch session while agent is busy');}

    const agent = this.getOrCreateAgent(sessionId);
    client.sessionId = sessionId;
    client.agent = agent;
  }

  /** Get the AgentOrchestrator for a connection. */
  getAgent(connectionId: string): AgentOrchestrator | undefined {
    return this.clients.get(connectionId)?.agent;
  }

  /** Get the session ID for a connection. */
  getSessionId(connectionId: string): string | undefined {
    return this.clients.get(connectionId)?.sessionId;
  }

  /** Check if a connection's agent is busy. */
  isBusy(connectionId: string): boolean {
    return this.clients.get(connectionId)?.busy ?? false;
  }

  /** Mark a connection's agent as busy/idle. */
  setBusy(connectionId: string, busy: boolean): void {
    const client = this.clients.get(connectionId);
    if (client) {client.busy = busy;}
  }

  /** Request cancellation for the current operation on a connection's agent. */
  cancel(connectionId: string): boolean {
    const client = this.clients.get(connectionId);
    if (!client || !client.busy) {
      return false;
    }

    client.agent.cancel();
    return true;
  }

  /** Set onToolUse callback for a connection's agent. */
  setOnToolUse(connectionId: string, callback: ((event: ToolUseEvent) => void) | undefined): void {
    const client = this.clients.get(connectionId);
    if (client) {client.agent.setOnToolUse(callback);}
  }

  /** List all active sessions with connection counts. */
  listSessions(): Array<{ sessionId: string; connections: number }> {
    const sessionMap = new Map<string, number>();
    for (const client of this.clients.values()) {
      sessionMap.set(client.sessionId, (sessionMap.get(client.sessionId) ?? 0) + 1);
    }
    return Array.from(sessionMap.entries()).map(([sessionId, connections]) => ({
      sessionId,
      connections,
    }));
  }

  /** Get count of active connections. */
  get connectionCount(): number {
    return this.clients.size;
  }

  /** Get usage stats for a specific connection's agent. */
  getUsage(connectionId: string): { inputTokens: number; outputTokens: number; calls: number } | undefined {
    const agent = this.clients.get(connectionId)?.agent;
    if (!agent) {return undefined;}
    const usage = agent.getUsage();
    return {
      inputTokens: usage.primary.inputTokens,
      outputTokens: usage.primary.outputTokens,
      calls: usage.primary.calls,
    };
  }

  /** Get usage stats for all active sessions. Returns an array of per-session usage entries. */
  getAllUsage(): Array<{
    sessionId: string;
    primary: { inputTokens: number; outputTokens: number; calls: number };
    delegation: Record<string, { inputTokens: number; outputTokens: number; calls: number }>;
    total: { inputTokens: number; outputTokens: number; calls: number; estimatedCost: number };
  }> {
    const results: Array<{
      sessionId: string;
      primary: { inputTokens: number; outputTokens: number; calls: number };
      delegation: Record<string, { inputTokens: number; outputTokens: number; calls: number }>;
      total: { inputTokens: number; outputTokens: number; calls: number; estimatedCost: number };
    }> = [];

    // De-duplicate by sessionId (multiple connections may share a session)
    const seen = new Set<string>();

    for (const client of this.clients.values()) {
      if (seen.has(client.sessionId)) {continue;}
      seen.add(client.sessionId);

      const usage = client.agent.getUsage();
      results.push({
        sessionId: client.sessionId,
        primary: usage.primary,
        delegation: usage.delegation,
        total: usage.total,
      });
    }

    return results;
  }

  private getOrCreateAgent(sessionId: string): AgentOrchestrator {
    let agent = this.agents.get(sessionId);
    if (!agent) {
      const session = this.config.sessionManager.getSession('ws', sessionId);
      const config = this.config.config;

      // Read per-session tier override from session config
      const sessionTier = session.getConfig?.('modelTier') as ModelTier | undefined;
      const primaryTier = sessionTier ?? config?.agents.primary_tier ?? 'default';

      const delegationConfig: DelegationConfig = {
        compaction: config?.agents.delegation.compaction ?? 'fast',
        memory_extraction: config?.agents.delegation.memory_extraction ?? 'fast',
        classification: config?.agents.delegation.classification ?? 'fast',
        tool_summarisation: config?.agents.delegation.tool_summarisation ?? 'fast',
        complex_reasoning: config?.agents.delegation.complex_reasoning ?? 'complex',
      };

      agent = new AgentOrchestrator({
        modelRouter: this.config.modelClient as ModelRouter,
        systemPrompt: this.config.systemPrompt,
        session,
        toolRegistry: this.config.toolRegistry,
        toolExecutor: this.config.toolExecutor,
        primaryTier,
        delegation: delegationConfig,
        maxDelegationDepth: config?.agents.max_delegation_depth ?? 3,
        maxIterations: config?.agents.max_iterations,
        compaction: config?.compaction.enabled ? {
          thresholdPct: config.compaction.threshold_pct,
          keepTurns: config.compaction.keep_turns,
          summaryMaxTokens: config.compaction.summary_max_tokens,
          importanceThreshold: config.compaction.importance_threshold,
        } : undefined,
        modelName: config?.models.default.model,
        contextWindow: config?.models.default.context_window,
        memoryStore: this.config.memoryStore,
        memoryInjectionStrategy: config?.memory?.injection_strategy,
        memoryMaxInjectionTokens: config?.memory?.max_injection_tokens,
        toolPolicyContext: {
          agent: primaryTier,
          provider: config?.models.default.provider,
          autonomyLevel: config?.agents.autonomy_level ?? 'standard',
        },
      });

      this.agents.set(sessionId, agent);
    }
    return agent;
  }
}
