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.
 *
 * Responsibilities:
 *   - register / lookup / close-by-task-id
 *   - periodic sweep for idle_timeout + duration_cap
 *   - 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, ConsoleSession>();
  private sweepTimer: ReturnType<typeof setInterval> | null = null;

  constructor(private readonly opts: SessionRegistryOptions) {}

  register(session: ConsoleSession): void {
    this.byTask.set(session.localTaskId, session);
  }

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

  listAll(): ConsoleSession[] {
    return [...this.byTask.values()].filter((s) => !s.isClosed);
  }

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

  async closeForTask(localTaskId: string, reason: SessionCloseReason): Promise<void> {
    const s = this.byTask.get(localTaskId);
    if (!s) return;
    this.byTask.delete(localTaskId);
    try {
      await 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 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']);
      }
    }
    await Promise.all(toClose.map(([s, r]) => this.closeForTask(s.localTaskId, 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 = [...this.byTask.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}`),
          ),
      ),
    );
  }
}
