import { Injectable, Logger, Inject, OnModuleDestroy } from '@nestjs/common';
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { randomUUID } from 'node:crypto';
import { z } from 'zod';
import type { Brain } from '@mosaicstack/brain';
import type { Memory } from '@mosaicstack/memory';
import { BRAIN } from '../brain/brain.tokens.js';
import { MEMORY } from '../memory/memory.tokens.js';
import { EmbeddingService } from '../memory/embedding.service.js';
import { CoordService } from '../coord/coord.service.js';

export const MCP_CALLER_IDENTITY_FIELDS = [
  'actorId',
  'actor',
  'authenticatedUserId',
  'channel',
  'organizationId',
  'ownerId',
  'sessionUserId',
  'teamId',
  'tenant',
  'tenantId',
  'user',
  'userId',
] as const;

type McpCallerIdentityField = (typeof MCP_CALLER_IDENTITY_FIELDS)[number];

export const MCP_TOOL_SCOPES = {
  brain_list_projects: 'brain:project:read',
  brain_get_project: 'brain:project:read',
  brain_list_tasks: 'brain:task:read',
  brain_create_task: 'brain:task:write',
  brain_update_task: 'brain:task:write',
  brain_list_missions: 'brain:mission:read',
  brain_list_conversations: 'brain:conversation:read',
  memory_search: 'memory:insight:read',
  memory_get_preferences: 'memory:preference:read',
  memory_save_preference: 'memory:preference:write',
  memory_save_insight: 'memory:insight:write',
  coord_mission_status: 'coord:read',
  coord_list_tasks: 'coord:read',
  coord_task_detail: 'coord:read',
} as const;

export type McpToolName = keyof typeof MCP_TOOL_SCOPES;
export type McpToolScope = (typeof MCP_TOOL_SCOPES)[McpToolName];

export interface McpActorContext {
  userId: string;
  tenantId: string;
  role: string;
  channel: 'mcp';
  correlationId: string;
  scopes: ReadonlySet<McpToolScope>;
}

interface SessionEntry {
  server: McpServer;
  transport: StreamableHTTPServerTransport;
  createdAt: Date;
  actor: McpActorContext;
}

const GLOBAL_ADMIN_MCP_SCOPES = new Set<McpToolScope>(Object.values(MCP_TOOL_SCOPES));
const TENANT_ADMIN_MCP_SCOPES = new Set<McpToolScope>([
  MCP_TOOL_SCOPES.brain_list_projects,
  MCP_TOOL_SCOPES.brain_get_project,
  MCP_TOOL_SCOPES.brain_list_tasks,
  MCP_TOOL_SCOPES.brain_create_task,
  MCP_TOOL_SCOPES.brain_update_task,
  MCP_TOOL_SCOPES.brain_list_missions,
  MCP_TOOL_SCOPES.brain_list_conversations,
  MCP_TOOL_SCOPES.memory_search,
  MCP_TOOL_SCOPES.memory_get_preferences,
  MCP_TOOL_SCOPES.memory_save_preference,
  MCP_TOOL_SCOPES.memory_save_insight,
]);
const MEMBER_MCP_SCOPES = new Set<McpToolScope>([
  MCP_TOOL_SCOPES.brain_list_projects,
  MCP_TOOL_SCOPES.brain_get_project,
  MCP_TOOL_SCOPES.brain_list_tasks,
  MCP_TOOL_SCOPES.brain_list_missions,
  MCP_TOOL_SCOPES.brain_list_conversations,
  MCP_TOOL_SCOPES.memory_search,
  MCP_TOOL_SCOPES.memory_get_preferences,
  MCP_TOOL_SCOPES.memory_save_preference,
  MCP_TOOL_SCOPES.memory_save_insight,
]);

export function deriveMcpToolScopesForUser(input: {
  role?: string | null;
}): ReadonlySet<McpToolScope> {
  if (input.role === 'platform-admin' || input.role === 'super-admin') {
    return new Set(GLOBAL_ADMIN_MCP_SCOPES);
  }
  if (input.role === 'admin') {
    return new Set(TENANT_ADMIN_MCP_SCOPES);
  }
  return new Set(MEMBER_MCP_SCOPES);
}

