/**
 * Image-rejection recovery transforms for the image-recovery plugin.
 *
 * When the provider rejects a turn because an attached image violates its
 * limits — too large on a side, over the payload cap, or below the minimum
 * size floor — the `post-model-call` hook rewrites the offending image blocks
 * in the working history and asks the loop to retry. {@link
 * recoverUnsendableImages} performs that in-memory transform for the immediate
 * retry; {@link persistUnsendableImageDowngrades} makes the same rewrite
 * durable, because the stored message row otherwise keeps the rejected image
 * block and it rehydrates on every later turn and keeps re-entering the
 * model's context. The durable rewrite replaces the unsendable block with its
 * resized form, or with a text note when it cannot be resized on this host,
 * so a rejected image cannot resurface and re-reject on every later turn.
 */

import {
  type ContentBlock,
  type Message,
  resolveMediaSourceData,
} from "@vellumai/plugin-api";

import {
  isBelowMinDimension,
  optimizeImageForTransport,
  upscaleImageToMinimum,
} from "../../../agent/image-optimize.js";
import { parseImageDimensions } from "../../../context/image-dimensions.js";
import {
  getMessages,
  updateMessageContent,
} from "../../../persistence/conversation-crud.js";
import { sniffBase64ImageMimeType } from "../../../util/image-conversion.js";
import { getLogger } from "../../../util/logger.js";

const log = getLogger("image-recovery");

// Anthropic rejects any image whose longest side exceeds this many pixels,
// regardless of payload size. Mirrors the user-facing message surfaced by
// `classifyConversationError` for the IMAGE_TOO_LARGE code.
// https://docs.anthropic.com/en/docs/build-with-claude/vision#image-size
const PROVIDER_MAX_IMAGE_DIMENSION = 8000;

// Anthropic rejects any single image whose base64 payload exceeds 5 MB.
// https://docs.anthropic.com/en/docs/build-with-claude/vision#image-size
const PROVIDER_MAX_IMAGE_PAYLOAD_BYTES = 5 * 1024 * 1024;

/**
 * Note left in place of an image that cannot be sent to the provider. Shared
 * with the in-memory recovery path so the persisted history matches what the
 * model saw on the turn the image was rejected.
 */
export const UNSENDABLE_IMAGE_NOTE =
  "(An image was attached but could not be sent — it does not meet the provider's image size limits and automatic resizing was not available. Please resize the image and try again.)";

/**
 * Replacement for an image that violates a provider hard limit (per-side pixel
 * cap, payload size, the minimum-size floor, or a declared media type that
 * disagrees with the actual bytes), or null when the image is within limits
 * and should be left untouched. Gating on the provider hard caps is what
 * keeps still-sendable images intact: a normally sized image is left alone
 * rather than being noted or needlessly rewritten.
 *
 * A mislabeled image (e.g. JPEG bytes declared `image/png` — clients derive
 * the MIME from the filename extension) is relabeled with its sniffed type,
 * bytes untouched. An unsendable image that can be resized is rewritten to
 * its resized form (downscaled when oversized, upscaled to the minimum floor
 * when undersized); one that cannot be resized on this host (resize is a
 * no-op, e.g. `sips` is absent off macOS or the format is unsupported) is
 * replaced with a text note.
 *
 * Shared by the in-memory recovery transform and the durable persist pass so
 * both apply the identical rule. Persisting the resized form is what lets a
 * poisoned conversation durably self-heal — the latest tool-result media is kept
 * in context, so without it the original rejected block rehydrates and
 * re-rejects on every later turn.
 */
export async function unsendableImageReplacement(
  block: Extract<ContentBlock, { type: "image" }>,
): Promise<ContentBlock | null> {
  // Resolve reference sources to their bytes so a reloaded (referenced) image
  // is gated on the same payload/dimension caps as an inline one. When the
  // attachment can no longer be read, leave the block untouched.
  const resolved = resolveMediaSourceData(block.source);
  if (!resolved) {
    return null;
  }
  const sniffed = sniffBase64ImageMimeType(resolved.data);
  const mediaTypeMismatch = sniffed != null && sniffed !== resolved.media_type;
  // The sniffed type also drives the resize paths: dimension parsing keyed on
  // the wrong declared type fails, which would misroute a mislabeled image to
  // the unsendable note.
  const effectiveMediaType = sniffed ?? resolved.media_type;
  const payloadBytes = resolved.data.length;
  const dims = mediaTypeMismatch
    ? parseImageDimensions(resolved.data, effectiveMediaType)
    : parseImageDimensions(block.source);
  const exceedsDimensionCap =
    dims != null &&
    (dims.width > PROVIDER_MAX_IMAGE_DIMENSION ||
      dims.height > PROVIDER_MAX_IMAGE_DIMENSION);
  const exceedsPayloadCap = payloadBytes > PROVIDER_MAX_IMAGE_PAYLOAD_BYTES;
  const belowMinDimension = isBelowMinDimension(dims);
  if (!exceedsDimensionCap && !exceedsPayloadCap && !belowMinDimension) {
    if (mediaTypeMismatch) {
      // Only the label is wrong, so keep the source shape: a workspace_ref
      // stays a reference (inlining it would bake the full payload into the
      // stored message row) and only media_type is corrected.
      return {
        type: "image",
        source: { ...block.source, media_type: effectiveMediaType },
      };
    }
    return null;
  }

  // The floor is undocumented, so undersized images are never touched
  // pre-send — the upscale runs only here, in response to an actual
  // provider rejection. Oversized images reuse the transport downscale.
  const optimized = belowMinDimension
    ? await upscaleImageToMinimum(resolved.data, effectiveMediaType)
    : await optimizeImageForTransport(resolved.data, effectiveMediaType);
  if (optimized && optimized.data !== resolved.data) {
    return {
      type: "image",
      source: {
        type: "base64",
        media_type: optimized.mediaType,
        data: optimized.data,
      },
    };
  }
  return { type: "text", text: UNSENDABLE_IMAGE_NOTE };
}

