import { z } from 'zod';

import { REASONING_OPTIONS } from '../services/reasoning-options.js';

// ---------------------------------------------------------------------------
// Zod schemas — passed to McpServer.registerTool (SDK converts to JSON Schema)
// ---------------------------------------------------------------------------

const reasoningEnum = z.enum(REASONING_OPTIONS);

export const spawnTaskSchema = z.object({
  prompt: z.string().min(1).describe('What the task should do. Be specific: include file paths, function or symbol names, the expected outcome, and any constraints.'),
  task_type: z.enum(['coder', 'planner', 'tester', 'researcher', 'general']).default('coder').describe('Routing hint: coder for code, planner for decomposing work, tester for tests, researcher for investigation, general for anything else.'),
  provider: z.enum(['codex', 'copilot', 'claude-cli']).optional().describe('Force a specific backend. Leave unset in almost all cases.'),
  reasoning: reasoningEnum.optional().describe('Model + reasoning effort. gpt-5.4(medium) is the default workhorse; use high for multi-file reasoning, xhigh for deep research, low for trivial edits.'),
  cwd: z.string().optional().describe('Absolute working directory. Defaults to server process cwd.'),
  timeout_ms: z.number().int().min(1000).max(3_600_000).optional().describe('Hard time limit in ms. Task is marked timed_out if exceeded.'),
  keep_alive: z.number().optional().describe('Retention window (ms) — how long the server keeps the completed result available.'),
  labels: z.array(z.string()).optional().describe('Free-form tags for filtering and grouping.'),
  depends_on: z.array(z.string()).optional().describe('Task IDs that must complete before this task starts.'),
  developer_instructions: z.string().optional().describe('System-level instructions injected ahead of the user prompt.'),
  context_files: z.array(z.object({
    path: z.string().describe('Absolute path of a file to include as context.'),
    description: z.string().optional().describe('Note explaining why this file is relevant.'),
  })).optional().describe('Files to prepend as additional context. Use sparingly.'),
});

export const waitTaskSchema = z.object({
  task_id: z.string().min(1).describe('ID of the task to wait on, as returned by spawn-task.'),
  timeout_ms: z.number().int().positive().max(300_000).default(30_000).optional().describe('Max time to block in ms. Returns early on terminal or input_required. Default 30s.'),
  poll_interval_ms: z.number().int().min(250).max(30_000).default(1000).optional().describe('Internal poll interval in ms. Keep the default unless tuning.'),
});

export const respondTaskSchema = z.discriminatedUnion('type', [
  z.object({
    task_id: z.string().min(1),
    type: z.literal('user_input'),
    answers: z.record(z.string(), z.string()).describe('Map of question id → answer string.'),
  }),
  z.object({
    task_id: z.string().min(1),
    type: z.literal('command_approval'),
    decision: z.enum(['accept', 'reject']).describe('Whether the agent may run the proposed command.'),
  }),
  z.object({
    task_id: z.string().min(1),
    type: z.literal('file_approval'),
    decision: z.enum(['accept', 'reject']).describe('Whether the agent may apply the proposed file edit.'),
  }),
  z.object({
    task_id: z.string().min(1),
    type: z.literal('elicitation'),
    action: z.enum(['accept', 'decline']).describe('Accept or decline the MCP elicitation request.'),
    content: z.record(z.string(), z.unknown()).optional().describe('Structured payload when accepting.'),
  }),
  z.object({
    task_id: z.string().min(1),
    type: z.literal('dynamic_tool'),
    result: z.string().optional().describe('Tool call result string returned to the agent.'),
    error: z.string().optional().describe('Error string if the tool call failed.'),
  }),
]);

export const messageTaskSchema = z.object({
  task_id: z.string().min(1).describe('ID of the task whose session should receive the follow-up.'),
  message: z.string().min(1).describe('The follow-up instruction or question. Be as specific as the original prompt.'),
  reasoning: reasoningEnum.optional().describe('Overrides reasoning for this follow-up turn only.'),
});

export const cancelTaskSchema = z.object({
  task_id: z.union([z.string().min(1), z.array(z.string().min(1)).min(1)]).describe('Task ID or array of task IDs to cancel.'),
});

