import { describe, expect, it, vi } from 'vitest';
import type { z } from 'zod';
import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import type { Brain } from '@mosaicstack/brain';
import type { Memory } from '@mosaicstack/memory';
import type { EmbeddingService } from '../memory/embedding.service.js';
import type { CoordService } from '../coord/coord.service.js';
import {
  assertMcpToolAuthorized,
  createMcpActorContext,
  deriveMcpToolScopesForUser,
  MCP_TOOL_SCOPES,
  McpService,
  type McpToolName,
} from './mcp.service.js';

type ToolResult = { content: Array<{ type: 'text'; text: string }> };
type ToolHandler = (params: Record<string, unknown>) => Promise<ToolResult>;

interface CapturedTool {
  inputSchema: z.ZodType;
  handler: ToolHandler;
}

function makeCapturingServer(): { server: McpServer; tools: Map<string, CapturedTool> } {
  const tools = new Map<string, CapturedTool>();
  const server = {
    registerTool(name: string, config: { inputSchema: z.ZodType }, handler: ToolHandler): void {
      tools.set(name, { inputSchema: config.inputSchema, handler });
    },
  };
  return { server: server as unknown as McpServer, tools };
}

function makeService(opts?: {
  projects?: Array<Record<string, unknown> & { id: string; ownerId?: string | null }>;
  missions?: Array<
    Record<string, unknown> & { id: string; projectId?: string | null; userId?: string | null }
  >;
  tasks?: Array<
    Record<string, unknown> & {
      id: string;
      projectId?: string | null;
      missionId?: string | null;
      status?: string;
    }
  >;
}) {
  const projects = opts?.projects ?? [];
  const missions = opts?.missions ?? [];
  const tasks = opts?.tasks ?? [];
  const brain = {
    projects: {
      findAll: vi.fn(async () => projects),
      findById: vi.fn(async (id: string) => projects.find((project) => project.id === id) ?? null),
    },
    tasks: {
      findAll: vi.fn(async () => tasks),
      findById: vi.fn(async (id: string) => tasks.find((task) => task.id === id) ?? null),
      findByProject: vi.fn(async (projectId: string) =>
        tasks.filter((task) => task.projectId === projectId),
      ),
      findByMission: vi.fn(async (missionId: string) =>
        tasks.filter((task) => task.missionId === missionId),
      ),
      findByStatus: vi.fn(async (status: string) => tasks.filter((task) => task.status === status)),
      create: vi.fn(async (task: Record<string, unknown>) => ({ id: 'task-1', ...task })),
      update: vi.fn(async (id: string, updates: Record<string, unknown>) => ({ id, ...updates })),
    },
    missions: {
      findAll: vi.fn(async () => missions),
      findById: vi.fn(async (id: string) => missions.find((mission) => mission.id === id) ?? null),
      findByProject: vi.fn(async (projectId: string) =>
        missions.filter((mission) => mission.projectId === projectId),
      ),
    },
    conversations: {
      findAll: vi.fn(async (userId: string) => [{ id: 'conversation-1', userId }]),
    },
  } as unknown as Brain;

  const memory = {
    insights: {
      searchByEmbedding: vi.fn(async (userId: string) => [{ id: 'insight-1', userId }]),
      create: vi.fn(async (insight: Record<string, unknown>) => ({ id: 'insight-2', ...insight })),
    },
    preferences: {
      findByUser: vi.fn(async (userId: string) => [{ key: 'theme', userId }]),
      findByUserAndCategory: vi.fn(async (userId: string, category: string) => [
        { key: 'theme', userId, category },
      ]),
      upsert: vi.fn(async (preference: Record<string, unknown>) => ({
        id: 'pref-1',
        ...preference,
      })),
    },
  } as unknown as Memory;

  const embeddings = {
    available: true,
    embed: vi.fn(async () => [0.1, 0.2, 0.3]),
  } as unknown as EmbeddingService;

  const coord = {
    getMissionStatus: vi.fn(async () => null),
    listTasks: vi.fn(async () => []),
    getTaskStatus: vi.fn(async () => null),
  } as unknown as CoordService;

  return {
    service: new McpService(brain, memory, embeddings, coord),
    brain: brain as unknown as {
      conversations: { findAll: ReturnType<typeof vi.fn> };
      tasks: {
        create: ReturnType<typeof vi.fn>;
        update: ReturnType<typeof vi.fn>;
      };
    },
    memory: memory as unknown as {
      insights: { searchByEmbedding: ReturnType<typeof vi.fn> };
    },
    coord: coord as unknown as {
      listTasks: ReturnType<typeof vi.fn>;
    },
  };
}

