import {
  AIMessage,
  BaseMessage,
  ToolMessage,
  HumanMessage,
  SystemMessage,
  MessageContentComplex,
} from '@langchain/core/messages';
import type { AnthropicMessage } from '@/types/messages';
import type Anthropic from '@anthropic-ai/sdk';
import { ContentTypes } from '@/common/enum';

type MessageWithContent = {
  content?: string | MessageContentComplex[];
};

/** Debug logger for cache operations - set ILLUMA_DEBUG_CACHE=true to enable */
const debugCache = (message: string, data?: unknown): void => {
  if (process.env.ILLUMA_DEBUG_CACHE === 'true') {
    // eslint-disable-next-line no-console
    console.log(
      `[Cache] ${message}`,
      data !== undefined ? JSON.stringify(data, null, 2) : ''
    );
  }
};

/**
 * Deep clones a message's content to prevent mutation of the original.
 */
function deepCloneContent<T extends string | MessageContentComplex[]>(
  content: T
): T {
  if (typeof content === 'string') {
    return content;
  }
  if (Array.isArray(content)) {
    return content.map((block) => ({ ...block })) as T;
  }
  return content;
}

/**
 * Clones a message with new content. For LangChain BaseMessage instances,
 * constructs a proper class instance so that `instanceof` checks are preserved
 * in downstream code (e.g., ensureThinkingBlockInMessages).
 * For plain objects (AnthropicMessage), uses object spread.
 */
function cloneMessage<T extends MessageWithContent>(
  message: T,
  content: string | MessageContentComplex[]
): T {
  if (message instanceof BaseMessage) {
    const baseParams = {
      content,
      additional_kwargs: { ...message.additional_kwargs },
      response_metadata: { ...message.response_metadata },
      id: message.id,
      name: message.name,
    };

    const msgType = message.getType();
    switch (msgType) {
    case 'ai':
      return new AIMessage({
        ...baseParams,
        tool_calls: (message as unknown as AIMessage).tool_calls,
      }) as unknown as T;
    case 'human':
      return new HumanMessage(baseParams) as unknown as T;
    case 'system':
      return new SystemMessage(baseParams) as unknown as T;
    case 'tool':
      return new ToolMessage({
        ...baseParams,
        tool_call_id: (message as unknown as ToolMessage).tool_call_id,
      }) as unknown as T;
    default:
      break;
    }
  }

  const {
    lc_kwargs: _lc_kwargs,
    lc_serializable: _lc_serializable,
    lc_namespace: _lc_namespace,
    ...rest
  } = message as T & {
    lc_kwargs?: unknown;
    lc_serializable?: unknown;
    lc_namespace?: unknown;
  };

  const cloned = { ...rest, content } as T;

  // Sync lc_kwargs.content with the new content to prevent LangChain coercion issues
  const lcKwargs = (message as Record<string, unknown>).lc_kwargs as
    | Record<string, unknown>
    | undefined;
  if (lcKwargs != null) {
    (cloned as Record<string, unknown>).lc_kwargs = {
      ...lcKwargs,
      content: content,
    };
  }

  // LangChain messages don't have a direct 'role' property - derive it from getType()
  if (
    'getType' in message &&
    typeof message.getType === 'function' &&
    !('role' in cloned)
  ) {
    const msgType = (message as unknown as BaseMessage).getType();
    const roleMap: Record<string, string> = {
      human: 'user',
      ai: 'assistant',
      system: 'system',
      tool: 'tool',
    };
    (cloned as Record<string, unknown>).role = roleMap[msgType] || msgType;
  }

  return cloned;
}

/**
 * Checks if a content block is a cache point
 */
function isCachePoint(block: MessageContentComplex): boolean {
  return 'cachePoint' in block && !('type' in block);
}

/**
 * Checks if a message's content needs cache control stripping.
 * Returns true if content has cachePoint blocks or cache_control fields.
 */
function needsCacheStripping(content: MessageContentComplex[]): boolean {
  for (let i = 0; i < content.length; i++) {
    const block = content[i];
    if (isCachePoint(block)) return true;
    if ('cache_control' in block) return true;
  }
  return false;
}

