/**
 * Per-agent JSON-RPC client. Talks to the deployed agent's POST /mcp endpoint.
 *
 * We hand-roll the wire calls (instead of using @modelcontextprotocol/sdk's
 * Client) because (a) the surface we need is tiny — tools/list, tools/call —
 * and (b) we want to send the user's bearer token as a plain Authorization
 * header without coupling to a transport class.
 */
const JSON_RPC = "2.0";
const DEFAULT_RPC_TIMEOUT_MS = 30_000;
const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 120_000;

export interface UpstreamConfig {
  name: string;
  url: string;
  token: string | null;
  requestTimeoutMs?: number;
  streamIdleTimeoutMs?: number;
  /** Forwarded to the upstream as ``params._meta.cp_jwt`` on tools/call.
   *  Lets the upstream agent act on the caller's behalf (file/agent CRUD
   *  on /v1/me/*) — same path the platform orchestrator uses. */
  cpJwt?: string | null;
  /** Paired with ``cpJwt``. The upstream agent uses this as the base URL
   *  when calling back into /v1/me/*. */
  cpUrl?: string | null;
  /** User's MinIO bucket (``user-<id>-files``). Forwarded as
   *  ``params._meta.bucket`` so agents that touch the workspace
   *  (agent-builder, file tools) get a context whose
   *  ``ctx.workspace.bucket`` resolves to the right bucket. */
  bucket?: string | null;
}

export interface UpstreamTool {
  name: string;
  description?: string;
  inputSchema: unknown;
  outputSchema?: unknown;
}

export interface UpstreamCallResult {
  content: Array<{ type: string; text?: string; [k: string]: unknown }>;
  structuredContent?: unknown;
  isError?: boolean;
}

export class UpstreamError extends Error {
  constructor(
    public agent: string,
    public status: number,
    message: string,
  ) {
    super(`upstream ${agent}: ${message}`);
    this.name = "UpstreamError";
  }
}

let _nextId = 1;

export class UpstreamAgent {
  constructor(private cfg: UpstreamConfig) {
    this.cfg.url = cfg.url.replace(/\/+$/, "");
  }

  get name(): string {
    return this.cfg.name;
  }

