"use client";

import { Client as LangGraphClient } from "@langchain/langgraph-sdk/client";

import { getAuthorizationHeaderValue } from "../auth";
import { getLangGraphBaseURL } from "../config";

import { withCsrfHeaders } from "./csrf";
import { sanitizeRunStreamOptions } from "./stream-mode";

function createCompatibleClient(isMock?: boolean): LangGraphClient {
  const client = new LangGraphClient({
    apiUrl: getLangGraphBaseURL(isMock),
    onRequest(_url, init) {
      const headers = withCsrfHeaders(init.headers);
      const authHeader = getAuthorizationHeaderValue();
      if (authHeader && !headers.has("Authorization")) {
        headers.set("Authorization", authHeader);
      }
      return {
        ...init,
        credentials: init.credentials ?? "include",
        headers,
      };
    },
  });

  const originalRunStream = client.runs.stream.bind(client.runs);
  client.runs.stream = ((threadId, assistantId, payload) =>
    originalRunStream(
      threadId,
      assistantId,
      sanitizeRunStreamOptions(payload),
    )) as typeof client.runs.stream;

  const originalJoinStream = client.runs.joinStream.bind(client.runs);
  client.runs.joinStream = ((threadId, runId, options) =>
    originalJoinStream(
      threadId,
      runId,
      sanitizeRunStreamOptions(options),
    )) as typeof client.runs.joinStream;

  return client;
}

const _clients = new Map<string, LangGraphClient>();
export function getAPIClient(isMock?: boolean): LangGraphClient {
  const cacheKey = isMock ? "mock" : "default";
  let client = _clients.get(cacheKey);

  if (!client) {
    client = createCompatibleClient(isMock);
    _clients.set(cacheKey, client);
  }

  return client;
}