import { createHash, randomBytes } from 'node:crypto';
import { mkdir, open, readFile, readdir, stat, unlink } from 'node:fs/promises';
import path from 'node:path';
import {
  acceptedContent,
  createRequestStateCodec,
  inputRequired,
  type CallToolResult,
  type InputRequiredResult,
  type RequestStateCodec,
  type ServerContext,
  type Tool,
} from '@modelcontextprotocol/server';
import { z } from 'zod';
import { POD_EXEC_TOOL_NAME, normalizeToolName } from '../utils/ToolPolicy.js';

export const GUARDED_TOOL_NAMES = [POD_EXEC_TOOL_NAME] as const;

type GuardedToolName = (typeof GUARDED_TOOL_NAMES)[number];

interface ApprovalState {
  kind: 'code-mode-sensitive-tool-approval';
  toolName: GuardedToolName;
  argumentsHash: string;
  nonce: string;
  expiresAt: number;
}

interface ApprovalReplayStore {
  claim(nonce: string, expiresAt: number): Promise<boolean>;
}

const APPROVAL_TTL_SECONDS = 600;

const approvalResponseSchema = z.object({
  approved: z.boolean(),
});

function approvalContextBinding(context: ServerContext): string {
  const method = context.mcpReq?.method || 'tools/call';
  const auth = context.http?.authInfo;
  if (auth) {
    const tokenHash = createHash('sha256').update(auth.token).digest('hex');
    return `${method}\0auth:${auth.clientId}:${tokenHash}`;
  }

  const clientInfo = (context.mcpReq?.envelope as any)?.['io.modelcontextprotocol/clientInfo'];
  if (clientInfo?.name) {
    return `${method}\0client:${clientInfo.name}:${clientInfo.version ?? ''}`;
  }

  return `${method}\0session:${context.sessionId ?? 'local'}`;
}

function createDefaultApprovalStateCodec(): RequestStateCodec<ApprovalState> {
  const configuredSecret = process.env.MCP_APPROVAL_STATE_SECRET;
  return createRequestStateCodec<ApprovalState>({
    key: configuredSecret || randomBytes(32),
    ttlSeconds: APPROVAL_TTL_SECONDS,
    bind: approvalContextBinding,
  });
}

class MemoryApprovalReplayStore implements ApprovalReplayStore {
  private readonly claimed = new Map<string, number>();

  async claim(nonce: string, expiresAt: number): Promise<boolean> {
    const now = Date.now();
    for (const [key, expiry] of this.claimed) {
      if (expiry <= now) this.claimed.delete(key);
    }
    if (this.claimed.has(nonce)) return false;
    this.claimed.set(nonce, expiresAt);
    return true;
  }
}

class FileApprovalReplayStore implements ApprovalReplayStore {
  constructor(private readonly directory: string) {}

  async claim(nonce: string, expiresAt: number): Promise<boolean> {
    await mkdir(this.directory, { recursive: true, mode: 0o700 });
    await this.cleanupExpired();
    const markerPath = path.join(this.directory, `${nonce}.used`);
    try {
      const handle = await open(markerPath, 'wx', 0o600);
      try {
        await handle.writeFile(String(expiresAt));
      } finally {
        await handle.close();
      }
      return true;
    } catch (error: any) {
      if (error?.code === 'EEXIST') return false;
      throw error;
    }
  }

  private async cleanupExpired(): Promise<void> {
    const now = Date.now();
    let entries: string[];
    try {
      entries = await readdir(this.directory);
    } catch {
      return;
    }
    await Promise.all(
      entries
        .filter((entry) => entry.endsWith('.used'))
        .map(async (entry) => {
          const markerPath = path.join(this.directory, entry);
          try {
            const content = (await readFile(markerPath, 'utf8')).trim();
            const expiry = content ? Number(content) : Number.NaN;
            const marker = await stat(markerPath);
            if ((Number.isFinite(expiry) && expiry <= now) || marker.mtimeMs < now - 86_400_000) {
              await unlink(markerPath);
            }
          } catch {
            // Another replica may have removed the marker concurrently.
          }
        }),
    );
  }
}

function createDefaultReplayStore(): ApprovalReplayStore {
  const directory = process.env.MCP_APPROVAL_REPLAY_DIR?.trim();
  return directory ? new FileApprovalReplayStore(directory) : new MemoryApprovalReplayStore();
}

export function assertSensitiveApprovalDeploymentConfig(
  transport: 'stdio' | 'http',
  env: NodeJS.ProcessEnv = process.env,
): void {
  if (transport !== 'http') return;
  const secret = env.MCP_APPROVAL_STATE_SECRET;
  if (!secret || Buffer.byteLength(secret, 'utf8') < 32) {
    throw new Error('HTTP mode requires MCP_APPROVAL_STATE_SECRET with at least 32 bytes');
  }
  const replayDir = env.MCP_APPROVAL_REPLAY_DIR?.trim();
  if (!replayDir || !path.isAbsolute(replayDir)) {
    throw new Error(
      'HTTP mode requires MCP_APPROVAL_REPLAY_DIR as an absolute path on storage shared by every replica',
    );
  }
}

export function isGuardedToolName(toolName: string): toolName is GuardedToolName {
  return (GUARDED_TOOL_NAMES as readonly string[]).includes(normalizeQualifiedToolName(toolName));
}

export function normalizeQualifiedToolName(toolName: string): string {
  return normalizeToolName(toolName);
}

export function withSensitiveToolAnnotations(tool: Tool): Tool {
  const toolName = normalizeQualifiedToolName(tool.name);
  if (toolName === POD_EXEC_TOOL_NAME) {
    return {
      ...tool,
      annotations: {
        ...tool.annotations,
        title: 'Execute command in pod (approval required)',
        readOnlyHint: false,
        destructiveHint: true,
        idempotentHint: false,
        openWorldHint: true,
      },
    };
  }

  return tool;
}

