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;
}

function mkRegistry(overrides: Partial<{
  idleTimeoutMs: number;
  maxSessionDurationMs: number;
  maxSessionsPerConnection: number;
}> = {}) {
  return new SessionRegistry({
    idleTimeoutMs: 60_000,
    maxSessionDurationMs: 3_600_000,
    maxSessionsPerConnection: 3,
    ...overrides,
  });
}

describe('SessionRegistry (multi-session, keyed by task+connection)', () => {
  describe('register / get / listForTask', () => {
    it('keeps two sessions for the same task on different connections', () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'connA');
      const b = fakeSession('t1', 'connB');
      r.register(a);
      r.register(b);

      expect(r.listForTask('t1')).toHaveLength(2);
      expect(r.get('t1', 'connA')).toBe(a);
      expect(r.get('t1', 'connB')).toBe(b);
    });

    it('returns null for an unknown task or unknown connection on a known task', () => {
      const r = mkRegistry();
      r.register(fakeSession('t1', 'connA'));

      expect(r.get('t-missing', 'connA')).toBeNull();
      expect(r.get('t1', 'conn-missing')).toBeNull();
    });

    it('does not return a closed session from get() or listForTask()', () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'connA');
      r.register(a);
      a.isClosed = true;

      expect(r.get('t1', 'connA')).toBeNull();
      expect(r.listForTask('t1')).toHaveLength(0);
    });

    it('registering a new session on an existing (task, connection) pair replaces it', () => {
      const r = mkRegistry();
      const a1 = fakeSession('t1', 'connA');
      const a2 = fakeSession('t1', 'connA');
      r.register(a1);
      r.register(a2);

      expect(r.get('t1', 'connA')).toBe(a2);
      expect(r.listForTask('t1')).toHaveLength(1);
    });
  });

  describe('getMostRecentForTask', () => {
    it('picks the session with the highest lastActivityAt', () => {
      const r = mkRegistry();
      const older = fakeSession('t1', 'connA', { lastActivityAt: 1000 });
      const newer = fakeSession('t1', 'connB', { lastActivityAt: 5000 });
      r.register(older);
      r.register(newer);

      expect(r.getMostRecentForTask('t1')).toBe(newer);
    });

    it('ignores closed sessions', () => {
      const r = mkRegistry();
      const newerButClosed = fakeSession('t1', 'connA', { lastActivityAt: 9000 });
      const older = fakeSession('t1', 'connB', { lastActivityAt: 1000 });
      r.register(newerButClosed);
      r.register(older);
      newerButClosed.isClosed = true;

      expect(r.getMostRecentForTask('t1')).toBe(older);
    });

    it('returns null when the task has no sessions', () => {
      const r = mkRegistry();
      expect(r.getMostRecentForTask('t-missing')).toBeNull();
    });
  });

  describe('closeSession', () => {
    it('closes only the targeted connection, leaving siblings on the same task alive', async () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'connA');
      const b = fakeSession('t1', 'connB');
      r.register(a);
      r.register(b);

      await r.closeSession('t1', 'connA', 'admin_kill');

      expect(a.close).toHaveBeenCalledWith('admin_kill');
      expect(b.close).not.toHaveBeenCalled();
      expect(r.get('t1', 'connA')).toBeNull();
      expect(r.get('t1', 'connB')).toBe(b);
      expect(r.listForTask('t1')).toEqual([b]);
    });

    it('is a no-op when the (task, connection) pair does not exist', async () => {
      const r = mkRegistry();
      await expect(r.closeSession('t-missing', 'conn-missing', 'admin_kill')).resolves.toBeUndefined();
    });

    it('removes the task entry entirely once its last session is closed', async () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'connA');
      r.register(a);
      await r.closeSession('t1', 'connA', 'admin_kill');
      expect(r.listForTask('t1')).toEqual([]);
      expect(r.listAll()).toEqual([]);
    });
  });

  describe('closeForTask', () => {
    it('closes ALL sessions for the task', async () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'connA');
      const b = fakeSession('t1', 'connB');
      r.register(a);
      r.register(b);

      await r.closeForTask('t1', 'admin_kill');

      expect(a.close).toHaveBeenCalledWith('admin_kill');
      expect(b.close).toHaveBeenCalledWith('admin_kill');
      expect(r.get('t1', 'connA')).toBeNull();
      expect(r.get('t1', 'connB')).toBeNull();
      expect(r.listForTask('t1')).toEqual([]);
    });

    it('does not touch other tasks', async () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'connA');
      const other = fakeSession('t2', 'connA');
      r.register(a);
      r.register(other);

      await r.closeForTask('t1', 'admin_kill');

      expect(other.close).not.toHaveBeenCalled();
      expect(r.get('t2', 'connA')).toBe(other);
    });

    it('is a no-op when the task has no sessions', async () => {
      const r = mkRegistry();
      await expect(r.closeForTask('t-missing', 'admin_kill')).resolves.toBeUndefined();
    });
  });

  describe('listAll / listForConnection', () => {
    it('listAll flattens sessions across all tasks, excluding closed ones', () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'connA');
      const b = fakeSession('t2', 'connB');
      const c = fakeSession('t2', 'connC');
      r.register(a);
      r.register(b);
      r.register(c);
      c.isClosed = true;

      const all = r.listAll();
      expect(all).toHaveLength(2);
      expect(all).toEqual(expect.arrayContaining([a, b]));
    });

    it('listForConnection filters listAll by connectionId across tasks', () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'c1');
      const b = fakeSession('t2', 'c1');
      const other = fakeSession('t3', 'c2');
      r.register(a);
      r.register(b);
      r.register(other);

      const forC1 = r.listForConnection('c1');
      expect(forC1).toHaveLength(2);
      expect(forC1).toEqual(expect.arrayContaining([a, b]));
    });
  });

  describe('sweep', () => {
    it('closes an idle session without touching a fresh session on the SAME task', async () => {
      const r = mkRegistry({ idleTimeoutMs: 1000 });
      const idle = fakeSession('t1', 'connA', { lastActivityAt: Date.now() - 2000 });
      const fresh = fakeSession('t1', 'connB', { lastActivityAt: Date.now() });
      r.register(idle);
      r.register(fresh);

      await r.sweep();

      expect(idle.close).toHaveBeenCalledWith('idle_timeout');
      expect(fresh.close).not.toHaveBeenCalled();
      expect(r.get('t1', 'connA')).toBeNull();
      expect(r.get('t1', 'connB')).toBe(fresh);
    });

    it('closes sessions over the duration cap', async () => {
      const r = mkRegistry({ idleTimeoutMs: 60_000, maxSessionDurationMs: 1000 });
      const old = fakeSession('t1', 'connA', {
        startedAt: Date.now() - 2000,
        lastActivityAt: Date.now(),
      });
      r.register(old);

      await r.sweep();

      expect(old.close).toHaveBeenCalledWith('duration_cap');
    });

    it('closes idle sessions across different tasks independently', async () => {
      const r = mkRegistry({ idleTimeoutMs: 1000 });
      const idle = fakeSession('t1', 'connA', { lastActivityAt: Date.now() - 2000 });
      const fresh = fakeSession('t2', 'connA', { lastActivityAt: Date.now() });
      r.register(idle);
      r.register(fresh);

      await r.sweep();

      expect(idle.close).toHaveBeenCalledWith('idle_timeout');
      expect(fresh.close).not.toHaveBeenCalled();
    });
  });

  describe('enforceCap', () => {
    it('returns the oldest sessions over the per-connection cap (across tasks)', () => {
      const r = mkRegistry({ 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);
      expect(r.enforceCap('c1')).toEqual([]);
      r.register(c);
      const evicted2 = r.enforceCap('c1');
      expect(evicted2.map((s: any) => s.localTaskId)).toEqual(['t1']);
    });
  });

  describe('revokeAccessFor', () => {
    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();
      expect(r.listForConnection('c1')).toHaveLength(1);
    });

    it('counts multiple viewer hits across sessions (different tasks) 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
    });
  });

  describe('shutdown', () => {
    it('closes every session across every task and clears the registry', async () => {
      const r = mkRegistry();
      const a = fakeSession('t1', 'connA');
      const b = fakeSession('t1', 'connB');
      const c = fakeSession('t2', 'connA');
      r.register(a);
      r.register(b);
      r.register(c);

      await r.shutdown();

      expect(a.close).toHaveBeenCalledWith('worker_shutdown');
      expect(b.close).toHaveBeenCalledWith('worker_shutdown');
      expect(c.close).toHaveBeenCalledWith('worker_shutdown');
      expect(r.get('t1', 'connA')).toBeNull();
      expect(r.get('t1', 'connB')).toBeNull();
      expect(r.get('t2', 'connA')).toBeNull();
      expect(r.listAll()).toEqual([]);
    });
  });
});