export function createMcpActorContext(input: {
  userId: string;
  tenantId?: string;
  role?: string | null;
  scopes?: Iterable<McpToolScope>;
  correlationId?: string;
}): McpActorContext {
  const userId = input.userId.trim();
  if (userId.length === 0) {
    throw new Error('MCP authenticated user is required');
  }

  return {
    userId,
    tenantId: input.tenantId?.trim() || `user:${userId}`,
    role: input.role ?? 'member',
    channel: 'mcp',
    correlationId: input.correlationId ?? randomUUID(),
    scopes: new Set(input.scopes ?? []),
  };
}

export function assertNoCallerControlledIdentity(params: unknown): void {
  if (params === null || typeof params !== 'object') return;

  const keys = new Set(Object.keys(params));
  const forbidden = MCP_CALLER_IDENTITY_FIELDS.find((field: McpCallerIdentityField) =>
    keys.has(field),
  );
  if (forbidden) {
    throw new Error(`MCP caller-controlled identity field is forbidden: ${forbidden}`);
  }
}

export function assertMcpToolAuthorized(
  actor: McpActorContext,
  toolName: McpToolName,
  params: unknown,
): void {
  assertNoCallerControlledIdentity(params);
  const requiredScope = MCP_TOOL_SCOPES[toolName];
  if (!actor.scopes.has(requiredScope)) {
    throw new Error(`MCP tool scope denied: ${requiredScope}`);
  }
}

function strictObject<T extends z.ZodRawShape>(shape: T): z.ZodObject<T> {
  return z.object(shape).strict();
}

type TenantScopedLike = {
  tenantId?: string | null;
  organizationId?: string | null;
  teamId?: string | null;
};
type ProjectLike = TenantScopedLike & { id: string; ownerId?: string | null };
type MissionLike = TenantScopedLike & {
  id: string;
  projectId?: string | null;
  userId?: string | null;
};
type TaskLike = TenantScopedLike & {
  projectId?: string | null;
  missionId?: string | null;
  userId?: string | null;
};

function isGlobalAdminActor(actor: McpActorContext): boolean {
  return actor.role === 'platform-admin' || actor.role === 'super-admin';
}

function isTenantAdminActor(actor: McpActorContext): boolean {
  return actor.role === 'admin';
}

function matchesTenant(actor: McpActorContext, record: TenantScopedLike): boolean {
  return (
    record.tenantId === actor.tenantId ||
    record.organizationId === actor.tenantId ||
    record.teamId === actor.tenantId
  );
}

function filterProjectsForActor<T extends ProjectLike>(actor: McpActorContext, projects: T[]): T[] {
  if (isGlobalAdminActor(actor)) return projects;
  return projects.filter(
    (project) =>
      project.ownerId === actor.userId ||
      (isTenantAdminActor(actor) && matchesTenant(actor, project)),
  );
}

function filterMissionsByDirectActorScope<T extends MissionLike>(
  actor: McpActorContext,
  missions: T[],
): T[] {
  if (isGlobalAdminActor(actor)) return missions;
  return missions.filter(
    (mission) =>
      mission.userId === actor.userId ||
      (isTenantAdminActor(actor) && matchesTenant(actor, mission)),
  );
}

function scopesEqual(left: ReadonlySet<McpToolScope>, right: ReadonlySet<McpToolScope>): boolean {
  if (left.size !== right.size) return false;
  for (const scope of left) {
    if (!right.has(scope)) return false;
  }
  return true;
}

function sameActorAuthorization(stored: McpActorContext, current: McpActorContext): boolean {
  return (
    stored.userId === current.userId &&
    stored.tenantId === current.tenantId &&
    stored.role === current.role &&
    scopesEqual(stored.scopes, current.scopes)
  );
}

@Injectable()
export class McpService implements OnModuleDestroy {
  private readonly logger = new Logger(McpService.name);
  private readonly sessions = new Map<string, SessionEntry>();

