import { createHash } from "node:crypto";

export type ContextViewMode = "observe" | "safe" | "aggressive";

export interface ContextMessage {
  role: string;
  content?: any;
  toolCallId?: string;
  toolName?: string;
  isError?: boolean;
}

export interface ContextViewTelemetry {
  supersededResults: number;
  shakenResults: number;
  fencedBlocks: number;
  estimatedTokensSaved: number;
  cachePrefixChanged: boolean;
}

export interface ContextViewInput<T extends ContextMessage> {
  messages: T[];
  mode: ContextViewMode;
  contextTokens: number;
  contextWindow: number;
  cacheAgeMs: number;
  cacheTtlMs: number;
  manual?: "tools" | "all";
}

export interface ContextViewResult<T extends ContextMessage> {
  messages: T[];
  telemetry: ContextViewTelemetry;
}

export interface ContextViewPolicy {
  apply<T extends ContextMessage>(input: ContextViewInput<T>): ContextViewResult<T>;
}

const CHARACTERS_PER_TOKEN = 4;

function textFromContent(content: unknown): string {
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";
  return content
    .filter((block) => block?.type === "text" && typeof block.text === "string")
    .map((block) => block.text as string)
    .join("\n");
}

function estimateMessageTokens(message: ContextMessage): number {
  const text = textFromContent(message.content);
  return Math.ceil((text.length + JSON.stringify(message).length - text.length) / CHARACTERS_PER_TOKEN);
}

function replaceText(message: ContextMessage, replacement: string): void {
  message.content = [{ type: "text", text: replacement }];
}

const FENCED_BLOCK = /(```|~~~)[^\n]*\n[\s\S]*?\1/g;
const MANUAL_BLOCK_TOKEN_THRESHOLD = 12_000;

export function historyBlockId(block: string): string {
  return createHash("sha256").update(block).digest("hex").slice(0, 24);
}

export function fencedHistoryBlocks(content: unknown): Array<{ id: string; text: string }> {
  const texts = typeof content === "string"
    ? [content]
    : Array.isArray(content)
      ? content
          .filter((block) => block?.type === "text" && typeof block.text === "string")
          .map((block) => block.text as string)
      : [];
  return texts.flatMap((text) =>
    [...text.matchAll(FENCED_BLOCK)].map((match) => ({
      id: historyBlockId(match[0]),
      text: match[0],
    })),
  );
}

function slimFencedBlocks(message: ContextMessage): {
  blocks: number;
  tokensSaved: number;
} {
  let blocks = 0;
  let tokensSaved = 0;
  const slimText = (text: string): string =>
    text.replace(FENCED_BLOCK, (block) => {
      const before = Math.ceil(block.length / CHARACTERS_PER_TOKEN);
      if (
        before <= MANUAL_BLOCK_TOKEN_THRESHOLD ||
        block.includes("cost-history-block://")
      ) {
        return block;
      }
      const replacement = `[Fenced block slimmed ~${before.toLocaleString()} tokens; recover: cost-history-block://${historyBlockId(block)}]`;
      blocks += 1;
      tokensSaved += Math.max(
        0,
        before - Math.ceil(replacement.length / CHARACTERS_PER_TOKEN),
      );
      return replacement;
    });

  if (typeof message.content === "string") {
    message.content = slimText(message.content);
  } else if (Array.isArray(message.content)) {
    message.content = message.content.map((block) =>
      block?.type === "text" && typeof block.text === "string"
        ? { ...block, text: slimText(block.text) }
        : block,
    );
  }
  return { blocks, tokensSaved };
}

