import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import type { Tool } from "@modelcontextprotocol/sdk/types.js";
import { HarnessError } from "./errors.js";

export interface GatewayConnection {
  list(signal: AbortSignal): Promise<Tool[]>;
  call(name: string, args: Record<string, unknown>, signal: AbortSignal): Promise<unknown>;
  close(): Promise<void>;
}
export type GatewayConnector = (signal: AbortSignal) => Promise<GatewayConnection>;
export interface GatewayConfig { url: URL; token: string | undefined; timeoutMs: number }
export const MAX_GATEWAY_BYTES = 2 * 1024 * 1024;
export function gatewayConfig(env: NodeJS.ProcessEnv = process.env): GatewayConfig {
  let url: URL;
  try { url = new URL(env.SWADS_MCP_URL || "https://ads.mincode.cn/mcp"); } catch { throw new HarnessError("SWADS_INVALID_CONFIG", "SWADS_MCP_URL must be a valid HTTP(S) URL"); }
  if (!["https:", "http:"].includes(url.protocol) || url.username || url.password || url.search || url.hash) throw new HarnessError("SWADS_INVALID_CONFIG", "SWADS_MCP_URL must not contain credentials, query or fragment");
  const timeoutMs = Number(env.SWADS_MCP_TIMEOUT_MS ?? 30_000);
  if (!Number.isInteger(timeoutMs) || timeoutMs < 100 || timeoutMs > 120_000) throw new HarnessError("SWADS_INVALID_CONFIG", "SWADS_MCP_TIMEOUT_MS must be between 100 and 120000");
  return { url, token: env.SWADS_MCP_TOKEN || undefined, timeoutMs };
}