function getTool(tools: Map<string, CapturedTool>, name: McpToolName): CapturedTool {
  const tool = tools.get(name);
  if (!tool) throw new Error(`Missing captured tool ${name}`);
  return tool;
}

function makeMemberActor(userId = 'authenticated-user') {
  return createMcpActorContext({
    userId,
    role: 'member',
    scopes: deriveMcpToolScopesForUser({ role: 'member' }),
  });
}

function makeAdminActor(userId = 'admin-user', tenantId?: string) {
  return createMcpActorContext({
    userId,
    tenantId,
    role: 'admin',
    scopes: deriveMcpToolScopesForUser({ role: 'admin' }),
  });
}

function makePlatformAdminActor(userId = 'platform-admin-user') {
  return createMcpActorContext({
    userId,
    role: 'platform-admin',
    scopes: deriveMcpToolScopesForUser({ role: 'platform-admin' }),
  });
}

describe('MCP actor identity and tool scope enforcement', () => {
  it('derives immutable actor, tenant, channel, correlation, and explicit tool scopes server-side', () => {
    const actor = createMcpActorContext({
      userId: ' user-authenticated ',
      role: 'member',
      scopes: deriveMcpToolScopesForUser({ role: 'member' }),
    });

    expect(actor.userId).toBe('user-authenticated');
    expect(actor.tenantId).toBe('user:user-authenticated');
    expect(actor.role).toBe('member');
    expect(actor.channel).toBe('mcp');
    expect(actor.correlationId).toMatch(/[0-9a-f-]{36}/i);
    expect(actor.scopes.has(MCP_TOOL_SCOPES.memory_search)).toBe(true);
    expect(actor.scopes.has(MCP_TOOL_SCOPES.coord_list_tasks)).toBe(false);
    expect(
      deriveMcpToolScopesForUser({ role: 'admin' }).has(MCP_TOOL_SCOPES.coord_list_tasks),
    ).toBe(false);
    expect(
      deriveMcpToolScopesForUser({ role: 'platform-admin' }).has(MCP_TOOL_SCOPES.coord_list_tasks),
    ).toBe(true);
  });

  it('fails closed when scopes are not supplied by the authenticated context policy', () => {
    const actor = createMcpActorContext({ userId: 'user-authenticated' });

    expect(actor.scopes.size).toBe(0);
    expect(() => assertMcpToolAuthorized(actor, 'memory_search', { query: 'notes' })).toThrow(
      'MCP tool scope denied: memory:insight:read',
    );
  });

  it('fails closed when a tool caller supplies actor or tenant identity fields', () => {
    const actor = createMcpActorContext({ userId: 'user-authenticated' });

    expect(() =>
      assertMcpToolAuthorized(actor, 'memory_search', {
        userId: 'victim-user',
        query: 'private data',
      }),
    ).toThrow('MCP caller-controlled identity field is forbidden: userId');

    expect(() =>
      assertMcpToolAuthorized(actor, 'coord_list_tasks', {
        tenantId: 'victim-tenant',
        projectPath: '/tmp/project',
      }),
    ).toThrow('MCP caller-controlled identity field is forbidden: tenantId');

    expect(() =>
      assertMcpToolAuthorized(actor, 'brain_create_task', {
        title: 'forged org',
        organizationId: 'victim-org',
      }),
    ).toThrow('MCP caller-controlled identity field is forbidden: organizationId');

    expect(() =>
      assertMcpToolAuthorized(actor, 'brain_update_task', {
        title: 'forged team',
        teamId: 'victim-team',
      }),
    ).toThrow('MCP caller-controlled identity field is forbidden: teamId');
  });

  it('fails closed when the server-derived actor lacks the required per-tool scope', () => {
    const actor = createMcpActorContext({ userId: 'user-authenticated', scopes: [] });

    expect(() => assertMcpToolAuthorized(actor, 'memory_search', { query: 'notes' })).toThrow(
      'MCP tool scope denied: memory:insight:read',
    );
  });

  it('removes caller-controlled userId from memory schemas and never queries victim memory', async () => {
    const { service, memory } = makeService();
    const { server, tools } = makeCapturingServer();
    const actor = makeMemberActor('authenticated-user');

    service.registerTools(server, actor);
    const tool = getTool(tools, 'memory_search');

    expect(tool.inputSchema.safeParse({ userId: 'victim-user', query: 'anything' }).success).toBe(
      false,
    );
    await expect(tool.handler({ userId: 'victim-user', query: 'anything' })).rejects.toThrow(
      'MCP caller-controlled identity field is forbidden: userId',
    );
    expect(memory.insights.searchByEmbedding).not.toHaveBeenCalled();

    await tool.handler({ query: 'only my notes' });
    expect(memory.insights.searchByEmbedding).toHaveBeenCalledWith(
      'authenticated-user',
      [0.1, 0.2, 0.3],
      5,
    );
  });

  it('binds conversation listing to the authenticated actor instead of a caller-supplied userId', async () => {
    const { service, brain } = makeService();
    const { server, tools } = makeCapturingServer();
    const actor = makeMemberActor('authenticated-user');

    service.registerTools(server, actor);
    const tool = getTool(tools, 'brain_list_conversations');

    expect(tool.inputSchema.safeParse({ userId: 'victim-user' }).success).toBe(false);
    await expect(tool.handler({ userId: 'victim-user' })).rejects.toThrow(
      'MCP caller-controlled identity field is forbidden: userId',
    );
    expect(brain.conversations.findAll).not.toHaveBeenCalled();

    await tool.handler({});
    expect(brain.conversations.findAll).toHaveBeenCalledWith('authenticated-user');
  });

  it('scopes brain project, mission, and task reads to the authenticated actor', async () => {
    const { service } = makeService({
      projects: [
        { id: 'project-owned', ownerId: 'authenticated-user', name: 'owned' },
        { id: 'project-victim', ownerId: 'victim-user', name: 'victim' },
      ],
      missions: [
        { id: 'mission-owned', projectId: 'project-owned' },
        { id: 'mission-victim', userId: 'victim-user', projectId: 'project-victim' },
      ],
      tasks: [
        { id: 'task-owned-project', projectId: 'project-owned', status: 'not-started' },
        { id: 'task-owned-mission', missionId: 'mission-owned', status: 'not-started' },
        { id: 'task-victim-project', projectId: 'project-victim', status: 'not-started' },
        { id: 'task-unowned', status: 'not-started' },
      ],
    });
    const { server, tools } = makeCapturingServer();
    const actor = makeMemberActor('authenticated-user');

    service.registerTools(server, actor);

    const projects = JSON.parse(
      (await getTool(tools, 'brain_list_projects').handler({})).content[0]!.text,
    );
    expect(projects.map((project: { id: string }) => project.id)).toEqual(['project-owned']);

    const missions = JSON.parse(
      (await getTool(tools, 'brain_list_missions').handler({})).content[0]!.text,
    );
    expect(missions.map((mission: { id: string }) => mission.id)).toEqual(['mission-owned']);

    const tasks = JSON.parse(
      (await getTool(tools, 'brain_list_tasks').handler({})).content[0]!.text,
    );
    expect(tasks.map((task: { id: string }) => task.id)).toEqual([
      'task-owned-project',
      'task-owned-mission',
    ]);
  });

  it('enforces tenant boundaries for tenant-admin brain project, mission, and task reads', async () => {
    const { service } = makeService({
      projects: [
        {
          id: 'project-tenant-a',
          ownerId: 'other-user-a',
          teamId: 'tenant-a',
          name: 'same tenant',
        },
        {
          id: 'project-tenant-b',
          ownerId: 'other-user-b',
          teamId: 'tenant-b',
          name: 'other tenant',
        },
      ],
      missions: [
        { id: 'mission-tenant-a', tenantId: 'tenant-a', projectId: 'project-tenant-a' },
        { id: 'mission-tenant-b', tenantId: 'tenant-b', projectId: 'project-tenant-b' },
      ],
      tasks: [
        { id: 'task-tenant-a', projectId: 'project-tenant-a', status: 'not-started' },
        { id: 'task-tenant-b', projectId: 'project-tenant-b', status: 'not-started' },
      ],
    });
    const { server, tools } = makeCapturingServer();
    const actor = makeAdminActor('tenant-admin-user', 'tenant-a');

    service.registerTools(server, actor);

    const projects = JSON.parse(
      (await getTool(tools, 'brain_list_projects').handler({})).content[0]!.text,
    );
    expect(projects.map((project: { id: string }) => project.id)).toEqual(['project-tenant-a']);

    const missions = JSON.parse(
      (await getTool(tools, 'brain_list_missions').handler({})).content[0]!.text,
    );
    expect(missions.map((mission: { id: string }) => mission.id)).toEqual(['mission-tenant-a']);

    const tasks = JSON.parse(
      (await getTool(tools, 'brain_list_tasks').handler({})).content[0]!.text,
    );
    expect(tasks.map((task: { id: string }) => task.id)).toEqual(['task-tenant-a']);
  });

  it('denies tenant-admin task writes outside the authenticated tenant', async () => {
    const { service, brain } = makeService({
      projects: [
        { id: 'project-tenant-a', ownerId: 'other-user-a', teamId: 'tenant-a' },
        { id: 'project-tenant-b', ownerId: 'other-user-b', teamId: 'tenant-b' },
      ],
      missions: [
        { id: 'mission-tenant-a', tenantId: 'tenant-a', projectId: 'project-tenant-a' },
        { id: 'mission-tenant-b', tenantId: 'tenant-b', projectId: 'project-tenant-b' },
      ],
      tasks: [
        { id: 'task-tenant-a', projectId: 'project-tenant-a', status: 'not-started' },
        { id: 'task-tenant-b', projectId: 'project-tenant-b', status: 'not-started' },
      ],
    });
    const { server, tools } = makeCapturingServer();
    const actor = makeAdminActor('tenant-admin-user', 'tenant-a');

    service.registerTools(server, actor);

    await expect(
      getTool(tools, 'brain_create_task').handler({
        title: 'unscoped tenant write',
      }),
    ).rejects.toThrow('MCP task scope denied');
    expect(brain.tasks.create).not.toHaveBeenCalled();

    await expect(
      getTool(tools, 'brain_create_task').handler({
        title: 'cross-tenant write',
        projectId: 'project-tenant-b',
      }),
    ).rejects.toThrow('MCP task project scope denied');
    expect(brain.tasks.create).not.toHaveBeenCalled();

    await expect(
      getTool(tools, 'brain_update_task').handler({
        id: 'task-tenant-a',
        projectId: 'project-tenant-b',
      }),
    ).rejects.toThrow('MCP task project scope denied');
    expect(brain.tasks.update).not.toHaveBeenCalled();

    const updateResult = await getTool(tools, 'brain_update_task').handler({
      id: 'task-tenant-b',
      title: 'cross-tenant update',
    });
    expect(updateResult.content[0]!.text).toBe('Task not found: task-tenant-b');
    expect(brain.tasks.update).not.toHaveBeenCalled();

    await getTool(tools, 'brain_create_task').handler({
      title: 'same-tenant write',
      projectId: 'project-tenant-a',
    });
    expect(brain.tasks.create).toHaveBeenCalledWith(
      expect.objectContaining({ projectId: 'project-tenant-a', title: 'same-tenant write' }),
    );
  });

  it('keeps admin-only coordination tools on server-derived paths', async () => {
    const { service, coord } = makeService();
    const { server, tools } = makeCapturingServer();
    const member = makeMemberActor('authenticated-user');
    const tenantAdmin = makeAdminActor('admin-user');
    const platformAdmin = makePlatformAdminActor('platform-admin-user');

    service.registerTools(server, member);
    const memberTool = getTool(tools, 'coord_list_tasks');
    expect(memberTool.inputSchema.safeParse({ projectPath: '/tmp/victim' }).success).toBe(false);
    await expect(memberTool.handler({})).rejects.toThrow('MCP tool scope denied: coord:read');

    tools.clear();
    service.registerTools(server, tenantAdmin);
    const tenantAdminTool = getTool(tools, 'coord_list_tasks');
    await expect(tenantAdminTool.handler({})).rejects.toThrow('MCP tool scope denied: coord:read');

    tools.clear();
    service.registerTools(server, platformAdmin);
    const platformAdminTool = getTool(tools, 'coord_list_tasks');
    await platformAdminTool.handler({ projectPath: '/tmp/victim' });
    expect(coord.listTasks).toHaveBeenCalledWith(process.cwd());
  });

  it('does not attach a guessed or stale-scope MCP session to another authenticated context', async () => {
    const { service } = makeService();
    const owner = makeMemberActor('owner-user');
    const attacker = makeMemberActor('attacker-user');
    const admin = makeAdminActor('admin-user');
    const downgradedAdmin = makeMemberActor('admin-user');

    const { sessionId, transport } = service.createSession(owner);
    const staleSession = service.createSession(admin);

    expect(service.getSession(sessionId, owner)).toBe(transport);
    expect(service.getSession(sessionId, attacker)).toBeNull();
    expect(service.getSession(sessionId, owner)).toBe(transport);
    expect(service.getSession(staleSession.sessionId, downgradedAdmin)).toBeNull();
    expect(service.getSession(staleSession.sessionId, admin)).toBe(staleSession.transport);

    await service.onModuleDestroy();
  });
});