import { McpServer, type CallToolResult } from '@modelcontextprotocol/server';
import * as z from 'zod/v4';
import { OrchestratorError, OrchestratorService } from '@platform-modules/chatgpt-orchestrator-core';
import { JsonFileRepository, type OrchestratorRepository } from '@platform-modules/chatgpt-orchestrator-persistence';
import {
  ContinuationPolicyPatchSchema,
  MessageDeliverySchema,
  MessagePrioritySchema,
  MessageTypeSchema,
  RunIdSchema,
  WorkerIdSchema,
  type Assignment,
  type ContinuationPolicy,
  type Run,
  type Worker,
  type RunId,
  type WorkerId,
} from '@platform-modules/chatgpt-orchestrator-protocol';
import { workerBootstrap } from './bootstrap.js';
import { commandForExecutionAttempt, ExecutionContinuityScheduler } from './execution-continuity.js';
import type { ExecutorBridge } from './executor-bridge.js';

export const serverInfo = { name: '@platform-modules/chatgpt-orchestrator-mcp', version: '0.1.0' } as const;

function success(structuredContent: Record<string, unknown>): CallToolResult {
  return { content: [{ type: 'text', text: JSON.stringify(structuredContent) }], structuredContent };
}

function failure(error: unknown): CallToolResult {
  const body = error instanceof OrchestratorError
    ? { code: error.code, message: error.message }
    : { code: 'INTERNAL_ERROR', message: error instanceof Error ? error.message : 'Unexpected orchestrator failure.' };
  return { content: [{ type: 'text', text: JSON.stringify({ error: body }) }], isError: true };
}

async function run(fn: () => Promise<Record<string, unknown>>): Promise<CallToolResult> {
  try { return success(await fn()); } catch (error) { return failure(error); }
}

function definedPolicyPatch(policy: Partial<ContinuationPolicy>): Partial<ContinuationPolicy> {
  return Object.fromEntries(Object.entries(policy).filter(([, value]) => value !== undefined)) as Partial<ContinuationPolicy>;
}

async function workerPromptContext(service: OrchestratorService, worker: Worker, assignment: Assignment): Promise<{ run: Run; worker: Worker; assignment: Assignment }> {
  return { run: await service.getRun(worker.runId), worker, assignment };
}

const assignmentInput = z.object({
  objective: z.string().min(1),
  constraints: z.array(z.string()).optional(),
  acceptanceCriteria: z.array(z.string()).optional(),
  dependencies: z.array(WorkerIdSchema).optional(),
  allowedScope: z.array(z.string()).optional(),
  forbiddenScope: z.array(z.string()).optional(),
  artifacts: z.array(z.string()).optional(),
});