export function createContextViewPolicy(): ContextViewPolicy {
  return {
    apply(input) {
      const messages = structuredClone(input.messages) as typeof input.messages;
      const toolCalls = new Map<string, { name: string; path?: string }>();

      for (const message of messages) {
        if (message.role !== "assistant" || !Array.isArray(message.content)) continue;
        for (const block of message.content) {
          if (block?.type !== "toolCall" || typeof block.id !== "string") continue;
          toolCalls.set(block.id, {
            name: String(block.name ?? ""),
            path:
              typeof block.arguments?.path === "string"
                ? block.arguments.path
                : undefined,
          });
        }
      }

      const suffixTokens = new Array<number>(messages.length);
      let accumulated = 0;
      for (let index = messages.length - 1; index >= 0; index--) {
        suffixTokens[index] = accumulated;
        accumulated += estimateMessageTokens(messages[index]);
      }

      const telemetry: ContextViewTelemetry = {
        supersededResults: 0,
        shakenResults: 0,
        fencedBlocks: 0,
        estimatedTokensSaved: 0,
        cachePrefixChanged: false,
      };
      const seenReadPaths = new Set<string>();
      const warmSuffixLimit = input.mode === "aggressive" ? 16_000 : 8_000;
      const cacheIsCold = input.cacheAgeMs >= input.cacheTtlMs;

      for (let index = messages.length - 1; index >= 0; index--) {
        const message = messages[index];
        if (
          message.role !== "toolResult" ||
          message.toolName !== "read" ||
          message.isError === true ||
          typeof message.toolCallId !== "string"
        ) {
          continue;
        }
        const call = toolCalls.get(message.toolCallId);
        if (!call?.path) continue;
        if (!seenReadPaths.has(call.path)) {
          seenReadPaths.add(call.path);
          continue;
        }
        if (!cacheIsCold && suffixTokens[index] > warmSuffixLimit) continue;

        const before = estimateMessageTokens(message);
        const replacement = `[Superseded by a newer read of ${call.path}; recover: cost-history://${message.toolCallId}]`;
        const after = Math.ceil(replacement.length / CHARACTERS_PER_TOKEN);
        telemetry.supersededResults += 1;
        telemetry.estimatedTokensSaved += Math.max(0, before - after);
        telemetry.cachePrefixChanged = true;
        if (input.mode !== "observe") replaceText(message, replacement);
      }

      const triggerRatio = input.mode === "aggressive" ? 0.55 : 0.65;
      if (
        input.manual !== undefined ||
        (input.contextWindow > 0 &&
          input.contextTokens / input.contextWindow >= triggerRatio)
      ) {
        const protectedRecentTokens = input.manual
          ? 0
          : input.mode === "aggressive"
            ? 16_000
            : Math.min(64_000, Math.floor(input.contextWindow * 0.25));
        const minimumSavings = input.manual
          ? 1
          : input.mode === "aggressive"
            ? 4_000
            : Math.floor(input.contextWindow * 0.08);
        const candidates: Array<{
          message: ContextMessage;
          toolCallId: string;
          beforeTokens: number;
          replacement: string;
          savedTokens: number;
        }> = [];

        for (let index = 0; index < messages.length; index++) {
          const message = messages[index];
          if (
            message.role !== "toolResult" ||
            !["bash", "read", "grep", "find"].includes(message.toolName ?? "") ||
            message.isError === true ||
            typeof message.toolCallId !== "string" ||
            suffixTokens[index] < protectedRecentTokens
          ) {
            continue;
          }
          const original = textFromContent(message.content);
          if (
            original.length === 0 ||
            original.startsWith("[Superseded by") ||
            original.includes("cost-artifact://")
          ) {
            continue;
          }
          const beforeTokens = Math.ceil(original.length / CHARACTERS_PER_TOKEN);
          const replacement = `[Context slimmed ~${beforeTokens.toLocaleString()} tokens; recover: cost-history://${message.toolCallId}]`;
          const afterTokens = Math.ceil(replacement.length / CHARACTERS_PER_TOKEN);
          candidates.push({
            message,
            toolCallId: message.toolCallId,
            beforeTokens,
            replacement,
            savedTokens: Math.max(0, beforeTokens - afterTokens),
          });
        }

        const totalSavings = candidates.reduce(
          (total, candidate) => total + candidate.savedTokens,
          0,
        );
        if (totalSavings >= minimumSavings) {
          telemetry.shakenResults += candidates.length;
          telemetry.estimatedTokensSaved += totalSavings;
          telemetry.cachePrefixChanged = candidates.length > 0 || telemetry.cachePrefixChanged;
          if (input.mode !== "observe") {
            for (const candidate of candidates) {
              replaceText(candidate.message, candidate.replacement);
            }
          }
        }
      }

      if (input.manual === "all") {
        for (const message of messages) {
          if (message.role !== "user" && message.role !== "assistant") continue;
          const slimmed = slimFencedBlocks(message);
          telemetry.fencedBlocks += slimmed.blocks;
          telemetry.estimatedTokensSaved += slimmed.tokensSaved;
          telemetry.cachePrefixChanged =
            slimmed.blocks > 0 || telemetry.cachePrefixChanged;
        }
      }

      return { messages, telemetry };
    },
  };
}