// ABOUTME: Message building utilities for converting durable events to provider messages

import type { ToolResult } from '@lace/ent-protocol';
import type { ContentBlock, ProviderMessage, ThinkingBlock } from '../providers/base-provider';
import { appendOrMergeUser } from './append-or-merge';
import { foldEvent, initialFoldState, type FoldState } from './fold-event';
import type { ToolCall as CoreToolCall, ToolResult as CoreToolResult } from '../tools/types';
import { estimateTokens } from '@lace/agent/utils/token-estimation';
import { logger } from '@lace/agent/utils/logger';
import { readParsedSessionEvents, type ParsedSessionEvent } from './parsed-events';

// Typed shapes for parsing event data
type TextBlock = { type: 'text'; text: string };
type ContentBlockShape = { type?: unknown; text?: unknown };
type ContextInjectedData = { content?: unknown[] };
type ContextCompactedData = { preserved?: unknown[] };
type SystemPromptSetData = { text?: unknown };
type MessageData = { content?: string | unknown[]; thinkingBlocks?: ThinkingBlock[] };
type ToolUseData = { toolCallId?: unknown; name?: unknown; input?: unknown; result?: ToolResult };
type PreservedMessage = {
  role?: unknown;
  content?: unknown;
  toolCalls?: CoreToolCall[];
  toolResults?: CoreToolResult[];
  thinkingBlocks?: ThinkingBlock[];
};

/**
 * Extracts text content from an array of content blocks.
 * Filters for text type blocks and joins them with newlines.
 */
function extractTextFromContentBlocks(content: unknown[]): string {
  if (!Array.isArray(content)) return '';
  return content
    .filter((b): b is TextBlock => {
      if (!b || typeof b !== 'object') return false;
      const block = b as ContentBlockShape;
      return block.type === 'text' && typeof block.text === 'string';
    })
    .map((b) => b.text)
    .join('\n');
}

/**
 * Extracts content blocks from a prompt, preserving both text and image blocks.
 * Returns string if only text blocks, array if any images present.
 */
function extractContentBlocks(content: unknown): string | ContentBlock[] {
  if (!Array.isArray(content)) {
    if (typeof content === 'string') return content;
    return '';
  }

  const blocks: ContentBlock[] = [];
  let hasImages = false;

  for (const block of content) {
    if (!block || typeof block !== 'object') continue;

    const b = block as Record<string, unknown>;
    if (b.type === 'text' && typeof b.text === 'string') {
      blocks.push({ type: 'text', text: b.text });
    } else if (b.type === 'image' && b.source && typeof b.source === 'object') {
      const source = b.source as Record<string, unknown>;
      if (
        source.type === 'base64' &&
        typeof source.media_type === 'string' &&
        typeof source.data === 'string'
      ) {
        blocks.push({
          type: 'image',
          source: {
            type: 'base64',
            media_type: source.media_type,
            data: source.data,
          },
        });
        hasImages = true;
      }
    }
  }

  // If no images, return simple string for backward compatibility
  if (!hasImages) {
    return blocks
      .filter((b): b is { type: 'text'; text: string } => b.type === 'text')
      .map((b) => b.text)
      .join('\n');
  }

  return blocks;
}

/**
 * Defensive post-pass for messages rebuilt from a `context_compacted` event's
 * `preserved` array. Strips orphan tool blocks in BOTH directions, because
 * Anthropic 400s on either:
 *
 *   - a `user` message carrying a `tool_result` whose `tool_use_id` has no
 *     matching `tool_use` in the immediately-prior `assistant` message
 *     ("unexpected tool_use_id found in tool_result blocks"); and
 *
 *   - an `assistant` message carrying a `tool_use` whose `id` has no matching
 *     `tool_result` in the immediately-following `user` message
 *     ("tool_use ids were found without tool_result blocks immediately after").
 *
 * Both halves of the pair can be lost to compaction, so we run two passes:
 * one that strips orphan user-side tool_results, one that strips orphan
 * assistant-side tool_uses. If a message ends up empty after a strip (no
 * tool blocks left and no text content), it is dropped entirely.
 *
 * Every dropped tool block is logged at WARN with its toolCallId so we can
 * spot compaction bugs that produce these orphans in the first place.
 */
