import { describe, it, expect, beforeEach, vi } from 'vitest';
import Database from 'better-sqlite3';
import { runMigrations } from '../db/migrate.js';
import { createRegistry } from './registry.js';
import { createTokenManager } from './token-manager.js';
import { createToolCache } from './tool-cache.js';
import { createAggregator } from './aggregator.js';

describe('createAggregator', () => {
  const validKey = 'a'.repeat(64);
  let db: Database.Database;

  beforeEach(() => {
    process.env.MCP_ENCRYPTION_KEY = validKey;
    db = new Database(':memory:');
    db.exec(`CREATE TABLE users (id TEXT PRIMARY KEY);`);
    db.exec(`CREATE TABLE jobs (id TEXT PRIMARY KEY, wait_reason TEXT);`);
    db.exec(`CREATE TABLE local_tasks (id INTEGER PRIMARY KEY AUTOINCREMENT);`);
    runMigrations(db);
    db.prepare('INSERT INTO users(id) VALUES(?)').run('u1');
    const reg = createRegistry(db);
    reg.upsert({
      id: 'canva',
      name: 'Canva',
      url: 'https://mcp.canva.example/mcp',
      authKind: 'oauth',
      ownerId: null,
      oauthClientId: 'cid',
      oauthClientSecret: 'cs',
      oauthScopes: null,
    });
    const cache = createToolCache(db, 600);
    cache.replaceForServer('canva', [
      { name: 'generate_designs', description: 'g', inputSchema: { type: 'object', properties: {} } },
    ]);
  });

  it('returns ToolDefs only for servers the user is connected to', async () => {
    const tm = createTokenManager(db, { doRefresh: async () => ({ access_token: 'x' }) });
    const agg = createAggregator({
      registry: createRegistry(db),
      tokenManager: tm,
      toolCache: createToolCache(db, 600),
      executeCall: vi.fn(),
    });

    // No token → empty
    let defs = await agg.getToolDefs('u1', ['mcp__canva__*']);
    expect(defs).toEqual([]);

    tm.saveTokens({
      userId: 'u1',
      serverId: 'canva',
      accessToken: 'at',
      refreshToken: 'rt',
      expiresAt: new Date(Date.now() + 3600_000).toISOString(),
      scope: null,
    });

    defs = await agg.getToolDefs('u1', ['mcp__canva__*']);
    expect(defs).toHaveLength(1);
    expect(defs[0].function.name).toBe('mcp__canva__generate_designs');
  });

  it('delegates executeTool for mcp__ names and returns null otherwise', async () => {
    const executeCall = vi.fn().mockResolvedValue({ output: 'ok', isError: false });
    const agg = createAggregator({
      registry: createRegistry(db),
      tokenManager: createTokenManager(db, { doRefresh: async () => ({ access_token: 'x' }) }),
      toolCache: createToolCache(db, 600),
      executeCall,
    });

    const other = await agg.executeTool('Read', {}, { workspacePath: '/tmp', ownerId: 'u1', jobId: 'j' } as never);
    expect(other).toBeNull();

    const bad = await agg.executeTool('mcp__bad__name__extra', {}, { workspacePath: '/tmp', ownerId: 'u1', jobId: 'j' } as never);
    expect(bad?.isError).toBe(true);
  });

  it('getToolDefs offers ONLY the job space\'s servers (cross-space negative)', async () => {
    // Two space-scoped servers; user u1 is connected to BOTH.
    const reg = createRegistry(db);
    reg.upsert({ id: 'srva', name: 'A', url: 'https://a.example/mcp', authKind: 'api_key', ownerId: 'u1', staticToken: 'tok', spaceId: 'space-A' });
    reg.upsert({ id: 'srvb', name: 'B', url: 'https://b.example/mcp', authKind: 'api_key', ownerId: 'u1', staticToken: 'tok', spaceId: 'space-B' });
    const cache = createToolCache(db, 600);
    cache.replaceForServer('srva', [{ name: 'tool_a', description: 'a', inputSchema: { type: 'object', properties: {} } }]);
    cache.replaceForServer('srvb', [{ name: 'tool_b', description: 'b', inputSchema: { type: 'object', properties: {} } }]);
    const tm = createTokenManager(db, { doRefresh: async () => ({ access_token: 'x' }) });
    // api_key servers report hasToken via static token, no per-user token needed.

    const agg = createAggregator({ registry: reg, tokenManager: tm, toolCache: cache, executeCall: vi.fn() });

    const inA = await agg.getToolDefs('u1', ['mcp__srva__*', 'mcp__srvb__*'], 'space-A');
    expect(inA.map(d => d.function.name)).toEqual(['mcp__srva__tool_a']);
    // SECURITY: space-B server must NOT be visible to a space-A job.
    expect(inA.map(d => d.function.name)).not.toContain('mcp__srvb__tool_b');

    const inB = await agg.getToolDefs('u1', ['mcp__srva__*', 'mcp__srvb__*'], 'space-B');
    expect(inB.map(d => d.function.name)).toEqual(['mcp__srvb__tool_b']);
  });

  it('getToolDefs with undefined spaceId falls back to the legacy NULL-space set', async () => {
    // canva (set up in beforeEach) is a global NULL-space server.
    const reg = createRegistry(db);
    reg.upsert({ id: 'srva', name: 'A', url: 'https://a.example/mcp', authKind: 'api_key', ownerId: 'u1', staticToken: 'tok', spaceId: 'space-A' });
    const cache = createToolCache(db, 600);
    cache.replaceForServer('canva', [{ name: 'generate_designs', description: 'g', inputSchema: { type: 'object', properties: {} } }]);
    cache.replaceForServer('srva', [{ name: 'tool_a', description: 'a', inputSchema: { type: 'object', properties: {} } }]);
    const tm = createTokenManager(db, { doRefresh: async () => ({ access_token: 'x' }) });
    tm.saveTokens({ userId: 'u1', serverId: 'canva', accessToken: 'at', refreshToken: 'rt', expiresAt: new Date(Date.now() + 3600_000).toISOString(), scope: null });

    const agg = createAggregator({ registry: reg, tokenManager: tm, toolCache: cache, executeCall: vi.fn() });

    // No spaceId → legacy fallback → only the NULL-space (global) server.
    const defs = await agg.getToolDefs('u1', ['mcp__canva__*', 'mcp__srva__*']);
    expect(defs.map(d => d.function.name)).toEqual(['mcp__canva__generate_designs']);
    expect(defs.map(d => d.function.name)).not.toContain('mcp__srva__tool_a');
  });

  it('executeTool blocks a cross-space MCP call (defense in depth)', async () => {
    const reg = createRegistry(db);
    reg.upsert({ id: 'srvb', name: 'B', url: 'https://b.example/mcp', authKind: 'api_key', ownerId: 'u1', staticToken: 'tok', spaceId: 'space-B' });
    const executeCall = vi.fn().mockResolvedValue({ output: 'ok', isError: false });
    const agg = createAggregator({
      registry: reg,
      tokenManager: createTokenManager(db, { doRefresh: async () => ({ access_token: 'x' }) }),
      toolCache: createToolCache(db, 600),
      executeCall,
    });

    // A job in space-A tries to call a space-B server → blocked, executeCall never runs.
    const blocked = await agg.executeTool('mcp__srvb__tool_b', {}, { workspacePath: '/tmp', ownerId: 'u1', jobId: 'j', spaceId: 'space-A' } as never);
    expect(blocked?.isError).toBe(true);
    expect(executeCall).not.toHaveBeenCalled();

    // Same job in space-B is allowed through to executeCall.
    const ok = await agg.executeTool('mcp__srvb__tool_b', {}, { workspacePath: '/tmp', ownerId: 'u1', jobId: 'j', spaceId: 'space-B' } as never);
    expect(ok?.isError).toBe(false);
    expect(executeCall).toHaveBeenCalledTimes(1);
  });
});
