#!/usr/bin/env node
import {McpServer} from '@modelcontextprotocol/sdk/server/mcp.js';
import {StdioServerTransport} from '@modelcontextprotocol/sdk/server/stdio.js';
import type {CallToolResult} from '@modelcontextprotocol/sdk/types.js';
import {spawn, type ChildProcess} from 'child_process';
import {randomUUID} from 'crypto';
import {z} from 'zod';

// Constants
const DEFAULT_TIMEOUT = 300000; // 5 minutes
const MAX_BUFFER_SIZE = 10000; // Max lines to buffer for background processes

// Helper functions
function jsonResult(data: any): CallToolResult {
  return {
    content: [{type: 'text', text: JSON.stringify(data)}],
  };
}

function errorResult(message: string, error: Error | string): CallToolResult {
  return {
    isError: true,
    content: [{type: 'text', text: `${message}: ${error}`}],
  };
}

function wrap<F extends (...args: any[]) => any>(fn: F): F {
  return ((...args: Parameters<F>) => {
    try {
      return fn(...args);
    } catch (err) {
      return errorResult('Failed to execute tool', err);
    }
  }) as F;
}

// ShellSession class - manages persistent shell sessions
class ShellSession {
  id: string;
  process: ChildProcess;
  cwd: string;
  env: Record<string, string>;
  createdAt: Date;
  private pendingCommand: {
    resolve: (value: {stdout: string; stderr: string; exitCode: number}) => void;
    reject: (error: Error) => void;
    stdout: string;
    stderr: string;
    timeout?: NodeJS.Timeout;
  } | null = null;

  constructor(id: string, options: {cwd?: string; env?: Record<string, string>}) {
    this.id = id;
    this.cwd = options.cwd || process.cwd();
    this.env = {...process.env, ...options.env} as Record<string, string>;
    this.createdAt = new Date();

    // Spawn persistent shell
    this.process = spawn('/bin/sh', ['-s'], {
      cwd: this.cwd,
      env: this.env,
      stdio: ['pipe', 'pipe', 'pipe'],
    });

    // Handle stdout
    this.process.stdout?.on('data', (data: Buffer) => {
      if (this.pendingCommand) {
        this.pendingCommand.stdout += data.toString();
      }
    });

    // Handle stderr
    this.process.stderr?.on('data', (data: Buffer) => {
      if (this.pendingCommand) {
        this.pendingCommand.stderr += data.toString();
      }
    });

    // Handle process exit
    this.process.on('exit', (code) => {
      if (this.pendingCommand) {
        const exitCode = code ?? -1;
        this.pendingCommand.resolve({
          stdout: this.pendingCommand.stdout,
          stderr: this.pendingCommand.stderr,
          exitCode,
        });
        this.pendingCommand = null;
      }
    });
  }

  async execute(
    command: string,
    timeout: number = DEFAULT_TIMEOUT
  ): Promise<{stdout: string; stderr: string; exitCode: number}> {
    return new Promise((resolve, reject) => {
      if (this.pendingCommand) {
        reject(new Error('Session busy - another command is running'));
        return;
      }

      // Create unique marker for command completion
      const marker = `__CMD_COMPLETE_${randomUUID()}__`;

      this.pendingCommand = {
        resolve: (result) => {
          if (timeoutHandle) clearTimeout(timeoutHandle);
          resolve(result);
        },
        reject: (error) => {
          if (timeoutHandle) clearTimeout(timeoutHandle);
          reject(error);
        },
        stdout: '',
        stderr: '',
      };

      // Set timeout
      const timeoutHandle = setTimeout(() => {
        if (this.pendingCommand) {
          const result = {
            stdout: this.pendingCommand.stdout,
            stderr: this.pendingCommand.stderr + '\n[Command timed out]',
            exitCode: 124, // Timeout exit code
          };
          this.pendingCommand.resolve(result);
          this.pendingCommand = null;
        }
      }, timeout);

      // Execute command and print marker when done
      const wrappedCommand = `${command}\necho "${marker}" >&1\necho "${marker}" >&2\n`;
      this.process.stdin?.write(wrappedCommand);

      // Watch for marker in output
      const checkComplete = () => {
        if (this.pendingCommand) {
          const hasStdoutMarker = this.pendingCommand.stdout.includes(marker);
          const hasStderrMarker = this.pendingCommand.stderr.includes(marker);

          if (hasStdoutMarker && hasStderrMarker) {
            // Remove markers from output
            const stdout = this.pendingCommand.stdout.replace(marker, '').trim();
            const stderr = this.pendingCommand.stderr.replace(marker, '').trim();

            this.pendingCommand.resolve({
              stdout,
              stderr,
              exitCode: 0,
            });
            this.pendingCommand = null;
          } else {
            setTimeout(checkComplete, 50);
          }
        }
      };

      setTimeout(checkComplete, 50);
    });
  }