// Read the wire body with a byte bound before handing it to the official SDK.
// A connection belongs to one command, so cancellation never affects another run.
export function mcpConnector(config: GatewayConfig, fetcher: typeof fetch = fetch): GatewayConnector {
  return async (signal) => {
    const boundedFetch: typeof fetch = async (input, init) => {
      const response = await fetcher(input, { ...init, redirect: "error", signal: AbortSignal.any([signal, ...(init?.signal ? [init.signal] : [])]) });
      if (!response.ok) { await response.body?.cancel(); throw new GatewayHttpError(response.status); }
      if (!response.body) return response;
      if (Number(response.headers.get("content-length")) > MAX_GATEWAY_BYTES) { await response.body.cancel(); throw failure("SWADS_RESPONSE_TOO_LARGE"); }
      const reader = response.body.getReader(); let total = 0;
      const body = new ReadableStream<Uint8Array>({
        async pull(controller) {
          try {
            const { value, done } = await reader.read();
            if (done) { controller.close(); return; }
            total += value.byteLength;
            if (total > MAX_GATEWAY_BYTES) { await reader.cancel(); controller.error(failure("SWADS_RESPONSE_TOO_LARGE")); return; }
            controller.enqueue(value);
          } catch (error) { controller.error(error); }
        },
        cancel: () => reader.cancel(),
      });
      return new Response(body, { status: response.status, headers: response.headers });
    };
    const client = new Client({ name: "agent-studio-swads-cli", version: "1.0.0" });
    const transport = new StreamableHTTPClientTransport(config.url, { fetch: boundedFetch, requestInit: { headers: { Authorization: `Bearer ${config.token}` } }, reconnectionOptions: { maxRetries: 0, initialReconnectionDelay: 1000, maxReconnectionDelay: 1000, reconnectionDelayGrowFactor: 1 } });
    const close = async () => { await client.close().catch(() => undefined); };
    const abort = () => { void close(); };
    signal.addEventListener("abort", abort, { once: true });
    try { await client.connect(transport as Parameters<Client["connect"]>[0], { signal, timeout: config.timeoutMs }); } catch (error) { signal.removeEventListener("abort", abort); await close(); throw error; }
    return {
      async list(requestSignal) {
        const all: Tool[] = []; let cursor: string | undefined;
        do {
          const page = await client.listTools(cursor ? { cursor } : {}, { signal: requestSignal, timeout: config.timeoutMs });
          all.push(...page.tools); cursor = page.nextCursor;
          if (all.length > 500 || Buffer.byteLength(JSON.stringify(all)) > MAX_GATEWAY_BYTES) throw failure("SWADS_RESPONSE_TOO_LARGE");
        } while (cursor);
        return all;
      },
      call: (name, args, requestSignal) => client.callTool({ name, arguments: args }, undefined, { signal: requestSignal, timeout: config.timeoutMs }),
      async close() { signal.removeEventListener("abort", abort); await close(); },
    };
  };
}
class GatewayHttpError extends Error { constructor(readonly status: number) { super("Gateway HTTP request failed"); } }
export function failure(code: string): HarnessError { return new HarnessError(code, code, 422); }
export function gatewayError(error: unknown, signal?: AbortSignal): HarnessError {
  if (error instanceof HarnessError && error.code.startsWith("SWADS_")) return failure(error.code);
  if (signal?.aborted) return failure(signal.reason?.name === "TimeoutError" ? "SWADS_TIMEOUT" : "SWADS_ABORTED");
  const status = error instanceof GatewayHttpError ? error.status : (error as { code?: number; status?: number } | null)?.status ?? (error as { code?: number } | null)?.code;
  return failure(status === 401 ? "SWADS_AUTH_FAILED" : status === 403 ? "SWADS_FORBIDDEN" : status === -32601 ? "SWADS_TOOL_MISSING" : status === -32602 ? "SWADS_TOOL_DRIFT" : status === -32001 ? "SWADS_TIMEOUT" : "SWADS_UNAVAILABLE");
}
export class SwadsGateway {
  private capabilities: { expires: number; tools: Tool[] } | undefined;
  private readonly connect: GatewayConnector;
  constructor(private readonly config = gatewayConfig(), connector?: GatewayConnector) { this.connect = connector ?? mcpConnector(config); }
  sanitize<T>(value: T): T {
    if (!this.config.token) return value;
    const scrub = (item: unknown): unknown => typeof item === "string" ? item.split(this.config.token!).join("[REDACTED]") : Array.isArray(item) ? item.map(scrub) : item && typeof item === "object" ? Object.fromEntries(Object.entries(item).map(([key, val]) => [String(scrub(key)), /token|authorization|secret|password/i.test(key) ? "[REDACTED]" : scrub(val)])) : item;
    return scrub(value) as T;
  }
  async execute<T>(operation: (connection: GatewayConnection, signal: AbortSignal) => Promise<T>, runSignal: AbortSignal): Promise<T> {
    if (!this.config.token) throw failure("SWADS_NOT_CONFIGURED");
    const signal = AbortSignal.any([runSignal, AbortSignal.timeout(this.config.timeoutMs)]);
    let connection: GatewayConnection | undefined;
    let abort: (() => void) | undefined;
    try {
      signal.throwIfAborted();
      const cancelled = new Promise<never>((_, reject) => { abort = () => reject(gatewayError(undefined, signal)); signal.addEventListener("abort", abort, { once: true }); });
      const work = (async () => {
        connection = await this.connect(signal);
        if (signal.aborted) { await connection.close(); throw gatewayError(undefined, signal); }
        return operation(connection, signal);
      })();
      return this.sanitize(await Promise.race([work, cancelled]));
    } catch (error) { throw gatewayError(error, signal); }
    finally { if (abort) signal.removeEventListener("abort", abort); await connection?.close().catch(() => undefined); }
  }
  async list(connection: GatewayConnection, signal: AbortSignal, refresh = false): Promise<Tool[]> {
    if (!refresh && this.capabilities && this.capabilities.expires > Date.now()) return this.capabilities.tools;
    const tools = await connection.list(signal);
    if (Buffer.byteLength(JSON.stringify(tools)) > MAX_GATEWAY_BYTES) throw failure("SWADS_RESPONSE_TOO_LARGE");
    this.capabilities = { tools, expires: Date.now() + 60_000 }; return tools;
  }
}