  constructor(
    @Inject(BRAIN) private readonly brain: Brain,
    @Inject(MEMORY) private readonly memory: Memory,
    @Inject(EmbeddingService) private readonly embeddings: EmbeddingService,
    @Inject(CoordService) private readonly coordService: CoordService,
  ) {}

  /**
   * Creates a new MCP session with its own server + transport pair.
   * Returns the transport for use by the controller.
   */
  createSession(actor: McpActorContext): {
    sessionId: string;
    transport: StreamableHTTPServerTransport;
  } {
    const sessionId = randomUUID();

    const transport = new StreamableHTTPServerTransport({
      sessionIdGenerator: () => sessionId,
      onsessioninitialized: (id) => {
        this.logger.log(
          `MCP session initialized: ${id} for actor=${actor.userId} tenant=${actor.tenantId} correlation=${actor.correlationId}`,
        );
      },
    });

    const server = new McpServer(
      { name: 'mosaic-gateway', version: '1.0.0' },
      { capabilities: { tools: {} } },
    );

    this.registerTools(server, actor);

    transport.onclose = () => {
      this.logger.log(`MCP session closed: ${sessionId}`);
      this.sessions.delete(sessionId);
    };

    server.connect(transport).catch((err: unknown) => {
      this.logger.error(
        `MCP server connect error for session ${sessionId}: ${err instanceof Error ? err.message : String(err)}`,
      );
    });

    this.sessions.set(sessionId, { server, transport, createdAt: new Date(), actor });
    return { sessionId, transport };
  }

  /**
   * Returns the transport for an existing session only when it belongs to the
   * currently authenticated MCP actor. Guessed or cross-tenant session IDs grant
   * no authority.
   */
  getSession(sessionId: string, actor: McpActorContext): StreamableHTTPServerTransport | null {
    const entry = this.sessions.get(sessionId);
    if (!entry) return null;
    if (!sameActorAuthorization(entry.actor, actor)) {
      this.logger.warn(
        `MCP session actor or scope mismatch: session=${sessionId} actor=${actor.userId} tenant=${actor.tenantId} role=${actor.role}`,
      );
      return null;
    }
    return entry.transport;
  }

  private async isProjectAuthorized(actor: McpActorContext, projectId: string): Promise<boolean> {
    if (isGlobalAdminActor(actor)) return true;
    const project = (await this.brain.projects.findById(projectId)) as ProjectLike | undefined;
    return project ? filterProjectsForActor(actor, [project]).length === 1 : false;
  }

  private async filterMissionsForActor<T extends MissionLike>(
    actor: McpActorContext,
    missions: T[],
  ): Promise<T[]> {
    if (isGlobalAdminActor(actor)) return missions;

    const projects = (await this.brain.projects.findAll()) as ProjectLike[];
    const projectIds = new Set(
      filterProjectsForActor(actor, projects).map((project) => project.id),
    );

    return missions.filter(
      (mission) =>
        filterMissionsByDirectActorScope(actor, [mission]).length === 1 ||
        (typeof mission.projectId === 'string' && projectIds.has(mission.projectId)),
    );
  }

  private async isMissionAuthorized(actor: McpActorContext, missionId: string): Promise<boolean> {
    if (isGlobalAdminActor(actor)) return true;
    const mission = (await this.brain.missions.findById(missionId)) as MissionLike | undefined;
    if (!mission) return false;
    return (await this.filterMissionsForActor(actor, [mission])).length === 1;
  }

  private async assertTaskReferencesAuthorized(
    actor: McpActorContext,
    refs: { projectId?: string | null; missionId?: string | null },
  ): Promise<void> {
    if (refs.projectId && !(await this.isProjectAuthorized(actor, refs.projectId))) {
      throw new Error('MCP task project scope denied');
    }
    if (refs.missionId && !(await this.isMissionAuthorized(actor, refs.missionId))) {
      throw new Error('MCP task mission scope denied');
    }
  }

  private async assertTaskCreateScopeAuthorized(
    actor: McpActorContext,
    refs: { projectId?: string | null; missionId?: string | null },
  ): Promise<void> {
    if (!isGlobalAdminActor(actor) && !refs.projectId && !refs.missionId) {
      throw new Error('MCP task scope denied');
    }
    await this.assertTaskReferencesAuthorized(actor, refs);
  }