  async kill(): Promise<void> {
    this.process.kill('SIGTERM');

    // Wait a bit for graceful shutdown
    await new Promise(resolve => setTimeout(resolve, 100));

    // Force kill if still running
    if (!this.process.killed) {
      this.process.kill('SIGKILL');
    }
  }
}

// BackgroundProcess class - manages detached background processes
class BackgroundProcess {
  id: string;
  command: string;
  process: ChildProcess;
  startedAt: Date;
  stdout: string[] = [];
  stderr: string[] = [];
  exitCode: number | null = null;
  maxBufferSize: number = MAX_BUFFER_SIZE;

  constructor(command: string, options: {cwd?: string; env?: Record<string, string>}) {
    this.id = randomUUID();
    this.command = command;
    this.startedAt = new Date();

    // Spawn detached process
    this.process = spawn('/bin/sh', ['-c', command], {
      cwd: options.cwd || process.cwd(),
      env: {...process.env, ...options.env},
      stdio: ['ignore', 'pipe', 'pipe'],
      detached: true,
    });

    // Buffer stdout (circular buffer)
    this.process.stdout?.on('data', (data: Buffer) => {
      const lines = data.toString().split('\n');
      this.stdout.push(...lines);
      if (this.stdout.length > this.maxBufferSize) {
        this.stdout = this.stdout.slice(-this.maxBufferSize);
      }
    });

    // Buffer stderr (circular buffer)
    this.process.stderr?.on('data', (data: Buffer) => {
      const lines = data.toString().split('\n');
      this.stderr.push(...lines);
      if (this.stderr.length > this.maxBufferSize) {
        this.stderr = this.stderr.slice(-this.maxBufferSize);
      }
    });

    // Track exit code
    this.process.on('exit', (code) => {
      this.exitCode = code;
    });

    // Unref to allow parent to exit
    this.process.unref();
  }

  isRunning(): boolean {
    return this.exitCode === null && !this.process.killed;
  }

  getOutput(): {stdout: string; stderr: string; exitCode: number | null} {
    return {
      stdout: this.stdout.join('\n'),
      stderr: this.stderr.join('\n'),
      exitCode: this.exitCode,
    };
  }

  async kill(): Promise<void> {
    if (!this.isRunning()) return;

    this.process.kill('SIGTERM');

    // Wait for graceful shutdown
    await new Promise(resolve => setTimeout(resolve, 5000));

    // Force kill if still running
    if (this.isRunning()) {
      this.process.kill('SIGKILL');
    }
  }
}

// Global registries
const sessions = new Map<string, ShellSession>();
const backgroundProcesses = new Map<string, BackgroundProcess>();

// Initialize MCP server
const server = new McpServer(
  {
    name: 'shell-mcp',
    version: '1.0.0',
  },
  {
    capabilities: {
      tools: {},
    },
  }
);