  private async rpc<T>(method: string, params: unknown): Promise<T> {
    const id = _nextId++;
    const timeoutMs = this.cfg.requestTimeoutMs ?? DEFAULT_RPC_TIMEOUT_MS;
    const controller = new AbortController();
    let timedOut = false;
    const timeout = setTimeout(() => {
      timedOut = true;
      controller.abort();
    }, timeoutMs);
    const headers: Record<string, string> = {
      "content-type": "application/json",
      accept: "application/json",
    };
    if (this.cfg.token) headers["authorization"] = `Bearer ${this.cfg.token}`;

    try {
      const resp = await fetch(`${this.cfg.url}/mcp`, {
        method: "POST",
        headers,
        body: JSON.stringify({ jsonrpc: JSON_RPC, id, method, params }),
        signal: controller.signal,
      });

      if (!resp.ok) {
        const text = await resp.text().catch(() => "");
        throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText);
      }

      const env: any = await resp.json();
      if (env?.error) {
        throw new UpstreamError(
          this.cfg.name,
          500,
          `${env.error.code}: ${env.error.message}`,
        );
      }
      return env?.result as T;
    } catch (err) {
      if (timedOut) {
        throw new UpstreamError(
          this.cfg.name,
          504,
          `request timed out after ${timeoutMs}ms`,
        );
      }
      throw err;
    } finally {
      clearTimeout(timeout);
    }
  }

  async listTools(): Promise<UpstreamTool[]> {
    const out = await this.rpc<{ tools: UpstreamTool[] }>("tools/list", {});
    return out?.tools ?? [];
  }

  async callTool(
    name: string,
    arguments_: Record<string, unknown>,
    opts: {
      onNotification?: (notif: {
        method: string;
        params?: Record<string, unknown>;
      }) => void | Promise<void>;
      progressToken?: string | number;
    } = {},
  ): Promise<UpstreamCallResult> {
    const meta: Record<string, unknown> = {};
    if (this.cfg.cpJwt) meta.cp_jwt = this.cfg.cpJwt;
    if (this.cfg.cpUrl) meta.cp_url = this.cfg.cpUrl;
    if (this.cfg.bucket) meta.bucket = this.cfg.bucket;
    if (opts.progressToken !== undefined) meta.progressToken = opts.progressToken;
    const params: Record<string, unknown> = { name, arguments: arguments_ };
    if (Object.keys(meta).length > 0) params._meta = meta;
    return this.streamingCallTool(params, opts.onNotification);
  }

  /**
   * tools/call always negotiates SSE on the wire — upstream agents can
   * take minutes (agent-builder.build runs an LLM loop + sandbox tests +
   * a deploy + a wait-for-live-card poll), and a plain JSON POST gets
   * killed by ingress idle timeouts (~60s default on traefik). The
   * upstream's SSE response interleaves any elicitation/create requests
   * and a final tools/call result; here we drain the stream and return
   * the final result. (Elicit relaying to the MCP client is not wired
   * yet — any elicit that fires will time out on the upstream side and
   * the skill will surface a soft error.)
   */
  private async streamingCallTool(
    params: Record<string, unknown>,
    onNotification?: (notif: {
      method: string;
      params?: Record<string, unknown>;
    }) => void | Promise<void>,
  ): Promise<UpstreamCallResult> {
    const id = _nextId++;
    const idleTimeoutMs =
      this.cfg.streamIdleTimeoutMs ?? DEFAULT_STREAM_IDLE_TIMEOUT_MS;
    const controller = new AbortController();
    let timedOut = false;
    let idleTimer: ReturnType<typeof setTimeout> | undefined;
    const resetIdleTimer = () => {
      if (idleTimer) clearTimeout(idleTimer);
      idleTimer = setTimeout(() => {
        timedOut = true;
        controller.abort();
      }, idleTimeoutMs);
    };
    resetIdleTimer();
    const headers: Record<string, string> = {
      "content-type": "application/json",
      accept: "text/event-stream, application/json",
    };
    if (this.cfg.token) headers["authorization"] = `Bearer ${this.cfg.token}`;

    let resp: Response;
    try {
      resp = await fetch(`${this.cfg.url}/mcp`, {
        method: "POST",
        headers,
        body: JSON.stringify({
          jsonrpc: JSON_RPC, id, method: "tools/call", params,
        }),
        signal: controller.signal,
      });
    } catch (err) {
      clearTimeout(idleTimer);
      if (timedOut) {
        throw new UpstreamError(
          this.cfg.name,
          504,
          `stream idle timeout after ${idleTimeoutMs}ms`,
        );
      }
      throw err;
    }
    resetIdleTimer();

    if (!resp.ok) {
      const text = await resp.text().catch(() => "");
      clearTimeout(idleTimer);
      throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText);
    }

    const ctype = resp.headers.get("content-type") || "";
    if (!ctype.includes("text/event-stream")) {
      // Server fell back to JSON (older a2a-pack or non-SSE upstream).
      const env: any = await resp.json();
      if (env?.error) {
        clearTimeout(idleTimer);
        throw new UpstreamError(
          this.cfg.name, 500,
          `${env.error.code}: ${env.error.message}`,
        );
      }
      clearTimeout(idleTimer);
      return env?.result as UpstreamCallResult;
    }

    if (!resp.body) {
      clearTimeout(idleTimer);
      throw new UpstreamError(this.cfg.name, 502, "empty stream body");
    }

    const reader = resp.body.getReader();
    const decoder = new TextDecoder();
    let buf = "";
    try {
      while (true) {
        let done: boolean;
        let value: Uint8Array | undefined;
        try {
          ({ done, value } = await reader.read());
        } catch (err) {
          if (timedOut) {
            throw new UpstreamError(
              this.cfg.name,
              504,
              `stream idle timeout after ${idleTimeoutMs}ms`,
            );
          }
          throw err;
        }
        if (value) {
          resetIdleTimer();
          buf += decoder.decode(value, { stream: true });
        }
        let sep = buf.indexOf("\n\n");
        while (sep !== -1) {
          const frame = buf.slice(0, sep);
          buf = buf.slice(sep + 2);
          const dataLines: string[] = [];
          for (const ln of frame.split("\n")) {
            if (ln.startsWith("data:")) dataLines.push(ln.slice(5).trimStart());
          }
          if (dataLines.length > 0) {
            const raw = dataLines.join("\n");
            if (raw === "[DONE]") {
              sep = buf.indexOf("\n\n");
              continue;
            }
            let msg: any;
            try { msg = JSON.parse(raw); } catch { msg = null; }
            if (msg && msg.id === id) {
              if (msg.error) {
                throw new UpstreamError(
                  this.cfg.name, 500,
                  `${msg.error.code}: ${msg.error.message}`,
                );
              }
              return msg.result as UpstreamCallResult;
            }
            // Otherwise it's a server→client message. Notifications
            // (no id) → forward to the MCP client so long-running
            // skills keep the client's tool-call timer alive. Requests
            // (has id, has method) — e.g. elicitation/create — are
            // not relayed yet, so they will time out upstream.
            if (msg && typeof msg.method === "string" && !("id" in msg) && onNotification) {
              try {
                await onNotification({
                  method: msg.method,
                  params: msg.params as Record<string, unknown> | undefined,
                });
              } catch {
                // Best-effort; never break the stream on a notify error.
              }
            }
          }
          sep = buf.indexOf("\n\n");
        }
        if (done) break;
      }
    } finally {
      clearTimeout(idleTimer);
      try { reader.releaseLock(); } catch { /* ignore */ }
    }
    throw new UpstreamError(this.cfg.name, 502, "stream ended without result");
  }
}