  private async filterTasksForActor<T extends TaskLike>(
    actor: McpActorContext,
    tasks: T[],
  ): Promise<T[]> {
    if (isGlobalAdminActor(actor)) return tasks;

    const [projects, missions] = await Promise.all([
      this.brain.projects.findAll(),
      this.brain.missions.findAll(),
    ]);
    const projectIds = new Set(
      filterProjectsForActor(actor, projects as ProjectLike[]).map((project) => project.id),
    );
    const missionIds = new Set(
      (await this.filterMissionsForActor(actor, missions as MissionLike[])).map(
        (mission) => mission.id,
      ),
    );

    return tasks.filter(
      (task) =>
        task.userId === actor.userId ||
        (isTenantAdminActor(actor) && matchesTenant(actor, task)) ||
        (typeof task.projectId === 'string' && projectIds.has(task.projectId)) ||
        (typeof task.missionId === 'string' && missionIds.has(task.missionId)),
    );
  }

  /**
   * Registers all platform tools on the given McpServer instance.
   */
  registerTools(server: McpServer, actor: McpActorContext): void {
    // ─── Brain: Project tools ────────────────────────────────────────────

    server.registerTool(
      'brain_list_projects',
      {
        description: 'List all projects in the brain.',
        inputSchema: strictObject({}),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'brain_list_projects', params);
        const projects = filterProjectsForActor(
          actor,
          (await this.brain.projects.findAll()) as ProjectLike[],
        );
        return {
          content: [{ type: 'text' as const, text: JSON.stringify(projects, null, 2) }],
        };
      },
    );

    server.registerTool(
      'brain_get_project',
      {
        description: 'Get a project by ID.',
        inputSchema: strictObject({
          id: z.string().describe('Project ID (UUID)'),
        }),
      },
      async ({ id, ...params }) => {
        assertMcpToolAuthorized(actor, 'brain_get_project', params);
        const project = (await this.brain.projects.findById(id)) as ProjectLike | undefined;
        const authorizedProject = project ? filterProjectsForActor(actor, [project])[0] : undefined;
        return {
          content: [
            {
              type: 'text' as const,
              text: authorizedProject
                ? JSON.stringify(authorizedProject, null, 2)
                : `Project not found: ${id}`,
            },
          ],
        };
      },
    );

    // ─── Brain: Task tools ───────────────────────────────────────────────

