import type { McpServer, McpInstance } from '@prisma/client';
import type { IMcpInstanceRepository, IMcpServerRepository } from '../repositories/interfaces.js';
import type { McpOrchestrator } from './orchestrator.js';

export interface HealthCheckSpec {
  tool: string;
  arguments?: Record<string, unknown>;
  intervalSeconds?: number;
  timeoutSeconds?: number;
  failureThreshold?: number;
}

export interface ProbeResult {
  healthy: boolean;
  latencyMs: number;
  message: string;
}

interface ProbeState {
  consecutiveFailures: number;
  lastProbeAt: number;
}

/**
 * Periodic health probe runner — calls MCP tools on running instances to verify
 * they are alive and responsive. Mirrors Kubernetes liveness probe semantics.
 *
 * For STDIO servers: runs `docker exec` with a disposable MCP client script
 * that sends initialize + tool/call via the package binary.
 *
 * For SSE/HTTP servers: sends HTTP JSON-RPC directly to the container port.
 */
export class HealthProbeRunner {
  private probeStates = new Map<string, ProbeState>();
  private timer: ReturnType<typeof setInterval> | null = null;

  constructor(
    private instanceRepo: IMcpInstanceRepository,
    private serverRepo: IMcpServerRepository,
    private orchestrator: McpOrchestrator,
    private logger?: { info: (msg: string) => void; error: (obj: unknown, msg: string) => void },
  ) {}

  /** Start the periodic probe loop. Runs every `tickIntervalMs` (default 15s). */
  start(tickIntervalMs = 15_000): void {
    if (this.timer) return;
    this.timer = setInterval(() => {
      this.tick().catch((err) => {
        this.logger?.error({ err }, 'Health probe tick failed');
      });
    }, tickIntervalMs);
  }

  stop(): void {
    if (this.timer) {
      clearInterval(this.timer);
      this.timer = null;
    }
  }

  /** Single tick: probe all RUNNING instances that have healthCheck configs and are due. */
  async tick(): Promise<void> {
    const instances = await this.instanceRepo.findAll();
    const running = instances.filter((i) => i.status === 'RUNNING' && i.containerId);

    // Cache servers by ID to avoid repeated lookups
    const serverCache = new Map<string, McpServer>();

    for (const inst of running) {
      let server = serverCache.get(inst.serverId);
      if (!server) {
        const s = await this.serverRepo.findById(inst.serverId);
        if (!s) continue;
        serverCache.set(inst.serverId, s);
        server = s;
      }

      const healthCheck = server.healthCheck as HealthCheckSpec | null;
      if (!healthCheck) continue;

      const intervalMs = (healthCheck.intervalSeconds ?? 60) * 1000;
      const state = this.probeStates.get(inst.id);
      const now = Date.now();

      // Skip if not due yet
      if (state && (now - state.lastProbeAt) < intervalMs) continue;

      await this.probeInstance(inst, server, healthCheck);
    }

    // Clean up states for instances that no longer exist
    const activeIds = new Set(running.map((i) => i.id));
    for (const key of this.probeStates.keys()) {
      if (!activeIds.has(key)) {
        this.probeStates.delete(key);
      }
    }
  }

  /** Probe a single instance and update its health status. */
  async probeInstance(
    instance: McpInstance,
    server: McpServer,
    healthCheck: HealthCheckSpec,
  ): Promise<ProbeResult> {
    const timeoutMs = (healthCheck.timeoutSeconds ?? 10) * 1000;
    const failureThreshold = healthCheck.failureThreshold ?? 3;
    const now = new Date();
    const start = Date.now();

    let result: ProbeResult;

    try {
      if (server.transport === 'SSE' || server.transport === 'STREAMABLE_HTTP') {
        result = await this.probeHttp(instance, server, healthCheck, timeoutMs);
      } else {
        result = await this.probeStdio(instance, server, healthCheck, timeoutMs);
      }
    } catch (err) {
      result = {
        healthy: false,
        latencyMs: Date.now() - start,
        message: err instanceof Error ? err.message : String(err),
      };
    }

    // Update probe state
    const state = this.probeStates.get(instance.id) ?? { consecutiveFailures: 0, lastProbeAt: 0 };
    state.lastProbeAt = Date.now();

    if (result.healthy) {
      state.consecutiveFailures = 0;
    } else {
      state.consecutiveFailures++;
    }
    this.probeStates.set(instance.id, state);

    // Determine health status
    const healthStatus = result.healthy
      ? 'healthy'
      : state.consecutiveFailures >= failureThreshold
        ? 'unhealthy'
        : 'degraded';

    // Build event
    const eventType = result.healthy ? 'Normal' : 'Warning';
    const eventMessage = result.healthy
      ? `Health check passed (${result.latencyMs}ms)`
      : `Health check failed: ${result.message}`;

    const existingEvents = (instance.events as Array<{ timestamp: string; type: string; message: string }>) ?? [];
    // Keep last 50 events
    const events = [
      ...existingEvents.slice(-49),
      { timestamp: now.toISOString(), type: eventType, message: eventMessage },
    ];

    // Update instance
    await this.instanceRepo.updateStatus(instance.id, instance.status as 'RUNNING', {
      healthStatus,
      lastHealthCheck: now,
      events,
    });

    this.logger?.info(
      `[health] ${(instance as unknown as { server?: { name: string } }).server?.name ?? instance.serverId}: ${healthStatus} (${result.latencyMs}ms) - ${eventMessage}`,
    );

    return result;
  }