// ---------------------------------------------------------------------------
// Inferred types — used by app.ts handlers
// ---------------------------------------------------------------------------

export type SpawnTaskInput = z.infer<typeof spawnTaskSchema>;
export type WaitTaskInput = z.infer<typeof waitTaskSchema>;
export type RespondTaskInput = z.infer<typeof respondTaskSchema>;
export type MessageTaskInput = z.infer<typeof messageTaskSchema>;
export type CancelTaskInput = z.infer<typeof cancelTaskSchema>;

// ---------------------------------------------------------------------------
// Tool descriptions — keyed by name, used by index.ts registerTool
// ---------------------------------------------------------------------------

export function getToolDescriptions(serverVersion?: string): Map<string, string> {
  const v = serverVersion ? ` (v${serverVersion})` : '';
  return new Map<string, string>([
    ['spawn-task', [
      `Create and start a provider-agnostic task, returning a task_id you can track.${v}`,
      '',
      'Dispatches the prompt to the provider registered for the given `task_type` (Codex, Copilot, Claude CLI) and returns immediately with a task_id. Use `wait-task` to block until the task reaches a terminal state or needs input, `respond-task` to unblock it, and `message-task` to send follow-ups on the same session.',
      '',
      'PARALLEL EXECUTION: Spawn multiple tasks in the same message to fan out work — each task runs in its own isolated agent workspace and reports back independently.',
      '',
      'AFTER SPAWNING: Always follow up with `wait-task`. The agent may pause almost immediately to request input; the bridge window surfaces that pending question so you can answer it without polling.',
      '',
      'WRITING A GOOD PROMPT: Name the exact files, functions, or symbols involved, state the expected behavior, and mention anything the agent must NOT touch.',
    ].join('\n')],
    ['wait-task', [
      'Block until a task settles or asks for input.',
      '',
      'Returns as soon as the task reaches a terminal state (`completed`, `failed`, `cancelled`) or enters `input_required`. If `timeout_ms` elapses first, it returns the current status.',
      '',
      'PATTERN: loop `wait-task` → if `input_required`, call `respond-task` → loop back. Give each call enough time to catch meaningful progress.',
    ].join('\n')],
    ['respond-task', [
      'Unblock a task that is in `input_required` because the agent requested input or an approval.',
      '',
      'Payload shape is discriminated by `type` — must match the pending question from `wait-task`:',
      '- `user_input` — `answers` map of question id → string',
      '- `command_approval` — `decision: "accept" | "reject"`',
      '- `file_approval` — `decision: "accept" | "reject"`',
      '- `elicitation` — `action: "accept" | "decline"`, optional `content`',
      '- `dynamic_tool` — `result` or `error`',
      '',
      'After responding the task resumes automatically. Follow up with `wait-task`.',
    ].join('\n')],
    ['message-task', [
      'Send a follow-up message to an existing task on its original session.',
      '',
      'Use this to steer a running task, refine completed work, or add instructions after reviewing partial results.',
      '',
      'After calling, follow up with `wait-task` exactly like after `spawn-task`.',
    ].join('\n')],
    ['cancel-task', [
      'Cancel one or more running tasks.',
      '',
      'Accepts a single task_id or an array. Running tasks are aborted and marked `cancelled`. Already-terminal tasks are listed under `already_terminal`; unknown IDs under `not_found`.',
    ].join('\n')],
  ]);
}

// ---------------------------------------------------------------------------
// Tool annotations — MCP SDK annotations per tool
// ---------------------------------------------------------------------------

export interface ToolAnnotation {
  readOnlyHint: boolean;
  destructiveHint: boolean;
  idempotentHint: boolean;
  openWorldHint: boolean;
}

const ANNOTATIONS = new Map<string, ToolAnnotation>([
  ['spawn-task', { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: true }],
  ['wait-task', { readOnlyHint: true, destructiveHint: false, idempotentHint: true, openWorldHint: false }],
  ['respond-task', { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: false }],
  ['message-task', { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: true }],
  ['cancel-task', { readOnlyHint: false, destructiveHint: true, idempotentHint: true, openWorldHint: false }],
]);

export function getToolAnnotation(name: string): ToolAnnotation {
  return ANNOTATIONS.get(name) ?? { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: false };
}