// Tool: execute
server.tool(
  'execute',
  'Execute a shell command in a persistent session (uses "default" session if session_id not provided) or as a background process',
  {
    command: z.string().describe('The shell command to execute'),
    session_id: z.string().optional().describe('Optional session ID to reuse or create a persistent shell session'),
    cwd: z.string().optional().describe('Working directory for the command'),
    env: z.record(z.string()).optional().describe('Environment variables to set'),
    timeout: z.number().optional().describe('Timeout in milliseconds (default: 300000 / 5 minutes)'),
    background: z.boolean().optional().describe('Run as detached background process'),
  },
  wrap(async ({command, session_id, cwd, env, timeout, background}) => {
    if (!command) {
      throw new Error('command is required');
    }

    // Background process
    if (background) {
      const proc = new BackgroundProcess(command, {cwd, env});
      backgroundProcesses.set(proc.id, proc);
      return jsonResult({
        process_id: proc.id,
        status: 'started',
        message: 'Background process started',
      });
    }

    // Use default session if no session_id provided
    const effectiveSessionId = session_id || 'default';
    let session = sessions.get(effectiveSessionId);
    if (!session) {
      session = new ShellSession(effectiveSessionId, {cwd, env});
      sessions.set(effectiveSessionId, session);
    }

    const result = await session.execute(command, timeout);
    return jsonResult({
      session_id: effectiveSessionId,
      ...result,
    });
  })
);

// Tool: list_sessions
server.tool(
  'list_sessions',
  'List all active sessions and background processes, optionally filtered',
  {
    id: z.string().optional().describe('Filter by session/process ID (fuzzy match)'),
    cwd: z.string().optional().describe('Filter sessions by working directory (fuzzy match)'),
    session: z.string().optional().describe('Filter by session ID (fuzzy match, alias for id)'),
    command: z.string().optional().describe('Filter background processes by command (fuzzy match)'),
  },
  wrap(async ({id, cwd, session, command}) => {
    // Helper for fuzzy matching (strip to alphanumerics and check contains)
    const fuzzyMatch = (value: string, filter: string): boolean => {
      const clean = (s: string) => s.toLowerCase().replace(/[^a-z0-9]/g, '');
      return clean(value).includes(clean(filter));
    };

    // Use either 'id' or 'session' parameter for ID filtering
    const idFilter = id || session;

    // Filter sessions
    let filteredSessions = [...sessions.entries()].map(([id, session]) => ({
      id,
      cwd: session.cwd,
      createdAt: session.createdAt.toISOString(),
    }));

    if (idFilter) {
      filteredSessions = filteredSessions.filter(s => fuzzyMatch(s.id, idFilter));
    }
    if (cwd) {
      filteredSessions = filteredSessions.filter(s => fuzzyMatch(s.cwd, cwd));
    }

    // Filter background processes
    let filteredProcesses = [...backgroundProcesses.entries()].map(([id, proc]) => ({
      id,
      command: proc.command,
      status: proc.isRunning() ? 'running' : 'completed',
      exitCode: proc.exitCode,
      startedAt: proc.startedAt.toISOString(),
    }));

    if (idFilter) {
      filteredProcesses = filteredProcesses.filter(p => fuzzyMatch(p.id, idFilter));
    }
    if (command) {
      filteredProcesses = filteredProcesses.filter(p => fuzzyMatch(p.command, command));
    }

    return jsonResult({
      sessions: filteredSessions,
      backgroundProcesses: filteredProcesses,
    });
  })
);

// Tool: kill_session
server.tool(
  'kill_session',
  'Kill a session or background process by ID',
  {
    id: z.string().describe('Session ID or background process ID to kill'),
  },
  wrap(async ({id}) => {
    if (!id) {
      throw new Error('id is required');
    }

    const session = sessions.get(id);
    const proc = backgroundProcesses.get(id);

    if (session) {
      await session.kill();
      sessions.delete(id);
      return jsonResult({success: true, message: `Session ${id} killed`});
    } else if (proc) {
      await proc.kill();
      backgroundProcesses.delete(id);
      return jsonResult({success: true, message: `Process ${id} killed`});
    } else {
      throw new Error(`No session or process found with id: ${id}`);
    }
  })
);

// Cleanup handler
async function cleanup() {
  console.error('Shutting down...');

  for (const session of sessions.values()) {
    await session.kill();
  }

  for (const proc of backgroundProcesses.values()) {
    await proc.kill();
  }

  process.exit(0);
}

process.on('SIGINT', cleanup);
process.on('SIGTERM', cleanup);

// Start server
(async () => {
  const transport = new StdioServerTransport();
  await server.connect(transport);
  console.error('Shell MCP Server running on stdio');
})();