  /** Probe an HTTP/SSE MCP server by sending a JSON-RPC tool call. */
  private async probeHttp(
    instance: McpInstance,
    server: McpServer,
    healthCheck: HealthCheckSpec,
    timeoutMs: number,
  ): Promise<ProbeResult> {
    if (!instance.containerId) {
      return { healthy: false, latencyMs: 0, message: 'No container ID' };
    }

    // Get container IP for internal network communication
    // (mcpd and MCP containers share the mcp-servers network)
    const containerInfo = await this.orchestrator.inspectContainer(instance.containerId);
    const containerPort = (server.containerPort as number | null) ?? 3000;

    let baseUrl: string;
    if (containerInfo.ip) {
      baseUrl = `http://${containerInfo.ip}:${containerPort}`;
    } else if (instance.port) {
      baseUrl = `http://localhost:${instance.port}`;
    } else {
      return { healthy: false, latencyMs: 0, message: 'No container IP or port' };
    }

    if (server.transport === 'SSE') {
      return this.probeSse(baseUrl, healthCheck, timeoutMs);
    }
    return this.probeStreamableHttp(baseUrl, healthCheck, timeoutMs);
  }

  /**
   * Probe a streamable-http MCP server (POST to root endpoint).
   */
  private async probeStreamableHttp(
    baseUrl: string,
    healthCheck: HealthCheckSpec,
    timeoutMs: number,
  ): Promise<ProbeResult> {
    const start = Date.now();
    const controller = new AbortController();
    const timer = setTimeout(() => controller.abort(), timeoutMs);

    try {
      const initResp = await fetch(baseUrl, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json', 'Accept': 'application/json, text/event-stream' },
        body: JSON.stringify({
          jsonrpc: '2.0', id: 1, method: 'initialize',
          params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'mcpctl-health', version: '0.0.1' } },
        }),
        signal: controller.signal,
      });

      if (!initResp.ok) {
        return { healthy: false, latencyMs: Date.now() - start, message: `Initialize HTTP ${initResp.status}` };
      }

      const sessionId = initResp.headers.get('mcp-session-id');
      const headers: Record<string, string> = { 'Content-Type': 'application/json', 'Accept': 'application/json, text/event-stream' };
      if (sessionId) headers['Mcp-Session-Id'] = sessionId;

      await fetch(baseUrl, {
        method: 'POST', headers,
        body: JSON.stringify({ jsonrpc: '2.0', method: 'notifications/initialized' }),
        signal: controller.signal,
      });

      const toolResp = await fetch(baseUrl, {
        method: 'POST', headers,
        body: JSON.stringify({
          jsonrpc: '2.0', id: 2, method: 'tools/call',
          params: { name: healthCheck.tool, arguments: healthCheck.arguments ?? {} },
        }),
        signal: controller.signal,
      });

      const latencyMs = Date.now() - start;

      if (!toolResp.ok) {
        return { healthy: false, latencyMs, message: `Tool call HTTP ${toolResp.status}` };
      }

      const body = await toolResp.text();
      try {
        const parsed = JSON.parse(body.includes('data: ') ? body.split('data: ')[1]!.split('\n')[0]! : body);
        if (parsed.error) {
          return { healthy: false, latencyMs, message: parsed.error.message ?? 'Tool call error' };
        }
      } catch {
        // If parsing fails but HTTP was ok, consider it healthy
      }

      return { healthy: true, latencyMs, message: 'ok' };
    } finally {
      clearTimeout(timer);
    }
  }

  /**
   * Probe an SSE-transport MCP server.
   * SSE protocol: GET /sse → endpoint event → POST /messages?session_id=...
   */
  private async probeSse(
    baseUrl: string,
    healthCheck: HealthCheckSpec,
    timeoutMs: number,
  ): Promise<ProbeResult> {
    const start = Date.now();
    const controller = new AbortController();
    const timer = setTimeout(() => controller.abort(), timeoutMs);

    try {
      // 1. Connect to SSE endpoint to get the message URL
      const sseResp = await fetch(`${baseUrl}/sse`, {
        method: 'GET',
        headers: { 'Accept': 'text/event-stream' },
        signal: controller.signal,
      });

      if (!sseResp.ok) {
        return { healthy: false, latencyMs: Date.now() - start, message: `SSE connect HTTP ${sseResp.status}` };
      }

      // 2. Read the SSE stream to find the endpoint event
      const reader = sseResp.body?.getReader();
      if (!reader) {
        return { healthy: false, latencyMs: Date.now() - start, message: 'No SSE stream body' };
      }

      const decoder = new TextDecoder();
      let buffer = '';
      let messagesUrl = '';

      // Read until we get the endpoint event
      while (!messagesUrl) {
        const { done, value } = await reader.read();
        if (done) break;
        buffer += decoder.decode(value, { stream: true });

        for (const line of buffer.split('\n')) {
          if (line.startsWith('data: ') && buffer.includes('event: endpoint')) {
            const endpoint = line.slice(6).trim();
            // Endpoint may be relative (e.g., /messages?session_id=...) or absolute
            messagesUrl = endpoint.startsWith('http') ? endpoint : `${baseUrl}${endpoint}`;
          }
        }
        // Keep only the last incomplete line
        const lines = buffer.split('\n');
        buffer = lines[lines.length - 1] ?? '';
      }

      if (!messagesUrl) {
        reader.cancel();
        return { healthy: false, latencyMs: Date.now() - start, message: 'No endpoint event from SSE' };
      }

      // 3. Initialize via the messages endpoint
      const postHeaders = { 'Content-Type': 'application/json' };

      const initResp = await fetch(messagesUrl, {
        method: 'POST', headers: postHeaders,
        body: JSON.stringify({
          jsonrpc: '2.0', id: 1, method: 'initialize',
          params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'mcpctl-health', version: '0.0.1' } },
        }),
        signal: controller.signal,
      });

      if (!initResp.ok) {
        reader.cancel();
        return { healthy: false, latencyMs: Date.now() - start, message: `Initialize HTTP ${initResp.status}` };
      }

      // 4. Send initialized notification
      await fetch(messagesUrl, {
        method: 'POST', headers: postHeaders,
        body: JSON.stringify({ jsonrpc: '2.0', method: 'notifications/initialized' }),
        signal: controller.signal,
      });

      // 5. Call health check tool
      const toolResp = await fetch(messagesUrl, {
        method: 'POST', headers: postHeaders,
        body: JSON.stringify({
          jsonrpc: '2.0', id: 2, method: 'tools/call',
          params: { name: healthCheck.tool, arguments: healthCheck.arguments ?? {} },
        }),
        signal: controller.signal,
      });

      const latencyMs = Date.now() - start;

      // 6. Read tool response from SSE stream
      // The response comes back on the SSE stream, not the POST response
      let responseBuffer = '';
      const readTimeout = setTimeout(() => reader.cancel(), 5000);

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        responseBuffer += decoder.decode(value, { stream: true });

        // Look for data lines containing our response (id: 2)
        for (const line of responseBuffer.split('\n')) {
          if (line.startsWith('data: ')) {
            try {
              const parsed = JSON.parse(line.slice(6));
              if (parsed.id === 2) {
                clearTimeout(readTimeout);
                reader.cancel();
                if (parsed.error) {
                  return { healthy: false, latencyMs, message: parsed.error.message ?? 'Tool call error' };
                }
                return { healthy: true, latencyMs, message: 'ok' };
              }
            } catch {
              // Not valid JSON, skip
            }
          }
        }
        const respLines = responseBuffer.split('\n');
        responseBuffer = respLines[respLines.length - 1] ?? '';
      }

      clearTimeout(readTimeout);
      reader.cancel();

      // If POST response itself was ok (202 for SSE), consider it healthy
      if (toolResp.ok) {
        return { healthy: true, latencyMs, message: 'ok' };
      }

      return { healthy: false, latencyMs, message: `Tool call HTTP ${toolResp.status}` };
    } finally {
      clearTimeout(timer);
    }
  }

  /**
   * Probe a STDIO MCP server by running `docker exec` with a disposable Node.js
   * script that pipes JSON-RPC messages into the package binary.
   */
  private async probeStdio(
    instance: McpInstance,
    server: McpServer,
    healthCheck: HealthCheckSpec,
    timeoutMs: number,
  ): Promise<ProbeResult> {
    if (!instance.containerId) {
      return { healthy: false, latencyMs: 0, message: 'No container ID' };
    }

    const start = Date.now();
    const packageName = server.packageName as string | null;
    const command = server.command as string[] | null;

    // Determine how to spawn the MCP server inside the container
    let spawnCmd: string[];
    if (packageName) {
      spawnCmd = ['npx', '--prefer-offline', '-y', packageName];
    } else if (command && command.length > 0) {
      spawnCmd = command;
    } else {
      return { healthy: false, latencyMs: 0, message: 'No packageName or command for STDIO server' };
    }

    // Build JSON-RPC messages for the health probe
    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: {
        protocolVersion: '2024-11-05',
        capabilities: {},
        clientInfo: { name: 'mcpctl-health', version: '0.0.1' },
      },
    });
    const initializedMsg = JSON.stringify({
      jsonrpc: '2.0', method: 'notifications/initialized',
    });
    const toolCallMsg = JSON.stringify({
      jsonrpc: '2.0', id: 2, method: 'tools/call',
      params: { name: healthCheck.tool, arguments: healthCheck.arguments ?? {} },
    });

    // Use a Node.js inline script that:
    // 1. Spawns the MCP server binary
    // 2. Sends initialize + initialized + tool call via stdin
    // 3. Reads responses from stdout
    // 4. Exits with 0 if tool call succeeds, 1 if it fails
    const spawnArgs = JSON.stringify(spawnCmd);
    const probeScript = `
const { spawn } = require('child_process');
const args = ${spawnArgs};
const proc = spawn(args[0], args.slice(1), { stdio: ['pipe', 'pipe', 'pipe'] });
let output = '';
let responded = false;
proc.stdout.on('data', d => {
  output += d;
  const lines = output.split('\\n');
  for (const line of lines) {
    if (!line.trim()) continue;
    try {
      const msg = JSON.parse(line);
      if (msg.id === 2) {
        responded = true;
        if (msg.error) {
          process.stdout.write('ERROR:' + (msg.error.message || 'unknown'));
          proc.kill();
          process.exit(1);
        } else {
          process.stdout.write('OK');
          proc.kill();
          process.exit(0);
        }
      }
    } catch {}
  }
  output = lines[lines.length - 1] || '';
});
proc.stderr.on('data', () => {});
proc.on('error', e => { process.stdout.write('ERROR:' + e.message); process.exit(1); });
proc.on('exit', (code) => { if (!responded) { process.stdout.write('ERROR:process exited ' + code); process.exit(1); } });
setTimeout(() => { if (!responded) { process.stdout.write('ERROR:timeout'); proc.kill(); process.exit(1); } }, ${timeoutMs - 2000});
proc.stdin.write(${JSON.stringify(initMsg)} + '\\n');
setTimeout(() => {
  proc.stdin.write(${JSON.stringify(initializedMsg)} + '\\n');
  setTimeout(() => {
    proc.stdin.write(${JSON.stringify(toolCallMsg)} + '\\n');
  }, 500);
}, 500);
`.trim();

    try {
      const result = await this.orchestrator.execInContainer(
        instance.containerId,
        ['node', '-e', probeScript],
        { timeoutMs },
      );

      const latencyMs = Date.now() - start;

      if (result.exitCode === 0 && result.stdout.includes('OK')) {
        return { healthy: true, latencyMs, message: 'ok' };
      }

      // Extract error message
      const errorMatch = result.stdout.match(/ERROR:(.*)/);
      const errorMsg = errorMatch?.[1] ?? (result.stderr.trim() || `exit code ${result.exitCode}`);
      return { healthy: false, latencyMs, message: errorMsg };
    } catch (err) {
      return {
        healthy: false,
        latencyMs: Date.now() - start,
        message: err instanceof Error ? err.message : String(err),
      };
    }
  }
}
