import type { ConsoleSession } from './console-session.js';
import type { SessionCloseReason } from './console-protocol.js';
import { logger } from '../logger.js';

export interface SessionRegistryOptions {
  idleTimeoutMs: number;
  maxSessionDurationMs: number;
  maxSessionsPerConnection: number;
}

/**
 * In-memory registry of live ConsoleSessions, keyed by (localTaskId, connectionId).
 *
 * A task can hold multiple concurrent sessions — one per connection — so a
 * user can have several SSH consoles open on the same task at once. Lookups
 * that used to be per-task now come in two flavors: exact `(task, connection)`
 * lookup via `get`, and task-scoped listing via `listForTask`.
 *
 * Responsibilities:
 *   - register / lookup / close-by-(task,connection) / close-all-for-task
 *   - periodic sweep for idle_timeout + duration_cap (per session, not per task)
 *   - enforce per-connection session caps (returns evict-candidates;
 *     caller decides whether to close them, since we want a clear audit
 *     reason like 'session_cap_evict' from the caller's context).
 *   - graceful shutdown on worker stop
 */
export class SessionRegistry {
  private byTask = new Map<string, Map<string, ConsoleSession>>();
  private sweepTimer: ReturnType<typeof setInterval> | null = null;

  constructor(private readonly opts: SessionRegistryOptions) {}

  register(session: ConsoleSession): void {
    let inner = this.byTask.get(session.localTaskId);
    if (!inner) {
      inner = new Map();
      this.byTask.set(session.localTaskId, inner);
    }
    inner.set(session.connectionId, session);
  }

  get(localTaskId: string, connectionId: string): ConsoleSession | null {
    const s = this.byTask.get(localTaskId)?.get(connectionId);
    return s && !s.isClosed ? s : null;
  }

  /** Non-closed sessions for a task, one per connection. */
  listForTask(localTaskId: string): ConsoleSession[] {
    return [...(this.byTask.get(localTaskId)?.values() ?? [])].filter((s) => !s.isClosed);
  }

  /** The most recently active non-closed session for a task, or null if none. */
  getMostRecentForTask(localTaskId: string): ConsoleSession | null {
    return (
      this.listForTask(localTaskId).sort((a, b) => b.lastActivityAt - a.lastActivityAt)[0] ?? null
    );
  }

  listAll(): ConsoleSession[] {
    const all: ConsoleSession[] = [];
    for (const inner of this.byTask.values()) {
      for (const s of inner.values()) {
        if (!s.isClosed) all.push(s);
      }
    }
    return all;
  }

  listForConnection(connectionId: string): ConsoleSession[] {
    return this.listAll().filter((s) => s.connectionId === connectionId);
  }

  /** Close a single (task, connection) session, leaving any other sessions on the same task alive. */
  async closeSession(
    localTaskId: string,
    connectionId: string,
    reason: SessionCloseReason,
  ): Promise<void> {
    const inner = this.byTask.get(localTaskId);
    const s = inner?.get(connectionId);
    if (!inner || !s) return;
    inner.delete(connectionId);
    if (inner.size === 0) this.byTask.delete(localTaskId);
    try {
      await s.close(reason);
    } catch (e) {
      logger.warn(`[console-registry] close error: ${(e as Error).message}`);
    }
  }

  /** Close ALL sessions registered for a task (every connection). */
  async closeForTask(localTaskId: string, reason: SessionCloseReason): Promise<void> {
    const inner = this.byTask.get(localTaskId);
    if (!inner) return;
    const sessions = [...inner.values()];
    this.byTask.delete(localTaskId);
    await Promise.all(
      sessions.map((s) =>
        s
          .close(reason)
          .catch((e) => logger.warn(`[console-registry] close error: ${(e as Error).message}`)),
      ),
    );
  }

  /**
   * Kick all active WebSocket viewers on `connectionId` that belong to
   * `userId`. The underlying SSH session is left alive so the agent and
   * other valid viewers (the connection owner, admins with bypass, viewers
   * with surviving grants) keep working. Returns the number of viewers kicked.
   *
   * Called by the grant-revocation hook in ssh-api.ts after a user-subject
   * grant is deleted. For org-subject grants, the caller would need to
   * expand to member userIds (deferred to a follow-up).
   */
  revokeAccessFor(args: {
    connectionId: string;
    userId: string;
    reason: SessionCloseReason;
  }): number {
    let kicked = 0;
    for (const session of this.listForConnection(args.connectionId)) {
      for (const v of session.listViewers()) {
        if (v.userId !== args.userId) continue;
        try {
          v.close(args.reason);
          kicked++;
        } catch (e) {
          logger.warn(`[console-registry] viewer close error: ${(e as Error).message}`);
        }
      }
    }
    if (kicked > 0) {
      logger.info(
        `[console-registry] revoked viewers connection=${args.connectionId} user=${args.userId} reason=${args.reason} kicked=${kicked}`,
      );
    }
    return kicked;
  }

  /**
   * Return the oldest sessions (across all tasks) that should be evicted to
   * keep the connection at or below its session cap. Does not mutate state —
   * the caller is responsible for closing the returned sessions (typically
   * with reason 'session_cap_evict').
   */
  enforceCap(connectionId: string): ConsoleSession[] {
    const sessions = this.listForConnection(connectionId);
    const over = sessions.length - this.opts.maxSessionsPerConnection;
    if (over <= 0) return [];
    sessions.sort((a, b) => a.startedAt - b.startedAt);
    return sessions.slice(0, over);
  }

  async sweep(): Promise<void> {
    const now = Date.now();
    const toClose: Array<[ConsoleSession, SessionCloseReason]> = [];
    for (const s of this.listAll()) {
      if (now - s.lastActivityAt > this.opts.idleTimeoutMs) {
        toClose.push([s, 'idle_timeout']);
        continue;
      }
      if (now - s.startedAt > this.opts.maxSessionDurationMs) {
        toClose.push([s, 'duration_cap']);
      }
    }
    // Close by (task, connection), NOT closeForTask — a task can have other,
    // still-healthy sessions on different connections that must stay alive.
    await Promise.all(
      toClose.map(([s, r]) => this.closeSession(s.localTaskId, s.connectionId, r)),
    );
  }

  startSweepTimer(intervalMs = 60_000): void {
    if (this.sweepTimer) return;
    this.sweepTimer = setInterval(() => {
      void this.sweep();
    }, intervalMs);
    if (typeof this.sweepTimer.unref === 'function') this.sweepTimer.unref();
  }

  stopSweepTimer(): void {
    if (this.sweepTimer) {
      clearInterval(this.sweepTimer);
      this.sweepTimer = null;
    }
  }

  async shutdown(): Promise<void> {
    this.stopSweepTimer();
    const all: ConsoleSession[] = [];
    for (const inner of this.byTask.values()) {
      all.push(...inner.values());
    }
    this.byTask.clear();
    await Promise.all(
      all.map((s) =>
        s
          .close('worker_shutdown')
          .catch((e) =>
            logger.warn(`[console-registry] shutdown close error: ${(e as Error).message}`),
          ),
      ),
    );
  }
}