function dropOrphanedToolBlocks(messages: ProviderMessage[]): void {
  // Pass A: drop user tool_results whose matching tool_use is missing in the
  // immediately-prior assistant message.
  for (let i = messages.length - 1; i >= 0; i--) {
    const m = messages[i]!;
    if (m.role !== 'user' || !Array.isArray(m.toolResults) || m.toolResults.length === 0) continue;

    const prev = messages[i - 1];
    const priorCallIds =
      prev && prev.role === 'assistant' && Array.isArray(prev.toolCalls)
        ? new Set(prev.toolCalls.map((c) => c.id))
        : new Set<string>();

    const kept: CoreToolResult[] = [];
    for (const tr of m.toolResults) {
      if (tr.id && priorCallIds.has(tr.id)) {
        kept.push(tr);
      } else {
        logger.warn('Dropping orphaned tool_result from compacted preserved messages', {
          toolCallId: tr.id ?? '<missing>',
        });
      }
    }

    if (kept.length === 0) {
      delete m.toolResults;
    } else {
      m.toolResults = kept;
    }

    const contentIsEmpty =
      typeof m.content === 'string' ? m.content.trim() === '' : m.content.length === 0;
    const hasToolResults = Array.isArray(m.toolResults) && m.toolResults.length > 0;
    if (!hasToolResults && contentIsEmpty) {
      messages.splice(i, 1);
    }
  }

  // Pass B: drop assistant tool_uses whose matching tool_result is missing in
  // the immediately-following user message. This is the case that bricked Ada
  // a tool_use survived into messages[1298] without its result and
  // every subsequent Anthropic call 400'd on the same toolu id.
  for (let i = messages.length - 1; i >= 0; i--) {
    const m = messages[i]!;
    if (m.role !== 'assistant' || !Array.isArray(m.toolCalls) || m.toolCalls.length === 0) continue;

    const next = messages[i + 1];
    const followingResultIds =
      next && next.role === 'user' && Array.isArray(next.toolResults)
        ? new Set(next.toolResults.map((r) => r.id))
        : new Set<string>();

    const kept: CoreToolCall[] = [];
    for (const tc of m.toolCalls) {
      if (tc.id && followingResultIds.has(tc.id)) {
        kept.push(tc);
      } else {
        logger.warn('Dropping orphaned tool_use from compacted preserved messages', {
          toolCallId: tc.id ?? '<missing>',
        });
      }
    }

    if (kept.length === 0) {
      delete m.toolCalls;
    } else {
      m.toolCalls = kept;
    }

    const contentIsEmpty =
      typeof m.content === 'string' ? m.content.trim() === '' : m.content.length === 0;
    const hasToolCalls = Array.isArray(m.toolCalls) && m.toolCalls.length > 0;
    if (!hasToolCalls && contentIsEmpty) {
      messages.splice(i, 1);
    }
  }
}

/**
 * Return type for buildProviderMessagesFromDurableEvents.
 * systemPrompt is sourced exclusively from system_prompt_set events (last one wins).
 * Sessions without a system_prompt_set event return systemPrompt: ''.
 */
export type BuiltProviderMessages = {
  messages: ProviderMessage[];
  systemPrompt: string;
};

/**
 * Builds provider messages from durable events stored in a session directory.
 * Reconstructs the conversation history by reading and parsing events.jsonl.
 * Exported for testing - converts durable events to provider message format.
 *
 * Returns { messages, systemPrompt } where systemPrompt comes from the last
 * system_prompt_set event in the log. Sessions without a system_prompt_set
 * event return systemPrompt: ''.
 */
export function buildProviderMessagesFromDurableEvents(sessionDir: string): BuiltProviderMessages {
  return buildProviderMessagesFromParsedEvents(readParsedSessionEvents(sessionDir));
}

/**
 * The mutable accumulator shared by the full rebuild and the incremental
 * projection fold. It carries the FoldState plus the system-prompt last-wins
 * value and the era-scoped system_prompt_set count. ONE per-event handler
 * (applyEventToProjection) mutates it, so the full builder and the incremental
 * tail-fold cannot diverge.
 */