/**
 * Checks if a message's content has Anthropic cache_control fields.
 */
function hasAnthropicCacheControl(content: MessageContentComplex[]): boolean {
  for (let i = 0; i < content.length; i++) {
    if ('cache_control' in content[i]) return true;
  }
  return false;
}

/**
 * Checks if a message's content has Bedrock cachePoint blocks.
 */
function hasBedrockCachePoint(content: MessageContentComplex[]): boolean {
  for (let i = 0; i < content.length; i++) {
    if (isCachePoint(content[i])) return true;
  }
  return false;
}

/**
 * Anthropic API: Adds cache control to the appropriate user messages in the payload.
 * Strips ALL existing cache control (both Anthropic and Bedrock formats) from all messages,
 * then adds fresh cache control to the last 2 user messages in a single backward pass.
 * This ensures we don't accumulate stale cache points across multiple turns.
 * Returns a new array - only clones messages that require modification.
 * @param messages - The array of message objects.
 * @returns - A new array of message objects with cache control added.
 */
export function addCacheControl<T extends AnthropicMessage | BaseMessage>(
  messages: T[]
): T[] {
  if (!Array.isArray(messages) || messages.length < 2) {
    return messages;
  }

  const updatedMessages: T[] = [...messages];
  let userMessagesModified = 0;

  for (let i = updatedMessages.length - 1; i >= 0; i--) {
    const originalMessage = updatedMessages[i];
    const content = originalMessage.content;
    const isUserMessage =
      ('getType' in originalMessage && originalMessage.getType() === 'human') ||
      ('role' in originalMessage && originalMessage.role === 'user');

    const hasArrayContent = Array.isArray(content);
    const needsStripping =
      hasArrayContent &&
      needsCacheStripping(content as MessageContentComplex[]);
    const needsCacheAdd =
      userMessagesModified < 2 &&
      isUserMessage &&
      (typeof content === 'string' || hasArrayContent);

    if (!needsStripping && !needsCacheAdd) {
      continue;
    }

    let workingContent: MessageContentComplex[];

    if (hasArrayContent) {
      workingContent = deepCloneContent(
        content as MessageContentComplex[]
      ).filter((block) => !isCachePoint(block as MessageContentComplex));

      for (let j = 0; j < workingContent.length; j++) {
        const block = workingContent[j] as Record<string, unknown>;
        if ('cache_control' in block) {
          delete block.cache_control;
        }
      }
    } else if (typeof content === 'string') {
      workingContent = [
        { type: 'text', text: content },
      ] as MessageContentComplex[];
    } else {
      workingContent = [];
    }

    if (userMessagesModified >= 2 || !isUserMessage) {
      updatedMessages[i] = cloneMessage(
        originalMessage as MessageWithContent,
        workingContent
      ) as T;
      continue;
    }

    for (let j = workingContent.length - 1; j >= 0; j--) {
      const contentPart = workingContent[j];
      if ('type' in contentPart && contentPart.type === 'text') {
        (contentPart as Anthropic.TextBlockParam).cache_control = {
          type: 'ephemeral',
        };
        userMessagesModified++;
        break;
      }
    }

    updatedMessages[i] = cloneMessage(
      originalMessage as MessageWithContent,
      workingContent
    ) as T;
  }

  return updatedMessages;
}

/**
 * Removes all Anthropic cache_control fields from messages
 * Used when switching from Anthropic to Bedrock provider
 * Returns a new array - only clones messages that require modification.
 */
export function stripAnthropicCacheControl<T extends MessageWithContent>(
  messages: T[]
): T[] {
  if (!Array.isArray(messages)) {
    return messages;
  }

  const updatedMessages: T[] = [...messages];

  for (let i = 0; i < updatedMessages.length; i++) {
    const originalMessage = updatedMessages[i];
    const content = originalMessage.content;

    if (!Array.isArray(content) || !hasAnthropicCacheControl(content)) {
      continue;
    }

    const clonedContent = deepCloneContent(content);
    for (let j = 0; j < clonedContent.length; j++) {
      const block = clonedContent[j] as Record<string, unknown>;
      if ('cache_control' in block) {
        delete block.cache_control;
      }
    }
    updatedMessages[i] = cloneMessage(originalMessage, clonedContent) as T;
  }

  return updatedMessages;
}

