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

export interface HealthCheckSpec {
  /** When set, probe sends initialize + tools/call (readiness). When omitted, probe sends tools/list only (liveness). */
  tool?: string;
  arguments?: Record<string, unknown>;
  intervalSeconds?: number;
  timeoutSeconds?: number;
  failureThreshold?: number;
}

/** Default liveness probe applied to any RUNNING instance whose server has no explicit healthCheck. */
export const DEFAULT_HEALTH_CHECK: HealthCheckSpec = {
  intervalSeconds: 30,
  timeoutSeconds: 8,
  failureThreshold: 3,
};

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 },
    /** Used for liveness probes (no explicit tool) — routes tools/list through the real production path. */
    private mcpProxyService?: McpProxyService,
  ) {}

  /** 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;
      }

      // Any server without an explicit healthCheck gets the default liveness probe.
      const healthCheck: HealthCheckSpec = (server.healthCheck as HealthCheckSpec | null) ?? DEFAULT_HEALTH_CHECK;

      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 (healthCheck.tool === undefined) {
        // Liveness probe: send tools/list through the real production path.
        // Mirrors exactly what mcplocal/client calls do, so synthetic and real
        // failures converge on the same signal.
        result = await this.probeLiveness(server, timeoutMs);
      } else {
        const readinessCheck = healthCheck as HealthCheckSpec & { tool: string };
        if (server.transport === 'STDIO') {
          // Route STDIO readiness through the proxy so probes hit the live
          // running container rather than spawning a fresh process inside
          // it. The legacy `probeStdio` (docker-exec a synthetic Node script
          // that re-spawns the package binary) only worked for
          // packageName-based servers — image-based STDIO servers (gitea,
          // docmost) returned a fake-unhealthy "No packageName or command"
          // before they even tried the tool. Going through mcpProxyService
          // also means readiness failures match production failures exactly.
          result = await this.probeReadinessViaProxy(server, readinessCheck, timeoutMs);
        } else {
          result = await this.probeHttp(instance, server, readinessCheck, 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;
  }

  /**
   * Readiness probe via McpProxyService — sends `tools/call` against the
   * configured probe tool through the live running instance. Used by
   * STDIO servers; HTTP/SSE servers go through the bespoke `probeHttp`
   * paths that connect directly to the container's IP+port (those work
   * fine and are kept as-is to minimise the diff in this PR).
   *
   * If the tool returns a JSON-RPC `error` (e.g. gitea-mcp-server's
   * "token is required" when GITEA_ACCESS_TOKEN didn't resolve), we mark
   * the instance unhealthy with the upstream error message. That's how
   * we catch broken-by-empty-secret cases that liveness (`tools/list`)
   * would otherwise pass.
   */
  private async probeReadinessViaProxy(
    server: McpServer,
    healthCheck: HealthCheckSpec & { tool: string },
    timeoutMs: number,
  ): Promise<ProbeResult> {
    const start = Date.now();
    if (!this.mcpProxyService) {
      return { healthy: false, latencyMs: 0, message: 'mcpProxyService not wired — cannot run readiness probe' };
    }

    const deadline = new Promise<ProbeResult>((resolve) => {
      setTimeout(() => resolve({
        healthy: false,
        latencyMs: timeoutMs,
        message: `Readiness probe timed out after ${timeoutMs}ms`,
      }), timeoutMs);
    });

    const probe = this.mcpProxyService
      .execute({
        serverId: server.id,
        method: 'tools/call',
        params: { name: healthCheck.tool, arguments: healthCheck.arguments ?? {} },
      })
      .then((response): ProbeResult => {
        const latencyMs = Date.now() - start;
        if (response.error) {
          return {
            healthy: false,
            latencyMs,
            message: response.error.message ?? `tools/call ${healthCheck.tool} returned error`,
          };
        }
        // Some servers report tool-level failures inside the result body
        // (`{ isError: true, content: [...] }`) rather than as JSON-RPC
        // errors. Treat that as unhealthy too.
        const result = response.result as { isError?: boolean; content?: Array<{ text?: string }> } | undefined;
        if (result?.isError) {
          const text = result.content?.[0]?.text ?? `${healthCheck.tool} returned isError`;
          return { healthy: false, latencyMs, message: text };
        }
        return { healthy: true, latencyMs, message: 'ok' };
      })
      .catch((err: unknown): ProbeResult => ({
        healthy: false,
        latencyMs: Date.now() - start,
        message: err instanceof Error ? err.message : String(err),
      }));

    return Promise.race([probe, deadline]);
  }

  /**
   * Liveness probe — sends tools/list via McpProxyService so the probe traverses
   * the exact code path production clients use. Works uniformly across every
   * transport (STDIO exec/attach, SSE, Streamable HTTP, external).
   */
  private async probeLiveness(server: McpServer, timeoutMs: number): Promise<ProbeResult> {
    const start = Date.now();
    if (!this.mcpProxyService) {
      return { healthy: false, latencyMs: 0, message: 'mcpProxyService not wired — cannot run default liveness probe' };
    }

    const deadline = new Promise<ProbeResult>((resolve) => {
      setTimeout(() => resolve({
        healthy: false,
        latencyMs: timeoutMs,
        message: `Liveness probe timed out after ${timeoutMs}ms`,
      }), timeoutMs);
    });

    const probe = this.mcpProxyService.execute({ serverId: server.id, method: 'tools/list' })
      .then((response): ProbeResult => {
        const latencyMs = Date.now() - start;
        if (response.error) {
          return { healthy: false, latencyMs, message: response.error.message ?? 'tools/list error' };
        }
        return { healthy: true, latencyMs, message: 'ok' };
      })
      .catch((err: unknown): ProbeResult => ({
        healthy: false,
        latencyMs: Date.now() - start,
        message: err instanceof Error ? err.message : String(err),
      }));

    return Promise.race([probe, deadline]);
  }

  /** Probe an HTTP/SSE MCP server by sending a JSON-RPC tool call. */
  private async probeHttp(
    instance: McpInstance,
    server: McpServer,
    healthCheck: HealthCheckSpec & { tool: string },
    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 & { tool: string },
    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 & { tool: string },
    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);
    }
  }

  // Note: a previous `probeStdio` implementation existed here that ran a
  // disposable Node script inside the container via `docker exec`,
  // re-spawning the package binary and piping JSON-RPC into it. It only
  // worked for packageName-based servers (the spawn step required an
  // npx-compatible package); image-based STDIO servers like
  // gitea-mcp-server fell through with "No packageName or command" and
  // were always reported unhealthy for the wrong reason. STDIO readiness
  // now goes through `probeReadinessViaProxy` which calls the live
  // running container — same code path as production traffic — and
  // surfaces the upstream error verbatim.
}
