/**
 * Monitoring Transport
 *
 * Records OpenTelemetry spans/metrics and logs tool calls to database.
 * Tracks in-flight requests to correlate requests with responses.
 */

import type { StudioContext } from "@/core/studio-context";
import { trace, context, type Context } from "@opentelemetry/api";
import type { Span } from "@opentelemetry/api";
import type {
  JSONRPCMessage,
  JSONRPCRequest,
  JSONRPCResponse,
  CallToolRequest,
  CallToolResult,
} from "@modelcontextprotocol/sdk/types.js";
import { WrapperTransport } from "@decocms/mcp-utils";
import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
import { emitMonitoringLog } from "@/monitoring/emit";
import { recordToolExecutionMetrics } from "@/monitoring/record-tool-execution-metrics";
import {
  extractCallToolErrorMessage,
  extractMetaProperties,
  mergeProperties,
} from "@/api/routes/proxy-monitoring";

interface MonitoringTransportOptions {
  ctx: StudioContext;
  connectionId: string;
  virtualMcpId?: string;
}

interface InflightRequest {
  startTime: number;
  method: string;
  toolName?: string;
  toolArguments?: Record<string, unknown>;
  span?: Span;
}

export class MonitoringTransport extends WrapperTransport {
  private inflightRequests = new Map<string | number, InflightRequest>();
  private requestContext: Context;

  constructor(
    innerTransport: Transport,
    private options: MonitoringTransportOptions,
  ) {
    super(innerTransport);
    // Capture the active OTel context at construction time so spans created
    // during async message handling are correctly parented to the HTTP request.
    this.requestContext = context.active();
  }

  protected override async handleOutgoingMessage(
    message: JSONRPCMessage,
  ): Promise<void> {
    if (this.isRequest(message)) {
      const request = message as JSONRPCRequest;
      this.onRequestStart(request);
    }

    return this.innerTransport.send(message);
  }

  protected override handleIncomingMessage(message: JSONRPCMessage): void {
    if (this.isResponse(message)) {
      const response = message as JSONRPCResponse;
      this.onResponseEnd(response);
    }

    // Forward to client
    this.onmessage?.(message);
  }

  private onRequestStart(request: JSONRPCRequest): void {
    const { ctx, connectionId } = this.options;

    // Extract tool name and arguments for tools/call
    let toolName: string | undefined;
    let toolArguments: Record<string, unknown> | undefined;
    if (request.method === "tools/call" && request.params) {
      const params = request.params as CallToolRequest["params"];
      toolName = params.name;
      toolArguments = params.arguments as Record<string, unknown> | undefined;
    }

    // Start OpenTelemetry span for tool calls, explicitly parented to the
    // context captured at construction time to survive async context loss.
    let span: Span | undefined;
    if (request.method === "tools/call" && toolName) {
      span = ctx.tracer.startSpan(
        "mcp.proxy.callTool",
        {
          attributes: {
            "connection.id": connectionId,
            "tool.name": toolName,
            "request.id": ctx.metadata.requestId,
            "jsonrpc.id": request.id,
            "jsonrpc.method": request.method,
          },
        },
        this.requestContext,
      );
    }

    // Only track if request has an ID
    if (request.id !== null && request.id !== undefined) {
      // Store request info with span
      this.inflightRequests.set(request.id, {
        startTime: Date.now(),
        method: request.method,
        toolName,
        toolArguments,
        span,
      });
    }
  }

  private onResponseEnd(response: JSONRPCResponse): void {
    // Skip if response has no ID
    if (response.id === null || response.id === undefined) return;

    const requestInfo = this.inflightRequests.get(response.id);
    if (!requestInfo) return;

    const { ctx, connectionId } = this.options;
    const { startTime, method, toolName, toolArguments, span } = requestInfo;
    const duration = Date.now() - startTime;

    // Clean up
    this.inflightRequests.delete(response.id);

    // Only record metrics/logging for tool calls
    if (method !== "tools/call" || !toolName) {
      return;
    }

    const organizationId = ctx.organization?.id ?? "";
    const transportError = "error" in response;
    const result = transportError ? response.error : response.result;

    // Convert to CallToolResult format for logging
    const callToolResult: CallToolResult = transportError
      ? {
          content: [
            {
              type: "text",
              text: response.error?.message || "Unknown error",
            },
          ],
          isError: true,
        }
      : (result as CallToolResult);

    // A tool call fails two ways: a JSON-RPC transport error, or a successful
    // response whose CallToolResult carries isError:true — the normal way MCP
    // tools signal failure. Both must count; checking only "error" in response
    // undercounted real failures ~35x in production telemetry.
    const isError = transportError || callToolResult?.isError === true;
    // Read response.error only inside the aliased type guard so TS narrows the
    // JSONRPCResponse union; downstream branches use the derived string.
    const transportErrorMessage = transportError
      ? response.error?.message
      : undefined;
    const errorMessage = extractCallToolErrorMessage(callToolResult) || null;

    // Record connection-level OpenTelemetry metrics
    ctx.meter.createHistogram("connection.proxy.duration").record(duration, {
      "connection.id": connectionId,
      "organization.id": organizationId,
      "tool.name": toolName,
      status: isError ? "error" : "success",
    });

    if (isError) {
      ctx.meter.createCounter("connection.proxy.errors").add(1, {
        "connection.id": connectionId,
        "organization.id": organizationId,
        "tool.name": toolName,
        error: transportErrorMessage ?? errorMessage ?? undefined,
      });
    } else {
      ctx.meter.createCounter("connection.proxy.requests").add(1, {
        "connection.id": connectionId,
        "organization.id": organizationId,
        "tool.name": toolName,
        status: "success",
      });
    }

    if (organizationId) {
      recordToolExecutionMetrics({
        ctx,
        organizationId,
        connectionId,
        toolName,
        durationMs: duration,
        isError,
        errorType: transportError ? "RemoteError" : isError ? "ToolError" : "",
      });
    }

    // Enrich and end OpenTelemetry span
    if (span) {
      if (isError) {
        span.recordException(
          new Error(transportErrorMessage ?? errorMessage ?? "Tool error"),
        );
      }

      const metaProps = extractMetaProperties(toolArguments);
      const properties = mergeProperties(ctx.metadata.properties, metaProps);

      const spanCtx = trace.setSpan(context.active(), span);

      emitMonitoringLog(
        {
          organizationId: ctx.organization?.id ?? "",
          connectionId,
          toolName,
          toolArguments,
          result: callToolResult,
          duration,
          isError,
          errorMessage,
          userId: ctx.auth.user?.id || ctx.auth.apiKey?.userId || null,
          requestId: ctx.metadata.requestId,
          userAgent: ctx.metadata.userAgent || null,
          virtualMcpId: this.options.virtualMcpId || null,
          properties: properties || null,
        },
        spanCtx,
      );

      // Keep span.end() — still needed for tracing
      span.end();
    }
  }

  // Clean up any dangling spans on close
  override async close(): Promise<void> {
    // End all in-flight spans
    for (const info of this.inflightRequests.values()) {
      if (info.span) {
        info.span.setAttributes({ "transport.closed": true });
        info.span.end();
      }
    }
    this.inflightRequests.clear();

    return super.close();
  }
}