/**
 * Removes all Bedrock cachePoint blocks from messages
 * Used when switching from Bedrock to Anthropic provider
 * Returns a new array - only clones messages that require modification.
 */
export function stripBedrockCacheControl<T extends MessageWithContent>(
  messages: T[]
): T[] {
  if (!Array.isArray(messages)) {
    return messages;
  }

  const updatedMessages: T[] = [...messages];

  for (let i = 0; i < updatedMessages.length; i++) {
    const originalMessage = updatedMessages[i];
    const content = originalMessage.content;

    if (!Array.isArray(content) || !hasBedrockCachePoint(content)) {
      continue;
    }

    const clonedContent = deepCloneContent(content).filter(
      (block) => !isCachePoint(block as MessageContentComplex)
    );
    updatedMessages[i] = cloneMessage(originalMessage, clonedContent) as T;
  }

  return updatedMessages;
}

/**
 * Adds Bedrock Converse API cache points using "Stable Prefix Caching" strategy.
 *
 * STRATEGY: Place cache point after the LAST ASSISTANT message only.
 * This ensures the prefix (everything before the cache point) remains STABLE
 * as the conversation grows, maximizing cache hits.
 *
 * Why this works:
 * - System message has its own cachePoint (added in AgentContext)
 * - Tools have their own cachePoint (added in IllumaBedrockConverse)
 * - Conversation history grows, but the PREFIX stays the same
 * - Only the NEW user message is uncached (it's always different)
 *
 * Example conversation flow:
 * Request 1: [System+cachePoint][Tools+cachePoint][User1] → No conversation cache yet
 * Request 2: [System][Tools][User1][Assistant1+cachePoint][User2] → Cache User1+Assistant1
 * Request 3: [System][Tools][User1][Assistant1][User2][Assistant2+cachePoint][User3]
 *            → Cache reads User1+A1+User2+A2, cache writes new portion
 *
 * Claude's "Simplified Cache Management" automatically looks back up to 20 content
 * blocks from the cache checkpoint to find the longest matching prefix.
 *
 * @param messages - The array of message objects (excluding system message).
 * @returns - The updated array with a single cache point after the last assistant message.
 */
export function addBedrockCacheControl<
  T extends Partial<BaseMessage> & MessageWithContent,
