import { getChunkSize, type ChunkingManager } from "./chunking";
import { countTokens, encoding } from "./tokens";

export interface SerializedToolResult {
  content: Array<{ type: "text"; text: string }>;
  details: Record<string, unknown>;
}

function normalizeForJson(value: unknown, seen = new WeakSet<object>()): unknown {
  if (value === null) {
    return null;
  }

  if (value === undefined) {
    return undefined;
  }

  const valueType = typeof value;
  if (valueType === "string" || valueType === "number" || valueType === "boolean") {
    return value;
  }

  if (valueType === "bigint") {
    return value.toString();
  }

  if (valueType === "function" || valueType === "symbol") {
    return String(value);
  }

  if (value instanceof Date) {
    return value.toISOString();
  }

  if (value instanceof Error) {
    return value.toString();
  }

  if (Array.isArray(value)) {
    return value.map((item) => normalizeForJson(item, seen));
  }

  if (value instanceof Map || value instanceof Set || ArrayBuffer.isView(value) || value instanceof ArrayBuffer) {
    return String(value);
  }

  if (valueType === "object") {
    if (seen.has(value as object)) {
      return "[Circular]";
    }
    seen.add(value as object);

    const prototype = Object.getPrototypeOf(value);
    if (prototype === Object.prototype || prototype === null) {
      const output: Record<string, unknown> = {};
      for (const key of Object.keys(value as Record<string, unknown>)) {
        output[key] = normalizeForJson((value as Record<string, unknown>)[key], seen);
      }
      return output;
    }

    return String(value);
  }

  return String(value);
}

export function serializeResponse(data: unknown): string {
  const normalized = normalizeForJson(data);
  return JSON.stringify(normalized === undefined ? null : normalized);
}

export function serializeResponseWithMeta(data: unknown): SerializedToolResult {
  const responseText = serializeResponse(data);
  const trackTokens = (process.env.RTFD_TRACK_TOKENS ?? "false").toLowerCase() === "true";

  if (!trackTokens) {
    return {
      content: [{ type: "text", text: responseText }],
      details: {},
    };
  }

  try {
    const tokenCount = countTokens(responseText);
    return {
      content: [{ type: "text", text: responseText }],
      details: {
        token_stats: {
          tokens_json: tokenCount,
          tokens_sent: tokenCount,
          format: "json",
          bytes_json: responseText.length,
        },
      },
    };
  } catch (error) {
    return {
      content: [{ type: "text", text: responseText }],
      details: {
        token_stats: {
          error: `Token counting failed: ${error instanceof Error ? error.message : String(error)}`,
        },
      },
    };
  }
}

export function chunkAndSerializeResponse(
  data: Record<string, unknown>,
  contentKey = "content",
  chunkingManager?: ChunkingManager,
): SerializedToolResult {
  const chunkSize = getChunkSize();

  if (chunkSize === 0 || !(contentKey in data)) {
    return serializeResponseWithMeta(data);
  }

  if (!chunkingManager) {
    return serializeResponseWithMeta(data);
  }

  const content = data[contentKey];
  if (typeof content !== "string" || !content) {
    return serializeResponseWithMeta(data);
  }

  const contentTokens = countTokens(content);
  if (contentTokens <= chunkSize) {
    const resultData: Record<string, unknown> = { ...data };
    resultData.chunking = {
      is_chunked: false,
      tokens_in_content: contentTokens,
    };
    return serializeResponseWithMeta(resultData);
  }

  const tokens = encoding.encode(content);
  const firstChunkTokens = tokens.slice(0, chunkSize);
  const remainingTokens = tokens.slice(chunkSize);

  const firstChunkContent = encoding.decode(firstChunkTokens);
  const remainingContent = encoding.decode(remainingTokens);

  const metadata = {
    chunk_number: 1,
    total_tokens: contentTokens,
    original_data: Object.fromEntries(Object.entries(data).filter(([key]) => key !== contentKey)),
  };

  const continuationToken = chunkingManager.storeContinuation(remainingContent, metadata);

  const resultData: Record<string, unknown> = { ...data };
  resultData[contentKey] = firstChunkContent;
  resultData.chunking = {
    is_chunked: true,
    chunk_number: 1,
    has_more: true,
    continuation_token: continuationToken,
    tokens_in_chunk: firstChunkTokens.length,
    remaining_tokens: remainingTokens.length,
    hint: `Call get_next_chunk('${continuationToken}') for more content`,
  };

  return serializeResponseWithMeta(resultData);
}