    server.registerTool(
      'brain_list_tasks',
      {
        description: 'List tasks, optionally filtered by project, mission, or status.',
        inputSchema: strictObject({
          projectId: z.string().optional().describe('Filter by project ID'),
          missionId: z.string().optional().describe('Filter by mission ID'),
          status: z.string().optional().describe('Filter by status'),
        }),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'brain_list_tasks', params);
        const { projectId, missionId, status } = params;
        type TaskStatus = 'not-started' | 'in-progress' | 'blocked' | 'done' | 'cancelled';
        let tasks;
        if (projectId) tasks = await this.brain.tasks.findByProject(projectId);
        else if (missionId) tasks = await this.brain.tasks.findByMission(missionId);
        else if (status) tasks = await this.brain.tasks.findByStatus(status as TaskStatus);
        else tasks = await this.brain.tasks.findAll();
        const scopedTasks = await this.filterTasksForActor(actor, tasks as TaskLike[]);
        return { content: [{ type: 'text' as const, text: JSON.stringify(scopedTasks, null, 2) }] };
      },
    );

    server.registerTool(
      'brain_create_task',
      {
        description: 'Create a new task in the brain.',
        inputSchema: strictObject({
          title: z.string().describe('Task title'),
          description: z.string().optional().describe('Task description'),
          projectId: z.string().optional().describe('Project ID'),
          missionId: z.string().optional().describe('Mission ID'),
          priority: z.string().optional().describe('Priority: low, medium, high, critical'),
        }),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'brain_create_task', params);
        await this.assertTaskCreateScopeAuthorized(actor, params);
        type Priority = 'low' | 'medium' | 'high' | 'critical';
        const task = await this.brain.tasks.create({
          ...params,
          priority: params.priority as Priority | undefined,
        });
        return { content: [{ type: 'text' as const, text: JSON.stringify(task, null, 2) }] };
      },
    );

    server.registerTool(
      'brain_update_task',
      {
        description: 'Update an existing task.',
        inputSchema: strictObject({
          id: z.string().describe('Task ID'),
          title: z.string().optional(),
          description: z.string().optional(),
          status: z
            .string()
            .optional()
            .describe('not-started, in-progress, blocked, done, cancelled'),
          priority: z.string().optional(),
          projectId: z.string().optional().describe('Project ID'),
          missionId: z.string().optional().describe('Mission ID'),
        }),
      },
      async ({ id, ...updates }) => {
        assertMcpToolAuthorized(actor, 'brain_update_task', updates);
        const existing = (await this.brain.tasks.findById(id)) as TaskLike | undefined;
        if (!existing || (await this.filterTasksForActor(actor, [existing])).length === 0) {
          return { content: [{ type: 'text' as const, text: `Task not found: ${id}` }] };
        }
        await this.assertTaskReferencesAuthorized(actor, updates);
        type TaskStatus = 'not-started' | 'in-progress' | 'blocked' | 'done' | 'cancelled';
        type Priority = 'low' | 'medium' | 'high' | 'critical';
        const task = await this.brain.tasks.update(id, {
          ...updates,
          status: updates.status as TaskStatus | undefined,
          priority: updates.priority as Priority | undefined,
        });
        return {
          content: [
            {
              type: 'text' as const,
              text: task ? JSON.stringify(task, null, 2) : `Task not found: ${id}`,
            },
          ],
        };
      },
    );

    // ─── Brain: Mission tools ────────────────────────────────────────────

    server.registerTool(
      'brain_list_missions',
      {
        description: 'List all missions, optionally filtered by project.',
        inputSchema: strictObject({
          projectId: z.string().optional().describe('Filter by project ID'),
        }),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'brain_list_missions', params);
        const { projectId } = params;
        const missions = await this.filterMissionsForActor(
          actor,
          (projectId
            ? await this.brain.missions.findByProject(projectId)
            : await this.brain.missions.findAll()) as MissionLike[],
        );
        return { content: [{ type: 'text' as const, text: JSON.stringify(missions, null, 2) }] };
      },
    );

    server.registerTool(
      'brain_list_conversations',
      {
        description: 'List conversations for the authenticated MCP actor.',
        inputSchema: strictObject({}),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'brain_list_conversations', params);
        const conversations = await this.brain.conversations.findAll(actor.userId);
        return {
          content: [{ type: 'text' as const, text: JSON.stringify(conversations, null, 2) }],
        };
      },
    );

    // ─── Memory tools ────────────────────────────────────────────────────

    server.registerTool(
      'memory_search',
      {
        description:
          'Search stored insights and knowledge for the authenticated MCP actor using natural language.',
        inputSchema: strictObject({
          query: z.string().describe('Natural language search query'),
          limit: z.number().optional().describe('Max results (default 5)'),
        }),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'memory_search', params);
        const { query, limit } = params;
        if (!this.embeddings.available) {
          return {
            content: [
              {
                type: 'text' as const,
                text: 'Semantic search unavailable — no embedding provider configured',
              },
            ],
          };
        }
        const embedding = await this.embeddings.embed(query);
        const results = await this.memory.insights.searchByEmbedding(
          actor.userId,
          embedding,
          limit ?? 5,
        );
        return { content: [{ type: 'text' as const, text: JSON.stringify(results, null, 2) }] };
      },
    );

    server.registerTool(
      'memory_get_preferences',
      {
        description: 'Retrieve stored preferences for the authenticated MCP actor.',
        inputSchema: strictObject({
          category: z
            .string()
            .optional()
            .describe('Filter by category: communication, coding, workflow, appearance, general'),
        }),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'memory_get_preferences', params);
        const { category } = params;
        type Cat = 'communication' | 'coding' | 'workflow' | 'appearance' | 'general';
        const prefs = category
          ? await this.memory.preferences.findByUserAndCategory(actor.userId, category as Cat)
          : await this.memory.preferences.findByUser(actor.userId);
        return { content: [{ type: 'text' as const, text: JSON.stringify(prefs, null, 2) }] };
      },
    );

    server.registerTool(
      'memory_save_preference',
      {
        description:
          'Store a learned preference for the authenticated MCP actor (e.g., "prefers tables over paragraphs").',
        inputSchema: strictObject({
          key: z.string().describe('Preference key'),
          value: z.string().describe('Preference value (JSON string)'),
          category: z
            .string()
            .optional()
            .describe('Category: communication, coding, workflow, appearance, general'),
        }),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'memory_save_preference', params);
        const { key, value, category } = params;
        type Cat = 'communication' | 'coding' | 'workflow' | 'appearance' | 'general';
        let parsedValue: unknown;
        try {
          parsedValue = JSON.parse(value);
        } catch {
          parsedValue = value;
        }
        const pref = await this.memory.preferences.upsert({
          userId: actor.userId,
          key,
          value: parsedValue,
          category: (category as Cat) ?? 'general',
          source: 'agent',
        });
        return { content: [{ type: 'text' as const, text: JSON.stringify(pref, null, 2) }] };
      },
    );

    server.registerTool(
      'memory_save_insight',
      {
        description:
          'Store a learned insight, decision, or knowledge for the authenticated MCP actor.',
        inputSchema: strictObject({
          content: z.string().describe('The insight or knowledge to store'),
          category: z
            .string()
            .optional()
            .describe('Category: decision, learning, preference, fact, pattern, general'),
        }),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'memory_save_insight', params);
        const { content, category } = params;
        type Cat = 'decision' | 'learning' | 'preference' | 'fact' | 'pattern' | 'general';
        const embedding = this.embeddings.available ? await this.embeddings.embed(content) : null;
        const insight = await this.memory.insights.create({
          userId: actor.userId,
          content,
          embedding,
          source: 'agent',
          category: (category as Cat) ?? 'learning',
        });
        return { content: [{ type: 'text' as const, text: JSON.stringify(insight, null, 2) }] };
      },
    );

    // ─── Coord tools ─────────────────────────────────────────────────────

    server.registerTool(
      'coord_mission_status',
      {
        description:
          'Get the current orchestration mission status including milestones, tasks, and active session.',
        inputSchema: strictObject({}),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'coord_mission_status', params);
        const status = await this.coordService.getMissionStatus(process.cwd());
        return {
          content: [
            {
              type: 'text' as const,
              text: status ? JSON.stringify(status, null, 2) : 'No active coord mission found.',
            },
          ],
        };
      },
    );

    server.registerTool(
      'coord_list_tasks',
      {
        description: 'List all tasks from the orchestration TASKS.md file.',
        inputSchema: strictObject({}),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'coord_list_tasks', params);
        const tasks = await this.coordService.listTasks(process.cwd());
        return { content: [{ type: 'text' as const, text: JSON.stringify(tasks, null, 2) }] };
      },
    );

    server.registerTool(
      'coord_task_detail',
      {
        description: 'Get detailed status for a specific orchestration task.',
        inputSchema: strictObject({
          taskId: z.string().describe('Task ID (e.g. P2-005)'),
        }),
      },
      async (params) => {
        assertMcpToolAuthorized(actor, 'coord_task_detail', params);
        const { taskId } = params;
        const detail = await this.coordService.getTaskStatus(process.cwd(), taskId);
        return {
          content: [
            {
              type: 'text' as const,
              text: detail
                ? JSON.stringify(detail, null, 2)
                : `Task ${taskId} not found in coord mission.`,
            },
          ],
        };
      },
    );
  }

  async onModuleDestroy(): Promise<void> {
    this.logger.log(`Closing ${this.sessions.size} MCP sessions on shutdown`);
    const closePromises = Array.from(this.sessions.values()).map(({ transport }) =>
      transport.close().catch((err: unknown) => {
        this.logger.warn(
          `Error closing MCP transport: ${err instanceof Error ? err.message : String(err)}`,
        );
      }),
    );
    await Promise.all(closePromises);
    this.sessions.clear();
  }
}