export type ProjectionAccumulator = {
  state: FoldState;
  systemPrompt: string;
  // system_prompt_set count scoped to the current compaction era; reset to 0 at
  // each context_compacted (a post-compaction rerender legitimately supersedes
  // earlier prompts). >1 within one era is the invariant violation worth surfacing.
  systemPromptCount: number;
};

export function newProjectionAccumulator(): ProjectionAccumulator {
  return { state: initialFoldState(), systemPrompt: '', systemPromptCount: 0 };
}

/**
 * Apply ONE durable event to the projection accumulator. This is the single
 * per-event handler used by BOTH buildProviderMessagesFromParsedEvents (full
 * rebuild) and foldTailIntoProjection (incremental tail). It folds
 * `prompt`/`message`/`tool_use` through the shared reducer (foldEvent) — the
 * Anthropic parallel-tool form — and layers the rebuild-only concerns:
 * system-prompt last-wins + era count (with reset on compaction), the empty-
 * `prompt` drop + image-preserving extraction, the `context_injected`
 * text-merge, and the `context_compacted` reset to `preserved[]` + orphan
 * cleanup. Those special events close the open tool batch (batch = null) so a
 * following tool_use starts a fresh turn rather than coalescing across them.
 */
export function applyEventToProjection(
  acc: ProjectionAccumulator,
  type: string,
  data: Record<string, unknown>
): void {
  if (type === 'system_prompt_set') {
    const eventData = data as SystemPromptSetData;
    if (typeof eventData.text === 'string') {
      acc.systemPrompt = eventData.text; // last-one-wins for defensive multi-event support
      acc.systemPromptCount++;
    }
    return;
  }

  if (type === 'prompt') {
    const content = extractContentBlocks((data as Record<string, unknown>).content);
    // Check if content is non-empty (string or array with items)
    const hasContent = typeof content === 'string' ? content.trim() : content.length > 0;
    // Image-preserving extraction (NOT text-flatten); drop empty prompts. Feed
    // the already-extracted content through the reducer verbatim so the empty
    // drop and image preservation stay rebuild-only concerns.
    if (hasContent) {
      acc.state = foldEvent(acc.state, { type: 'prompt', data: { content } });
    }
    return;
  }

  if (type === 'context_injected') {
    // Emit as role:user so runtime context is visible to the model in the
    // messages array while remaining distinct from the stable system prompt.
    // Use appendOrMergeUser to avoid consecutive role:user messages when the
    // prior entry is a user[toolResults] turn — Anthropic combines consecutive
    // same-role messages in implementation-defined ways and it disrupts cache reach.
    const eventData = data as ContextInjectedData;
    const contentArr = Array.isArray(eventData.content) ? eventData.content : [];
    const content = extractTextFromContentBlocks(contentArr);
    if (content.trim()) {
      const merged = appendOrMergeUser(acc.state.messages, content);
      acc.state = { messages: merged, batch: null };
    } else {
      // Even an empty injection closes the open tool batch.
      acc.state = { messages: acc.state.messages, batch: null };
    }
    return;
  }

  if (type === 'context_compacted') {
    const eventData = data as ContextCompactedData;
    const preserved = Array.isArray(eventData.preserved) ? eventData.preserved : [];

    // A compaction starts a new era: the warning count resets so a post-
    // compaction rerender of the system prompt is not flagged as a violation.
    acc.systemPromptCount = 0;

    // Reset messages: the preserved array is the complete rebuilt conversation.
    // The summary (if any) lives as the first entry in preserved — placed there
    // by the compaction strategy. Do not read a separate data.summary field;
    // no writer produces that shape and doing so would duplicate the summary.
    const messages: ProviderMessage[] = [];

    for (const msg of preserved) {
      if (!msg || typeof msg !== 'object') continue;
      const msgObj = msg as PreservedMessage;
      const role = msgObj.role;
      const rawContent = msgObj.content;
      if (role !== 'user' && role !== 'assistant' && role !== 'system') continue;
      // Accept string or ContentBlock[] — images in tail preserved messages
      // must be passed through without being flattened to text.
      if (typeof rawContent !== 'string' && !Array.isArray(rawContent)) continue;
      const content: string | ContentBlock[] = Array.isArray(rawContent)
        ? (rawContent as ContentBlock[])
        : rawContent;

      const toolCalls = Array.isArray(msgObj.toolCalls) ? msgObj.toolCalls : undefined;
      const toolResults = Array.isArray(msgObj.toolResults) ? msgObj.toolResults : undefined;
      const thinkingBlocks = Array.isArray(msgObj.thinkingBlocks)
        ? msgObj.thinkingBlocks
        : undefined;

      messages.push({
        role,
        content,
        ...(toolCalls ? { toolCalls } : {}),
        ...(toolResults ? { toolResults } : {}),
        ...(thinkingBlocks && thinkingBlocks.length > 0 ? { thinkingBlocks } : {}),
      });
    }

    dropOrphanedToolBlocks(messages);

    // The reset replaces the whole conversation and closes the open tool batch.
    acc.state = { messages, batch: null };
    return;
  }

  if (type === 'message') {
    const eventData = data as MessageData;
    // Verbatim content (string or ContentBlock[]) — the reducer never flattens
    // images. For the common text-only assistant turn the converter extracts the
    // same text either way, so this is shape-preserving for the wire.
    acc.state = foldEvent(acc.state, {
      type: 'message',
      data: {
        content: eventData.content,
        ...(Array.isArray(eventData.thinkingBlocks)
          ? { thinkingBlocks: eventData.thinkingBlocks }
          : {}),
      },
    });
    return;
  }

  if (type === 'tool_use') {
    const eventData = data as ToolUseData;
    acc.state = foldEvent(acc.state, {
      type: 'tool_use',
      data: {
        toolCallId: eventData.toolCallId,
        name: eventData.name,
        input: eventData.input,
        result: eventData.result,
      },
    });
    return;
  }
}

