import { describe, it, expect, vi } from 'vitest';
import { SessionRegistry } from './console-registry.js';

type FakeViewer = { userId: string; close: ReturnType<typeof vi.fn> };

function fakeSession(
  localTaskId: string,
  connectionId: string,
  opts: { lastActivityAt?: number; startedAt?: number; viewers?: FakeViewer[] } = {},
) {
  return {
    localTaskId,
    connectionId,
    lastActivityAt: opts.lastActivityAt ?? Date.now(),
    startedAt: opts.startedAt ?? Date.now(),
    isClosed: false,
    close: vi.fn(async (_reason: string) => {}),
    listViewers: vi.fn(() => opts.viewers ?? []),
  } as any;
}

describe('SessionRegistry', () => {
  it('store / get / closeForTask', async () => {
    const r = new SessionRegistry({
      idleTimeoutMs: 60_000,
      maxSessionDurationMs: 3_600_000,
      maxSessionsPerConnection: 3,
    });
    const s = fakeSession('t1', 'c1');
    r.register(s);
    expect(r.get('t1')).toBe(s);
    await r.closeForTask('t1', 'admin_kill');
    expect(s.close).toHaveBeenCalledWith('admin_kill');
    expect(r.get('t1')).toBeNull();
  });

  it('sweep closes idle sessions', async () => {
    const r = new SessionRegistry({
      idleTimeoutMs: 1000,
      maxSessionDurationMs: 3_600_000,
      maxSessionsPerConnection: 3,
    });
    const idle = fakeSession('t1', 'c1', { lastActivityAt: Date.now() - 2000 });
    const fresh = fakeSession('t2', 'c1', { lastActivityAt: Date.now() });
    r.register(idle);
    r.register(fresh);
    await r.sweep();
    expect(idle.close).toHaveBeenCalledWith('idle_timeout');
    expect(fresh.close).not.toHaveBeenCalled();
  });

  it('sweep closes sessions over duration cap', async () => {
    const r = new SessionRegistry({
      idleTimeoutMs: 60_000,
      maxSessionDurationMs: 1000,
      maxSessionsPerConnection: 3,
    });
    const old = fakeSession('t1', 'c1', {
      startedAt: Date.now() - 2000,
      lastActivityAt: Date.now(),
    });
    r.register(old);
    await r.sweep();
    expect(old.close).toHaveBeenCalledWith('duration_cap');
  });

  it('max_sessions_per_connection evicts oldest', () => {
    const r = new SessionRegistry({
      idleTimeoutMs: 60_000,
      maxSessionDurationMs: 3_600_000,
      maxSessionsPerConnection: 2,
    });
    const a = fakeSession('t1', 'c1', { startedAt: 1 });
    const b = fakeSession('t2', 'c1', { startedAt: 2 });
    const c = fakeSession('t3', 'c1', { startedAt: 3 });
    r.register(a);
    r.register(b);
    const evicted = r.enforceCap('c1');
    expect(evicted).toEqual([]);
    r.register(c);
    const evicted2 = r.enforceCap('c1');
    expect(evicted2.map((s: any) => s.localTaskId)).toEqual(['t1']);
  });

  describe('revokeAccessFor', () => {
    function mkRegistry() {
      return new SessionRegistry({
        idleTimeoutMs: 60_000,
        maxSessionDurationMs: 3_600_000,
        maxSessionsPerConnection: 3,
      });
    }

    it('kicks viewers matching userId on the target connection', () => {
      const r = mkRegistry();
      const vAlice: FakeViewer = { userId: 'alice', close: vi.fn() };
      const vBob: FakeViewer = { userId: 'bob', close: vi.fn() };
      const session = fakeSession('t1', 'c-target', { viewers: [vAlice, vBob] });
      r.register(session);
      const kicked = r.revokeAccessFor({ connectionId: 'c-target', userId: 'alice', reason: 'access_revoked' });
      expect(kicked).toBe(1);
      expect(vAlice.close).toHaveBeenCalledWith('access_revoked');
      expect(vBob.close).not.toHaveBeenCalled();
    });

    it('does not kick viewers on other connections', () => {
      const r = mkRegistry();
      const vAlice: FakeViewer = { userId: 'alice', close: vi.fn() };
      const sessionOther = fakeSession('t2', 'c-other', { viewers: [vAlice] });
      r.register(sessionOther);
      const kicked = r.revokeAccessFor({ connectionId: 'c-target', userId: 'alice', reason: 'access_revoked' });
      expect(kicked).toBe(0);
      expect(vAlice.close).not.toHaveBeenCalled();
    });

    it('leaves the underlying session alive (does not call session.close)', () => {
      const r = mkRegistry();
      const vAlice: FakeViewer = { userId: 'alice', close: vi.fn() };
      const session = fakeSession('t1', 'c1', { viewers: [vAlice] });
      r.register(session);
      r.revokeAccessFor({ connectionId: 'c1', userId: 'alice', reason: 'access_revoked' });
      expect(session.close).not.toHaveBeenCalled();
      // Session still listed for the connection
      expect(r.listForConnection('c1')).toHaveLength(1);
    });

    it('counts multiple viewer hits across sessions on the same connection', () => {
      const r = mkRegistry();
      const vA1: FakeViewer = { userId: 'alice', close: vi.fn() };
      const vA2: FakeViewer = { userId: 'alice', close: vi.fn() };
      const vB: FakeViewer = { userId: 'bob', close: vi.fn() };
      r.register(fakeSession('t1', 'c1', { viewers: [vA1, vB] }));
      r.register(fakeSession('t2', 'c1', { viewers: [vA2] }));
      const kicked = r.revokeAccessFor({ connectionId: 'c1', userId: 'alice', reason: 'access_revoked' });
      expect(kicked).toBe(2);
      expect(vA1.close).toHaveBeenCalled();
      expect(vA2.close).toHaveBeenCalled();
      expect(vB.close).not.toHaveBeenCalled();
    });

    it('returns 0 when the connection has no active sessions', () => {
      const r = mkRegistry();
      const kicked = r.revokeAccessFor({ connectionId: 'c-missing', userId: 'alice', reason: 'access_revoked' });
      expect(kicked).toBe(0);
    });

    it('survives a viewer close() that throws (other viewers still kicked)', () => {
      const r = mkRegistry();
      const bad: FakeViewer = { userId: 'alice', close: vi.fn(() => { throw new Error('ws gone'); }) };
      const good: FakeViewer = { userId: 'alice', close: vi.fn() };
      r.register(fakeSession('t1', 'c1', { viewers: [bad, good] }));
      const kicked = r.revokeAccessFor({ connectionId: 'c1', userId: 'alice', reason: 'access_revoked' });
      expect(bad.close).toHaveBeenCalled();
      expect(good.close).toHaveBeenCalled();
      expect(kicked).toBe(1); // only the good one counted as successful kick
    });
  });

  it('shutdown closes all and clears map', async () => {
    const r = new SessionRegistry({
      idleTimeoutMs: 60_000,
      maxSessionDurationMs: 3_600_000,
      maxSessionsPerConnection: 3,
    });
    const a = fakeSession('t1', 'c1');
    const b = fakeSession('t2', 'c2');
    r.register(a);
    r.register(b);
    await r.shutdown();
    expect(a.close).toHaveBeenCalledWith('worker_shutdown');
    expect(b.close).toHaveBeenCalledWith('worker_shutdown');
    expect(r.get('t1')).toBeNull();
    expect(r.get('t2')).toBeNull();
  });
});
