import type { BaseMessage } from '@langchain/core/messages';
import { SystemMessage, HumanMessage, AIMessage } from '@langchain/core/messages';
import type { StructuredTool } from '@langchain/core/tools';

/**
 * Provider-agnostic hook to preprocess messages before an LLM call.
 * Applied transparently by the harness; implementations are provider-specific.
 */
export interface ModelMiddleware {
  processMessages(messages: BaseMessage[], tools: StructuredTool[]): BaseMessage[];
}

/**
 * No-op implementation for providers that don't support prompt caching.
 */
export class NoopMiddleware implements ModelMiddleware {
  processMessages(messages: BaseMessage[]): BaseMessage[] {
    return messages;
  }
}

/**
 * Mirrors Python's AnthropicPromptCachingMiddleware logic.
 *
 * Tags with cache_control: { type: 'ephemeral' }:
 * 1. The system message last content block (stable prompt prefix — always a cache hit after turn 1)
 * 2. The last non-current cacheable message (AIMessage or HumanMessage before the final user message)
 *    so the full conversation prefix is cached on the next turn.
 *
 * Requires ChatAnthropic to be configured with:
 *   clientOptions: { defaultHeaders: { 'anthropic-beta': 'prompt-caching-2024-07-31' } }
 */
export class AnthropicCachingMiddleware implements ModelMiddleware {
  processMessages(messages: BaseMessage[], _tools: StructuredTool[]): BaseMessage[] {
    if (messages.length === 0) return messages;

    const result = messages.map(msg => cloneMessage(msg));

    // 1. Tag system message
    const systemMsg = result.find(m => m._getType() === 'system');
    if (systemMsg) {
      addCacheControl(systemMsg);
    }

    // 2. Tag the last AIMessage (assistant turn) before the current user input.
    // Workspace state is injected as a HumanMessage just before the current user input;
    // by matching only 'ai' we ensure the cache boundary is the last assistant response,
    // leaving workspace context and the new user input uncached (always fresh).
    const candidates = result.slice(0, -1);
    for (let i = candidates.length - 1; i >= 0; i--) {
      const t = candidates[i]._getType();
      if (t === 'ai') {
        addCacheControl(candidates[i]);
        break;
      }
    }

    return result;
  }
}

/**
 * Shallow-clone a message so we don't mutate history objects.
 */
function cloneMessage(msg: BaseMessage): BaseMessage {
  const type = msg._getType();
  const content = typeof msg.content === 'string'
    ? msg.content
    : JSON.parse(JSON.stringify(msg.content));

  if (type === 'system') return new SystemMessage({ content, additional_kwargs: { ...msg.additional_kwargs } });
  if (type === 'human') return new HumanMessage({ content, additional_kwargs: { ...msg.additional_kwargs } });
  if (type === 'ai') return new AIMessage({ content, additional_kwargs: { ...msg.additional_kwargs }, tool_calls: (msg as AIMessage).tool_calls });
  // For other types (tool messages etc.), return as-is — we don't tag them
  return msg;
}

/**
 * Add cache_control to the last content block of a message.
 * Converts string content to a block array if needed.
 */
function addCacheControl(msg: BaseMessage): void {
  if (typeof msg.content === 'string') {
    // Convert to block array
    (msg as any).content = [{ type: 'text', text: msg.content, cache_control: { type: 'ephemeral' } }];
  } else if (Array.isArray(msg.content) && msg.content.length > 0) {
    const last = msg.content[msg.content.length - 1] as any;
    last.cache_control = { type: 'ephemeral' };
  }
}