>(messages: T[]): T[] {
  if (!Array.isArray(messages) || messages.length < 1) {
    debugCache('addBedrockCacheControl: Skipping - no messages', {
      count: messages.length,
    });
    return messages;
  }

  debugCache(
    'addBedrockCacheControl: Processing messages with stable prefix strategy',
    {
      count: messages.length,
    }
  );

  // Clone messages to avoid mutating originals
  const updatedMessages: T[] = messages.map((msg) => {
    const content = msg.content;
    if (Array.isArray(content)) {
      // Strip existing cachePoint blocks and Anthropic-style cache_control
      const stripped = content
        .filter((block) => !isCachePoint(block))
        .map((block) => {
          const rec = block as Record<string, unknown>;
          if ('cache_control' in rec) {
            const { cache_control: _, ...rest } = rec;
            return rest as MessageContentComplex;
          }
          return block;
        });
      if (needsCacheStripping(content) || hasAnthropicCacheControl(content)) {
        return cloneMessage(msg, stripped) as T;
      }
    }
    return cloneMessage(msg, content as string | MessageContentComplex[]) as T;
  });

  // Helper function to check if a message contains reasoning/thinking blocks
  const hasReasoningBlock = (message: T): boolean => {
    const content = message.content;
    if (!Array.isArray(content)) {
      return false;
    }
    for (const block of content) {
      const type = (block as { type?: string }).type;
      if (
        type === 'reasoning_content' ||
        type === 'reasoning' ||
        type === 'thinking' ||
        type === 'redacted_thinking'
      ) {
        return true;
      }
    }
    return false;
  };

  /**
   * Helper to add a cachePoint to a message's content.
   * For strings: wraps into [{type: 'text', text}, {cachePoint}]
   * For arrays: inserts cachePoint after the last non-whitespace text block
   * Returns the new message or null if no suitable insertion point.
   */
  const addCachePointToMessage = (message: T): T | null => {
    const msgContent = message.content;

    if (typeof msgContent === 'string' && msgContent !== '') {
      const newContent = [
        { type: ContentTypes.TEXT, text: msgContent },
        { cachePoint: { type: 'default' } },
      ] as MessageContentComplex[];
      return cloneMessage(message, newContent) as T;
    }

    if (Array.isArray(msgContent) && msgContent.length > 0) {
      if (hasReasoningBlock(message)) {
        return null;
      }
      // Find the last text block and insert cache point after it
      for (let j = msgContent.length - 1; j >= 0; j--) {
        const type = (msgContent[j] as { type?: string }).type;
        if (type === ContentTypes.TEXT || type === 'text') {
          const text = (msgContent[j] as { text?: string }).text;
          if (text != null && text.trim() !== '') {
            const newContent = [
              ...msgContent.slice(0, j + 1),
              { cachePoint: { type: 'default' } } as MessageContentComplex,
              ...msgContent.slice(j + 1),
            ];
            return cloneMessage(message, newContent) as T;
          }
        }
      }
    }

    return null;
  };

  // Add cache points to the last 2 messages (from the end) that have eligible content.
  // This mirrors Anthropic's two-breakpoint strategy for Bedrock's cache API.
  // Skip messages with only whitespace, empty content, or reasoning blocks.
  const MAX_CACHE_POINTS = 2;
  let applied = 0;

  for (
    let i = updatedMessages.length - 1;
    i >= 0 && applied < MAX_CACHE_POINTS;
    i--
  ) {
    const message = updatedMessages[i];
    const msgContent = message.content;

    // Skip empty/whitespace-only content
    if (msgContent == null) continue;
    if (typeof msgContent === 'string' && msgContent.trim() === '') continue;
    if (Array.isArray(msgContent) && msgContent.length === 0) continue;

    // Skip non-string, non-array content (e.g., number, object without type)
    if (typeof msgContent !== 'string' && !Array.isArray(msgContent)) continue;

    // Skip AI messages with only whitespace text and reasoning blocks (tool-call scenario)
    const messageType =
      'getType' in message && typeof message.getType === 'function'
        ? message.getType()
        : 'unknown';
    const role = (message as Record<string, unknown>).role as
      | string
      | undefined;
    const isAi = messageType === 'ai' || role === 'assistant';
    if (isAi && hasReasoningBlock(message)) {
      // Check if all text blocks are whitespace-only
      if (Array.isArray(msgContent)) {
        const hasNonWhitespaceText = msgContent.some((block) => {
          const type = (block as { type?: string }).type;
          if (type === ContentTypes.TEXT || type === 'text') {
            const text = (block as { text?: string }).text;
            return text != null && text.trim() !== '';
          }
          return false;
        });
        if (!hasNonWhitespaceText) {
          debugCache(
            `⚠️ Message cachePoint SKIPPED at index ${i} (AI with whitespace-only text + reasoning)`
          );
          continue;
        }
      }
    }

    // Skip messages that already have cache points (multi-agent scenarios)
    if (
      Array.isArray(msgContent) &&
      msgContent.some((b) => isCachePoint(b as MessageContentComplex))
    ) {
      continue;
    }

    const updated = addCachePointToMessage(message);
    if (updated != null) {
      updatedMessages[i] = updated;
      applied++;
      debugCache(
        `📍 Message cachePoint at index ${i} (${typeof message.content === 'string' ? 'string' : 'array'})`
      );
    }
  }

  debugCache(
    'addBedrockCacheControl: Complete - stable prefix caching applied',
    {
      appliedCachePoints: applied,
      totalMessages: updatedMessages.length,
    }
  );

  return updatedMessages;
}