import {
  ContentBlockParam,
  MessageCreateParams,
  MessageParam,
} from "@anthropic-ai/sdk/resources";
import {
  ChatCompletionCreateParams,
  ChatCompletionMessageParam,
} from "openai/resources/index";

import {
  CACHING_STRATEGIES,
  CachingStrategyName,
} from "./AnthropicCachingStrategies.js";
import {
  addCacheControlToLastTwoUserMessages,
  openaiToolToAnthropicTool,
} from "./AnthropicUtils.js";

interface SystemMapping {
  openaiIndex: number;
  start: number;
  length: number;
  wasString: boolean;
  originalContent: ChatCompletionMessageParam["content"];
  textPartIndices: (number | null)[];
}

interface MessageMapping {
  openaiIndex: number;
  anthropicIndex: number;
  role: string;
  wasString: boolean;
  originalContent: ChatCompletionMessageParam["content"];
  textPartIndices: (number | null)[];
}

interface ConversionResult {
  anthropicBody: MessageCreateParams;
  systemMappings: SystemMapping[];
  messageMappings: MessageMapping[];
}

const convertContentToBlocks = (
  content: ChatCompletionMessageParam["content"],
): {
  blocks: ContentBlockParam[];
  textPartIndices: (number | null)[];
  wasString: boolean;
} => {
  if (typeof content === "string" || typeof content === "number") {
    const text = String(content);
    return {
      blocks: [
        {
          type: "text",
          text,
        } as ContentBlockParam,
      ],
      textPartIndices: [null],
      wasString: true,
    };
  }

  if (!Array.isArray(content)) {
    return {
      blocks: [],
      textPartIndices: [],
      wasString: false,
    };
  }

  const blocks: ContentBlockParam[] = [];
  const textPartIndices: (number | null)[] = [];

  content.forEach((part: any, idx: number) => {
    if (part?.type === "text") {
      blocks.push({
        type: "text",
        text: part.text ?? "",
      } as ContentBlockParam);
      textPartIndices.push(idx);
    } else {
      blocks.push({ ...(part ?? {}) } as any);
      textPartIndices.push(null);
    }
  });

  return {
    blocks,
    textPartIndices,
    wasString: false,
  };
};

const convertToAnthropic = (
  body: ChatCompletionCreateParams,
): ConversionResult => {
  const systemBlocks: ContentBlockParam[] = [];
  const systemMappings: SystemMapping[] = [];
  const messages: MessageParam[] = [];
  const messageMappings: MessageMapping[] = [];

  let systemOffset = 0;

  body.messages.forEach((message, index) => {
    const { blocks, textPartIndices, wasString } = convertContentToBlocks(
      message.content,
    );

    if (message.role === "system") {
      const length = blocks.length;
      systemMappings.push({
        openaiIndex: index,
        start: systemOffset,
        length,
        wasString,
        originalContent: message.content,
        textPartIndices,
      });
      systemBlocks.push(...blocks);
      systemOffset += length;
    } else {
      messages.push({
        role: message.role as MessageParam["role"],
        content: blocks as any,
      });
      messageMappings.push({
        openaiIndex: index,
        anthropicIndex: messages.length - 1,
        role: message.role,
        wasString,
        originalContent: message.content,
        textPartIndices,
      });
    }
  });

  const tools = body.tools
    ?.filter((tool) => tool.type === "function")
    .map((tool) => openaiToolToAnthropicTool(tool));

  const anthropicBody: MessageCreateParams = {
    model: body.model,
    messages,
    max_tokens: body.max_tokens ?? 1,
    system: systemBlocks.length > 0 ? (systemBlocks as any) : undefined,
    tools,
  };

  return { anthropicBody, systemMappings, messageMappings };
};

export const applyAnthropicCachingToOpenRouterBody = (
  body: ChatCompletionCreateParams,
  strategy: CachingStrategyName,
): void => {
  const { anthropicBody, systemMappings, messageMappings } =
    convertToAnthropic(body);

  const cachingStrategy =
    CACHING_STRATEGIES[strategy] ?? CACHING_STRATEGIES.systemAndTools;
  const cachedBody = cachingStrategy({ ...anthropicBody });

  cachedBody.messages = cachedBody.messages ?? [];
  addCacheControlToLastTwoUserMessages(cachedBody.messages);

  const cachedSystem = Array.isArray(cachedBody.system)
    ? cachedBody.system
    : [];

  systemMappings.forEach((mapping) => {
    const openaiMessage = body.messages[mapping.openaiIndex] as any;
    if (!openaiMessage) {
      return;
    }

    const slice = cachedSystem.slice(
      mapping.start,
      mapping.start + mapping.length,
    );
    const hasCache = slice.some((block: any) => block?.cache_control);

    if (!hasCache) {
      openaiMessage.content = mapping.originalContent;
      return;
    }

    if (mapping.wasString) {
      openaiMessage.content = slice.map((block: any) => ({
        type: "text",
        text: block?.text ?? "",
        ...(block?.cache_control ? { cache_control: block.cache_control } : {}),
      }));
      return;
    }

    if (Array.isArray(mapping.originalContent)) {
      const newParts = mapping.originalContent.map((part: any) => ({
        ...part,
      }));

      slice.forEach((block: any, idx: number) => {
        const originalIndex = mapping.textPartIndices[idx];
        if (
          originalIndex === null ||
          originalIndex === undefined ||
          !block?.cache_control
        ) {
          return;
        }

        newParts[originalIndex] = {
          ...newParts[originalIndex],
          cache_control: block.cache_control,
          ...(block.text !== undefined ? { text: block.text } : {}),
        };
      });

      openaiMessage.content = newParts;
    }
  });

  const cachedMessages = cachedBody.messages ?? [];
  messageMappings.forEach((mapping) => {
    const openaiMessage = body.messages[mapping.openaiIndex] as any;
    const cachedMessage = cachedMessages[mapping.anthropicIndex] as any;
    if (!openaiMessage || !cachedMessage) {
      return;
    }

    if (cachedMessage.role !== "user") {
      openaiMessage.content = mapping.originalContent;
      return;
    }

    const contentArray = Array.isArray(cachedMessage.content)
      ? cachedMessage.content
      : [];
    const hasCache = contentArray.some((block: any) => block?.cache_control);

    if (!hasCache) {
      openaiMessage.content = mapping.originalContent;
      return;
    }

    if (mapping.wasString) {
      openaiMessage.content = contentArray.map((block: any) => ({
        type: "text",
        text: block?.text ?? "",
        ...(block?.cache_control ? { cache_control: block.cache_control } : {}),
      }));
      return;
    }

    if (Array.isArray(mapping.originalContent)) {
      const newParts = mapping.originalContent.map((part: any) => ({
        ...part,
      }));

      contentArray.forEach((block: any, idx: number) => {
        const originalIndex = mapping.textPartIndices[idx];
        if (
          originalIndex === null ||
          originalIndex === undefined ||
          !block?.cache_control
        ) {
          return;
        }

        newParts[originalIndex] = {
          ...newParts[originalIndex],
          cache_control: block.cache_control,
          ...(block.text !== undefined ? { text: block.text } : {}),
        };
      });

      openaiMessage.content = newParts;
    }
  });

  if (body.tools?.length && cachedBody.tools?.length) {
    body.tools = body.tools.map((tool, idx) => {
      const cachedTool = (cachedBody.tools ?? [])[idx] as any;
      if (!cachedTool?.cache_control) {
        return tool;
      }
      return {
        ...tool,
        cache_control: cachedTool.cache_control,
      } as any;
    });
  }
};