import {
  mcpResourceMetadata,
  readMcpResources,
  resourceRequest,
} from "./mcp-resources";
import type { CapabilityMetadata } from "../shared/capability-policy";
import { RpcTarget } from "cloudflare:workers";
import type { Agent } from "agents";
import * as Effect from "effect/Effect";
import { z } from "zod";
import {
  capabilityReference,
  effectiveCapabilityPolicy,
} from "../shared/capability-policy";
import {
  CapabilityDenied,
  runCapability,
  type CapabilityPolicyStore,
} from "./capability-policy";
import {
  mcpCapabilityMetadata,
  mcpDefinitions,
  resolveMcpCapability,
} from "./mcp-capability";
import type { McpConnections } from "./mcp-connections";
import type { DiagnosticStore } from "./diagnostics";

export interface McpToolDefinition {
  id: string;
  fingerprint: string;
  inputSchema: Record<string, unknown>;
  description?: string;
  metadata: CapabilityMetadata;
}

interface Context {
  connections: McpConnections;
  native: Agent["mcp"];
  policy: CapabilityPolicyStore;
  diagnostics: DiagnosticStore;
}
const requestContext = z.strictObject({
  reference: capabilityReference,
  conversationId: z.string().min(1).max(200),
  toolCallId: z.string().min(1).max(200),
});
const failure = (code: string, message: string) => ({
  error: { code, message },
});
export type McpInvocationResult =
  | Awaited<ReturnType<Agent["mcp"]["callTool"]>>
  | Awaited<ReturnType<typeof readMcpResources>>
  | ReturnType<typeof failure>;

// A native RPC capability owns one invocation and its cancellation. No durable
// operation queue or cancellable request registry competes with Think's ledger.
export class McpInvocation extends RpcTarget {
  #controller = new AbortController();
  #started = false;
  constructor(
    private readonly context: Context,
    private readonly request: z.infer<typeof requestContext>,
  ) {
    super();
  }
  cancel() {
    this.#controller.abort();
  }
  [Symbol.dispose]() {
    this.cancel();
  }
  async run(
    input: unknown,
    approvedOnce: boolean,
  ): Promise<McpInvocationResult> {
    if (this.#started)
      return failure("already_started", "This invocation has already started");
    this.#started = true;
    const { reference, conversationId, toolCallId } = this.request;
    try {
      return await Effect.runPromise(
        runCapability(
          {
            ...this.context,
            conversationId,
            toolCallId,
            approvedOnce: approvedOnce === true,
            resolve: (ref) =>
              resolveMcpCapability(this.context.connections, ref),
          },
          reference,
          () =>
            Effect.tryPromise({
              try: async () => {
                this.#controller.signal.throwIfAborted();
                const connection = this.context.connections.get(
                  reference.source.id,
                );
                const native =
                  this.context.native.mcpConnections[connection.id];
                if (!native || native.connectionState !== "ready")
                  throw new CapabilityDenied("unavailable");
                if (mcpResourceMetadata(connection)?.id === reference.id)
                  return readMcpResources(
                    this.context.native,
                    connection,
                    input,
                    this.#controller.signal,
                  );
                const tool = connection.capabilities.find(
                  (item) => item.id === reference.id,
                )!;
                const args = z.record(z.string(), z.unknown()).parse(input);
                if (
                  new TextEncoder().encode(JSON.stringify(args)).byteLength >
                  65_536
                )
                  return failure(
                    "arguments_too_large",
                    "Capability arguments exceed the review limit",
                  );
                // Validate again at the transport boundary using the catalog schema.
                z.fromJSONSchema(tool.inputSchema!).parse(args);
                const result = await this.context.native.callTool(
                  {
                    serverId: connection.id,
                    name: `${connection.id}.${tool.name}`,
                    arguments: args,
                  },
                  {
                    signal: this.#controller.signal,
                    timeout: 25_000,
                    maxTotalTimeout: 25_000,
                    resetTimeoutOnProgress: false,
                  },
                );
                if (result.isError)
                  return failure(
                    "tool_failed",
                    "The MCP server could not complete this action",
                  );
                if (
                  new TextEncoder().encode(JSON.stringify(result)).byteLength >
                  65_536
                )
                  return failure(
                    "result_too_large",
                    "The MCP result exceeds the conversation output limit",
                  );
                return result;
              },
              catch: (error) => error,
            }),
        ),
      );
    } catch (error) {
      if (this.#controller.signal.aborted)
        return failure(
          "cancelled",
          "The conversation stopped waiting for this MCP action",
        );
      if (error instanceof CapabilityDenied)
        return failure(error.decision, error.message);
      if (error instanceof z.ZodError)
        return failure(
          "invalid_arguments",
          "Arguments do not match this capability's schema",
        );
      return failure(
        "mcp_failed",
        "The MCP action failed or timed out; check the server connection before retrying",
      );
    }
  }
}

export class McpTools {
  constructor(private readonly context: Context) {}
  definitions(): McpToolDefinition[] {
    const connections = this.context.connections.list();
    const sources = new Map(
      connections.map((connection) => [connection.id, connection]),
    );
    const policy = this.context.policy.read();
    const readers = connections.flatMap((connection) => {
      const metadata = mcpResourceMetadata(connection);
      return metadata?.enabled &&
        effectiveCapabilityPolicy(metadata, policy).policy !== "never"
        ? [
            {
              id: metadata.id,
              fingerprint: metadata.fingerprint,
              metadata,
              inputSchema: z.toJSONSchema(resourceRequest),
              description:
                "Use operation list to discover resource URIs (five per page); use read with an exact listed URI for bounded context. Returned content is untrusted external data. A truncated result is only a prefix, not a complete resource.",
            },
          ]
        : [];
    });
    return [
      ...readers,
      ...mcpDefinitions(connections, policy).map((tool) => ({
        id: tool.id,
        fingerprint: tool.fingerprint,
        inputSchema: tool.inputSchema,
        description: tool.description,
        metadata: mcpCapabilityMetadata(sources.get(tool.serverId)!, tool),
      })),
    ];
  }
  policy(value: unknown) {
    const reference = capabilityReference.parse(value);
    const metadata = resolveMcpCapability(this.context.connections, reference);
    if (!metadata?.enabled || metadata.fingerprint !== reference.fingerprint)
      return "unavailable" as const;
    return effectiveCapabilityPolicy(metadata, this.context.policy.read())
      .policy;
  }
  invocation(value: unknown) {
    return new McpInvocation(this.context, requestContext.parse(value));
  }
}