/**
 * Surface the system_prompt_set invariant violation (more than one in the
 * current compaction era). Called once after a whole-log build/seed examines
 * every event; NOT called on the incremental tail-fold hot path.
 */
export function finalizeProjectionWarnings(acc: ProjectionAccumulator): void {
  if (acc.systemPromptCount > 1) {
    // More than one system_prompt_set since the last compaction (or session
    // start). The "written once" invariant has been violated within this era.
    // We still use the last value (defensive last-wins), but we surface the
    // violation so it can be investigated.
    logger.warn('Multiple system_prompt_set events found in session — invariant violation', {
      count: acc.systemPromptCount,
    });
  }
}

/**
 * Pure core of buildProviderMessagesFromDurableEvents: derives the provider
 * message prefix + system prompt from an already-parsed, sorted event array.
 * No I/O — the durable log is read and parsed once by the caller. Folds every
 * event through the shared applyEventToProjection handler.
 */
export function buildProviderMessagesFromParsedEvents(
  parsedEvents: ParsedSessionEvent[]
): BuiltProviderMessages {
  const acc = newProjectionAccumulator();
  for (const e of parsedEvents) applyEventToProjection(acc, e.type, e.data);
  finalizeProjectionWarnings(acc);
  return { messages: acc.state.messages, systemPrompt: acc.systemPrompt };
}

/**
 * Estimates the total token count for a set of provider messages.
 * Includes content tokens, tool call tokens, and tool result tokens.
 */
export function estimateProviderTokens(messages: ProviderMessage[]): number {
  let total = 0;
  for (const message of messages) {
    if (typeof message.content === 'string') {
      total += estimateTokens(message.content);
    } else {
      // Count tokens for text blocks only (images don't count as text tokens)
      const textContent = message.content
        .filter((b): b is ContentBlock & { type: 'text' } => b.type === 'text')
        .map((b) => b.text)
        .join('\n');
      total += estimateTokens(textContent);
    }
    if (message.toolCalls) {
      total += estimateTokens(JSON.stringify(message.toolCalls));
    }
    if (message.toolResults) {
      total += estimateTokens(JSON.stringify(message.toolResults));
    }
  }
  return total;
}