export function createOrchestratorMcpServer(
  service: OrchestratorService,
  bridge?: ExecutorBridge,
  scheduler?: ExecutionContinuityScheduler,
): McpServer {
  const server = new McpServer(serverInfo);
  const continuity = scheduler ?? new ExecutionContinuityScheduler(service, bridge);

  server.registerTool('run.create', {
    title: 'Create Orchestration Run',
    description: 'Create one durable orchestration run and its /root logical worker.',
    inputSchema: z.object({ title: z.string().min(1), taskName: z.string().min(1).optional(), rootConversationName: z.string().min(1).optional(), idempotencyKey: z.string().min(1).optional() }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ title, taskName, rootConversationName, idempotencyKey }) => run(async () => {
    const result = await service.createRun({ title, ...(taskName ? { taskName } : {}), ...(rootConversationName ? { rootConversationName } : {}), ...(idempotencyKey ? { idempotencyKey } : {}) });
    const attempt = (await service.listExecutionAttempts(result.rootWorker.workerId))
      .find((candidate) => candidate.reason === 'initial');
    if (!attempt) throw new OrchestratorError('INTERNAL_ERROR', 'Root initial execution attempt is missing.');
    const context = await service.getWorkerAssignmentContext(result.rootWorker.workerId);
    return { ...result, bootstrap: workerBootstrap(context), executionId: attempt.executionId };
  }));

  server.registerTool('run.get', {
    title: 'Get Orchestration Run',
    description: 'Read current durable run state.',
    inputSchema: z.object({ runId: RunIdSchema }),
    annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ runId }) => run(async () => ({ run: await service.getRun(runId as RunId) })));

  server.registerTool('run.update_context', {
    title: 'Update Run Labels',
    description: 'Set the concise main task name and source/root ChatGPT conversation name used in owner-visible worker labels and prompts.',
    inputSchema: z.object({ runId: RunIdSchema, taskName: z.string().min(1).optional(), rootConversationName: z.string().min(1).optional() })
      .refine((value) => Boolean(value.taskName || value.rootConversationName), { message: 'At least one context field is required.' }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ runId, taskName, rootConversationName }) => run(async () => ({
    ...await service.updateRunContext(runId as RunId, { ...(taskName ? { taskName } : {}), ...(rootConversationName ? { rootConversationName } : {}) }),
  })));

  server.registerTool('run.cancel', {
    title: 'Cancel Orchestration Run',
    description: 'Cancel one active orchestration run and all of its non-terminal workers.',
    inputSchema: z.object({ runId: RunIdSchema, reason: z.string().min(1) }),
    annotations: { readOnlyHint: false, destructiveHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ runId, reason }) => run(async () => ({ run: await service.cancelRun(runId as RunId, reason) })));



  server.registerTool('run.continuation.get', {
    title: 'Get Run Continuation Policy',
    description: 'Read the durable default execution-continuity policy for a run.',
    inputSchema: z.object({ runId: RunIdSchema }),
    annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ runId }) => run(async () => ({ policy: await service.getRunContinuationPolicy(runId as RunId) })));

  server.registerTool('run.continuation.set_default', {
    title: 'Set Run Continuation Policy',
    description: 'Update the durable default execution-continuity policy for workers without an override.',
    inputSchema: z.object({ runId: RunIdSchema, policy: ContinuationPolicyPatchSchema }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ runId, policy }) => run(async () => {
    const updated = await service.setRunContinuationPolicy(runId as RunId, definedPolicyPatch(policy as Partial<ContinuationPolicy>));
    await continuity.reconcile();
    return { policy: updated };
  }));

  server.registerTool('worker.spawn', {
    title: 'Spawn Logical Worker',
    description: 'Create a logical child worker assignment and request the ChatGPT Web executor to launch a managed conversation.',
    inputSchema: z.object({
      runId: RunIdSchema,
      parentWorkerId: WorkerIdSchema,
      name: z.string().regex(/^[a-zA-Z0-9._-]+$/),
      displayName: z.string().min(1).optional(),
      taskSummary: z.string().min(1).optional(),
      assignment: assignmentInput,
      idempotencyKey: z.string().min(1).optional(),
    }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async (args) => run(async () => {
    const result = await service.spawnWorker({
      runId: args.runId as RunId,
      parentWorkerId: args.parentWorkerId as WorkerId,
      name: args.name,
      ...(args.displayName ? { displayName: args.displayName } : {}),
      ...(args.taskSummary ? { taskSummary: args.taskSummary } : {}),
      assignment: args.assignment as Parameters<OrchestratorService['spawnWorker']>[0]['assignment'],
      ...(args.idempotencyKey ? { idempotencyKey: args.idempotencyKey } : {}),
    });
    const attempt = await continuity.dispatchAttemptForReason(result.worker.workerId, 'initial');
    const command = commandForExecutionAttempt(attempt);
    const context = await workerPromptContext(service, result.worker, result.assignment);
    return { ...result, bootstrap: workerBootstrap(context), executionId: attempt.executionId, executorCommandId: command.commandId };
  }));

  server.registerTool('worker.attach', {
    title: 'Attach Worker',
    description: 'Attach the current ChatGPT conversation to an existing logical worker and retrieve its authoritative assignment.',
    inputSchema: z.object({ workerId: WorkerIdSchema }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId }) => run(async () => ({ ...await service.attachWorker(workerId as WorkerId) })));

  server.registerTool('worker.list', {
    title: 'List Workers',
    description: 'List all logical workers in a run with lifecycle and managed-conversation state.',
    inputSchema: z.object({ runId: RunIdSchema }),
    annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ runId }) => run(async () => {
    const workers = await service.listWorkers(runId as RunId);
    return { workers: await Promise.all(workers.map(async (worker) => ({ worker, conversation: await service.getConversationBinding(worker.workerId) }))) };
  }));



  server.registerTool('worker.continuation.get', {
    title: 'Get Worker Continuation State',
    description: 'Read durable worker execution-continuity state, effective policy, policy source, and latest execution attempt.',
    inputSchema: z.object({ workerId: WorkerIdSchema }),
    annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId }) => run(async () => ({ ...await service.getWorkerContinuity(workerId as WorkerId) })));

  server.registerTool('worker.continuation.set', {
    title: 'Set Worker Continuation Policy',
    description: 'Set a worker-specific execution-continuity policy override, or pass null to inherit the run default.',
    inputSchema: z.object({ workerId: WorkerIdSchema, policy: ContinuationPolicyPatchSchema.nullable() }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId, policy }) => run(async () => {
    const view = await service.setWorkerContinuationPolicy(workerId as WorkerId, policy === null ? null : definedPolicyPatch(policy as Partial<ContinuationPolicy>));
    await continuity.reconcile();
    return { ...view };
  }));

  server.registerTool('worker.execution_attempts', {
    title: 'List Worker Execution Attempts',
    description: 'List durable execution-attempt history for a worker in sequence order.',
    inputSchema: z.object({ workerId: WorkerIdSchema }),
    annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId }) => run(async () => ({ attempts: await service.listExecutionAttempts(workerId as WorkerId) })));

  server.registerTool('worker.await_human', {
    title: 'Await Human Input',
    description: 'Durably block a worker because genuine human input or approval is required.',
    inputSchema: z.object({
      workerId: WorkerIdSchema,
      reason: z.string().min(1),
      request: z.string().min(1),
      choices: z.array(z.string().min(1)).optional(),
      contextRefs: z.array(z.string().min(1)).optional(),
    }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId, reason, request, choices, contextRefs }) => run(async () => ({
    ...await service.awaitHuman(workerId as WorkerId, {
      reason,
      request,
      ...(choices ? { choices } : {}),
      ...(contextRefs ? { contextRefs } : {}),
    }),
  })));

  server.registerTool('worker.await_dependency', {
    title: 'Await Dependency',
    description: 'Durably block a worker on an external or orchestrated dependency.',
    inputSchema: z.object({
      workerId: WorkerIdSchema,
      reason: z.string().min(1),
      dependencies: z.array(WorkerIdSchema).optional(),
    }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId, reason, dependencies }) => run(async () => ({
    ...await service.awaitDependency(workerId as WorkerId, {
      reason,
      ...(dependencies ? { dependencies: dependencies as WorkerId[] } : {}),
    }),
  })));

  server.registerTool('worker.pause', {
    title: 'Pause Worker Continuation',
    description: 'Pause an unfinished worker without cancelling it; an active generation may finish but no new turn is scheduled.',
    inputSchema: z.object({ workerId: WorkerIdSchema }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId }) => run(async () => ({ ...await service.pauseWorker(workerId as WorkerId) })));

  server.registerTool('worker.resume', {
    title: 'Resume Worker Continuation',
    description: 'Clear a worker pause/stall state and make unfinished work runnable again under its effective continuation policy.',
    inputSchema: z.object({ workerId: WorkerIdSchema }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId }) => run(async () => {
    const view = await service.resumeWorker(workerId as WorkerId);
    await continuity.reconcile();
    return { ...view };
  }));

  server.registerTool('worker.continue_now', {
    title: 'Continue Worker Now',
    description: 'Request exactly one manual continuation turn; blocked human/dependency waits require explicit override.',
    inputSchema: z.object({
      workerId: WorkerIdSchema,
      overrideBlock: z.boolean().default(false),
      idempotencyKey: z.string().min(1).optional(),
    }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId, overrideBlock, idempotencyKey }) => run(async () => ({
    ...await continuity.continueNow(workerId as WorkerId, {
      overrideBlock,
      ...(idempotencyKey ? { idempotencyKey } : {}),
    }),
  })));

  server.registerTool('worker.progress', {
    title: 'Report Worker Progress',
    description: 'Append a durable progress event for an active worker.',
    inputSchema: z.object({ workerId: WorkerIdSchema, detail: z.string().min(1) }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: false },
  }, async ({ workerId, detail }) => run(async () => ({ event: await service.progress(workerId as WorkerId, detail) })));

  server.registerTool('worker.complete', {
    title: 'Complete Worker',
    description: 'Mark an active worker completed and append its terminal summary. Completing /root also completes the run after children are terminal.',
    inputSchema: z.object({ workerId: WorkerIdSchema, summary: z.string().min(1) }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: false },
  }, async ({ workerId, summary }) => run(async () => ({ worker: await service.completeWorker(workerId as WorkerId, summary) })));

  server.registerTool('worker.fail', {
    title: 'Fail Worker',
    description: 'Mark a non-terminal worker failed with a durable reason.',
    inputSchema: z.object({ workerId: WorkerIdSchema, reason: z.string().min(1) }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: false },
  }, async ({ workerId, reason }) => run(async () => ({ worker: await service.failWorker(workerId as WorkerId, reason) })));

  server.registerTool('worker.interrupt', {
    title: 'Interrupt Worker',
    description: 'Cancel one non-terminal worker in the control plane.',
    inputSchema: z.object({ workerId: WorkerIdSchema, reason: z.string().min(1) }),
    annotations: { readOnlyHint: false, destructiveHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ workerId, reason }) => run(async () => ({ worker: await service.interruptWorker(workerId as WorkerId, reason) })));

  server.registerTool('worker.followup', {
    title: 'Assign Follow-up Work',
    description: 'Replace a completed or active managed worker assignment with explicit follow-up work and wake its existing ChatGPT conversation.',
    inputSchema: z.object({ workerId: WorkerIdSchema, assignment: assignmentInput }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: false },
  }, async ({ workerId, assignment }) => run(async () => {
    const result = await service.followupWorker(workerId as WorkerId, assignment as Parameters<OrchestratorService['followupWorker']>[1]);
    const attempt = await continuity.dispatchAttemptForReason(result.worker.workerId, 'followup');
    const command = commandForExecutionAttempt(attempt);
    return { ...result, executionId: attempt.executionId, executorCommandId: command.commandId };
  }));

  server.registerTool('message.send', {
    title: 'Send Worker Message',
    description: 'Persist one durable inter-worker message first, then apply its requested delivery policy. steer_now uses internal mid-turn steering only while the target is generating.',
    inputSchema: z.object({
      runId: RunIdSchema,
      fromWorkerId: WorkerIdSchema,
      toWorkerId: WorkerIdSchema,
      type: MessageTypeSchema,
      body: z.string().min(1),
      delivery: MessageDeliverySchema.default('next_turn'),
      priority: MessagePrioritySchema.default('normal'),
      idempotencyKey: z.string().min(1).optional(),
      notBefore: z.string().datetime().optional(),
    }).superRefine((value, ctx) => {
      if (value.delivery === 'steer_now' && !value.idempotencyKey) {
        ctx.addIssue({ code: 'custom', path: ['idempotencyKey'], message: 'steer_now requires idempotencyKey for exactly-once durable delivery.' });
      }
    }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
  }, async (args) => run(async () => ({ message: await continuity.sendMessage({
    runId: args.runId as RunId,
    fromWorkerId: args.fromWorkerId as WorkerId,
    toWorkerId: args.toWorkerId as WorkerId,
    type: args.type,
    body: args.body,
    delivery: args.delivery,
    priority: args.priority,
    ...(args.idempotencyKey ? { idempotencyKey: args.idempotencyKey } : {}),
    ...(args.notBefore ? { notBefore: args.notBefore } : {}),
  }) })));

  server.registerTool('message.broadcast', {
    title: 'Broadcast Worker Message',
    description: 'Send one durable message from a worker to every other worker in the run.',
    inputSchema: z.object({ runId: RunIdSchema, fromWorkerId: WorkerIdSchema, type: MessageTypeSchema, body: z.string().min(1) }),
    annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: false },
  }, async (args) => run(async () => {
    const messages = await service.broadcastMessage({
      runId: args.runId as RunId, fromWorkerId: args.fromWorkerId as WorkerId, type: args.type, body: args.body,
    });
    await Promise.all(messages.map((message) => continuity.dispatchMessage(message)));
    return { messages: await Promise.all(messages.map(async (message) => (await service.getMessageDeliveryContext(message.messageId)).message)) };
  }));

  server.registerTool('events.list', {
    title: 'List Run Events',
    description: 'Read durable run events after an optional cursor.',
    inputSchema: z.object({ runId: RunIdSchema, afterCursor: z.number().int().nonnegative().default(0) }),
    annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ runId, afterCursor }) => run(async () => ({ events: await service.listEvents(runId as RunId, afterCursor) })));

  server.registerTool('events.wait', {
    title: 'Wait for Run Events',
    description: 'Bounded long-poll wait for durable orchestration events after a cursor. Use repeatedly while coordinating active workers.',
    inputSchema: z.object({
      runId: RunIdSchema,
      afterCursor: z.number().int().nonnegative().default(0),
      timeoutMs: z.number().int().min(1).max(60_000).default(30_000),
    }),
    annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false },
  }, async ({ runId, afterCursor, timeoutMs }) => run(async () => ({ ...await service.waitEvents(runId as RunId, afterCursor, timeoutMs) })));

  return server;
}

export function createDefaultRepository(statePath = process.env.CHATGPT_ORCHESTRATOR_STATE ?? `${process.env.HOME ?? '.'}/.local/state/chatgpt-orchestrator/state.json`): JsonFileRepository {
  return new JsonFileRepository(statePath);
}

export function createDefaultService(statePath?: string): OrchestratorService {
  return new OrchestratorService(createDefaultRepository(statePath));
}

export function createOrchestratorMcpServerFactory(repository?: OrchestratorRepository, bridge?: ExecutorBridge): () => McpServer {
  const service = repository ? new OrchestratorService(repository) : createDefaultService();
  return () => createOrchestratorMcpServer(service, bridge);
}