/**
 * Rewrite every stored message in a conversation that holds an image the
 * provider rejects — whether a top-level block or one nested in a
 * tool_result's contentBlocks — replacing it with its resized form, or with
 * {@link UNSENDABLE_IMAGE_NOTE} when it cannot be resized on this host. Reads
 * stored content directly (not the in-memory, injection-enriched copy) so
 * injected prefixes and hydrated source paths are never written back.
 *
 * Idempotent: a resized image is within limits and a note is no longer an
 * image, so neither matches on a second run. Returns the number of rewritten
 * messages.
 */
export async function persistUnsendableImageDowngrades(
  conversationId: string,
): Promise<number> {
  let rewritten = 0;
  for (const row of getMessages(conversationId)) {
    const hasImageBlock = row.content.some(
      (b) =>
        b.type === "image" ||
        (b.type === "tool_result" &&
          b.contentBlocks?.some((cb) => cb.type === "image")),
    );
    if (!hasImageBlock) {
      continue;
    }

    const parsed = row.content;

    let changed = false;
    const next: ContentBlock[] = [];
    for (const block of parsed as ContentBlock[]) {
      if (block.type === "image") {
        const replacement = await unsendableImageReplacement(block);
        if (!replacement) {
          next.push(block);
          continue;
        }
        changed = true;
        next.push(replacement);
        continue;
      }
      // Images returned by a tool (e.g. browser_screenshot) live inside the
      // tool_result's contentBlocks, not as top-level blocks. Downgrade them
      // in place so the tool_use/tool_result pairing stays intact.
      if (block.type === "tool_result" && block.contentBlocks?.length) {
        let nestedChanged = false;
        const contentBlocks: ContentBlock[] = [];
        for (const cb of block.contentBlocks) {
          if (cb.type !== "image") {
            contentBlocks.push(cb);
            continue;
          }
          const replacement = await unsendableImageReplacement(cb);
          if (!replacement) {
            contentBlocks.push(cb);
            continue;
          }
          nestedChanged = true;
          contentBlocks.push(replacement);
        }
        if (!nestedChanged) {
          next.push(block);
          continue;
        }
        changed = true;
        next.push({ ...block, contentBlocks });
        continue;
      }
      next.push(block);
    }
    if (!changed) {
      continue;
    }

    updateMessageContent(row.id, JSON.stringify(next));
    rewritten++;
    log.info(
      { conversationId, messageId: row.id },
      "Persisted unsendable-image downgrade so it cannot resurface on later turns",
    );
  }
  return rewritten;
}

/**
 * True when a message's content holds an image the provider may have rejected
 * — either a top-level image block (user upload) or one nested inside a
 * tool_result's contentBlocks (e.g. a browser screenshot).
 */
function messageHasImageBlock(content: ReadonlyArray<ContentBlock>): boolean {
  return content.some(
    (b) =>
      b.type === "image" ||
      (b.type === "tool_result" &&
        (b.contentBlocks?.some((cb) => cb.type === "image") ?? false)),
  );
}

/**
 * Resize every unsendable image in the working history for an immediate
 * retry, leaving still-sendable images untouched. Recovers both top-level image
 * blocks (user uploads) and images nested inside a tool_result's contentBlocks
 * (e.g. a browser screenshot) in place, so the tool_use/tool_result pairing
 * stays intact rather than dropping the whole tool_result. Applies the same
 * provider-cap gate as {@link persistUnsendableImageDowngrades} so the in-memory
 * retry and the durable rewrite agree on which images are unsendable.
 *
 * Reports whether any block was actually rewritten. A provider rejection can
 * be classified as recoverable yet leave nothing to fix — e.g. a media-type
 * mismatch on a format {@link sniffBase64ImageMimeType} cannot identify (a
 * renamed BMP/TIFF/SVG, or an unconvertible HEIF) yields no replacement. The
 * caller must not retry an unchanged history: it would resend the identical
 * rejected image and falsely report a correction.
 */
export async function recoverUnsendableImages(
  messages: ReadonlyArray<Message>,
): Promise<{
  messages: Message[];
  changed: boolean;
}> {
  let changed = false;
  const recovered: Message[] = [];
  for (const msg of messages) {
    if (!Array.isArray(msg.content) || !messageHasImageBlock(msg.content)) {
      recovered.push(msg);
      continue;
    }
    const content: ContentBlock[] = [];
    for (const b of msg.content) {
      if (b.type === "image") {
        const replacement = await unsendableImageReplacement(b);
        if (replacement) {
          changed = true;
          content.push(replacement);
        } else {
          content.push(b);
        }
        continue;
      }
      if (b.type === "tool_result" && b.contentBlocks?.length) {
        const contentBlocks: ContentBlock[] = [];
        for (const cb of b.contentBlocks) {
          if (cb.type !== "image") {
            contentBlocks.push(cb);
            continue;
          }
          const replacement = await unsendableImageReplacement(cb);
          if (replacement) {
            changed = true;
            contentBlocks.push(replacement);
          } else {
            contentBlocks.push(cb);
          }
        }
        content.push({ ...b, contentBlocks });
        continue;
      }
      content.push(b);
    }
    recovered.push({ ...msg, content });
  }
  return { messages: recovered, changed };
}