export class SensitiveToolApprovalGuard {
  constructor(
    private readonly stateCodec: RequestStateCodec<ApprovalState> = createDefaultApprovalStateCodec(),
    private readonly replayStore: ApprovalReplayStore = createDefaultReplayStore(),
  ) {}

  public get requestStateVerifier() {
    return (state: string, context: ServerContext) => this.stateCodec.verify(state, context);
  }

  public async authorizeAndExecute(
    toolName: string,
    params: unknown,
    context: ServerContext | undefined,
    execute: () => Promise<unknown>,
  ): Promise<unknown | InputRequiredResult> {
    const normalizedToolName = normalizeQualifiedToolName(toolName);
    if (!isGuardedToolName(normalizedToolName)) {
      throw new Error(`Tool '${toolName}' is not configured for sensitive-tool approval`);
    }

    if (!context) {
      return this.denied(
        normalizedToolName,
        'Sensitive tools require an interactive MCP approval context.',
      );
    }

    const argumentsHash = hashArguments(params);
    const state = context.mcpReq.requestState<ApprovalState>();

    if (state === undefined) {
      const expiresAt = Date.now() + APPROVAL_TTL_SECONDS * 1000;
      const requestState = await this.stateCodec.mint(
        {
          kind: 'code-mode-sensitive-tool-approval',
          toolName: normalizedToolName,
          argumentsHash,
          nonce: randomBytes(16).toString('hex'),
          expiresAt,
        },
        context,
      );

      return inputRequired({
        inputRequests: {
          approval: inputRequired.elicit({
            message: buildApprovalMessage(normalizedToolName, params),
            requestedSchema: approvalResponseSchema,
          }),
        },
        requestState,
      });
    }

    if (!isMatchingApprovalState(state, normalizedToolName, argumentsHash)) {
      return this.denied(
        normalizedToolName,
        'Approval state does not match the requested tool and arguments.',
      );
    }

    if (!(await this.replayStore.claim(state.nonce, state.expiresAt))) {
      return this.denied(normalizedToolName, 'Approval state has already been used.');
    }

    const response = acceptedContent(
      context.mcpReq.inputResponses,
      'approval',
      approvalResponseSchema,
    );
    if (response?.approved !== true) {
      return this.denied(normalizedToolName, 'User approval was declined or cancelled.');
    }

    return execute();
  }

  private denied(toolName: GuardedToolName, reason: string): CallToolResult {
    return {
      isError: true,
      content: [
        {
          type: 'text',
          text: JSON.stringify({
            success: false,
            tool: toolName,
            approval: 'denied',
            error: reason,
          }),
        },
      ],
    };
  }
}

export const sensitiveToolApprovalGuard = new SensitiveToolApprovalGuard();

function isMatchingApprovalState(
  state: ApprovalState,
  toolName: GuardedToolName,
  argumentsHash: string,
): boolean {
  return (
    state !== null &&
    typeof state === 'object' &&
    state.kind === 'code-mode-sensitive-tool-approval' &&
    state.toolName === toolName &&
    state.argumentsHash === argumentsHash &&
    typeof state.nonce === 'string' &&
    /^[a-f0-9]{32}$/.test(state.nonce) &&
    typeof state.expiresAt === 'number' &&
    state.expiresAt >= Date.now()
  );
}

function hashArguments(params: unknown): string {
  const canonical = JSON.stringify(sortJsonValue(params ?? {}));
  return createHash('sha256').update(canonical).digest('hex');
}

function sortJsonValue(value: unknown): unknown {
  if (Array.isArray(value)) {
    return value.map(sortJsonValue);
  }
  if (value && typeof value === 'object') {
    return Object.fromEntries(
      Object.entries(value as Record<string, unknown>)
        .filter(([, entryValue]) => entryValue !== undefined)
        .sort(([left], [right]) => left.localeCompare(right))
        .map(([key, entryValue]) => [key, sortJsonValue(entryValue)]),
    );
  }
  return value;
}

function buildApprovalMessage(toolName: GuardedToolName, params: unknown): string {
  const input = params && typeof params === 'object' ? (params as Record<string, unknown>) : {};

  if (toolName === POD_EXEC_TOOL_NAME) {
    return [
      'Approve execution of a command inside a Kubernetes pod?',
      `Namespace: ${display(input.namespace, 'default')}`,
      `Pod: ${display(input.podName, '(missing)')}`,
      `Container: ${display(input.container, 'first container')}`,
      `Command: ${displayExecCommand(input)}`,
      `Stdin: ${typeof input.stdin === 'string' ? JSON.stringify(input.stdin) : '(none)'}`,
      `Timeout: ${display(input.timeoutSeconds, 60)} seconds`,
      'This command can modify container files or processes and may access workload credentials or networks.',
    ].join('\n');
  }

  throw new Error(`Unsupported guarded tool: ${toolName}`);
}

function displayExecCommand(input: Record<string, unknown>): string {
  if (Array.isArray(input.args) && input.args.length > 0) {
    return `direct argv ${JSON.stringify(input.args)}`;
  }
  if (typeof input.argv === 'string' && input.argv.trim()) {
    return `direct argv ${JSON.stringify(input.argv.trim().split(/\s+/))}`;
  }
  if (typeof input.command === 'string' && input.command) {
    return `shell ${display(input.shell, '/bin/sh')} -c ${JSON.stringify(input.command)}`;
  }
  return '(missing)';
}

function display(value: unknown, fallback: string | number): string {
  if (value === undefined || value === null || value === '') return String(fallback);
  return typeof value === 'string' ? value : JSON.stringify(value);
}