import { writeFile } from "node:fs/promises";
import type {
  AgentToolResult,
  ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import {
  convertToPng,
  formatDimensionNote,
  resizeImage,
} from "@earendil-works/pi-coding-agent";
import type { KernelToHostMessage } from "../bridge/protocol.ts";
import type { TruncationMeta } from "../output/output-meta.ts";
import {
  artifactNotice,
  DEFAULT_MAX_BYTES,
  DEFAULT_MAX_LINES,
  formatMiddleElisionMarker,
  OutputSink,
  type OutputSummary,
  TailBuffer,
  truncateHeadBytes,
  truncateTail,
} from "../output/streaming-output.ts";

const MAX_DISPLAY_TEXT_BYTES = 8000;
const MAX_DISPLAY_IMAGES = 8;
const MAX_JSON_OUTPUTS = 8;
const MAX_DISPLAY_PAYLOAD_BYTES = 4 * 1024 * 1024;
/** Failure budget: bounded visible output for failed cells (Step B1). */
const FAILURE_OUTPUT_MAX_BYTES = 20_000;
/** Marker + separator room reserved so the elided body stays in budget. */
const FAILURE_MARKER_RESERVE_BYTES = 64;

export interface EvalImageContent {
  readonly data: string;
  readonly mimeType: string;
  readonly type: "image";
}

export interface EvalImageResizeResult {
  readonly dimensionNote?: string;
  readonly image: EvalImageContent;
}

export type EvalImageResizer = (
  image: EvalImageContent,
  model: ExtensionContext["model"]
) => Promise<EvalImageResizeResult>;

export interface EvalOutputOptions {
  readonly artifactPath?: string;
  /** Coalesce streaming preview updates; see CellResultBuilderOptions.chunkThrottleMs. */
  readonly chunkThrottleMs?: number;
  readonly headBytes: number;
  readonly imageResizer?: EvalImageResizer;
  readonly maxColumns: number;
  readonly model: ExtensionContext["model"];
  readonly onChunk: (aggregateText: string, cellText: string) => void;
}

export interface EvalOutputResult {
  readonly hasMarkdown: boolean;
  readonly images: readonly EvalImageContent[];
  readonly jsonOutputs: readonly unknown[];
  readonly meta?: TruncationMeta;
  readonly notice?: string;
  readonly output: string;
  readonly truncated: boolean;
}

type DisplayMessage = Extract<KernelToHostMessage, { type: "display" }>;
type WebpModel =
  | { readonly provider: string; readonly api: string }
  | undefined;

class DisplayPayloadError extends Error {
  readonly name = "DisplayPayloadError";

  constructor(mimeType: string, options: { cause: SyntaxError }) {
    super(`Invalid ${mimeType} display payload: ${options.cause.message}`, options);
  }
}

export class EvalOutputCollector {
  readonly #options: EvalOutputOptions;
  readonly #sink: OutputSink;
  readonly #aggregateTail = new TailBuffer(DEFAULT_MAX_BYTES * 2);
  readonly #cellTail = new TailBuffer(DEFAULT_MAX_BYTES * 2);
  readonly #displayImages: EvalImageContent[] = [];
  readonly #images: EvalImageContent[] = [];
  readonly #jsonOutputs: unknown[] = [];
  #hasMarkdown = false;
  #imagesProcessed = false;
  #displayImagesOmitted = 0;
  #jsonOutputsOmitted = 0;
  #jsonDisplayCount = 0;

  constructor(options: EvalOutputOptions) {
    this.#options = options;
    this.#sink = new OutputSink({
      artifactPath: options.artifactPath,
      headBytes: options.headBytes,
      maxColumns: options.maxColumns,
      chunkThrottleMs: options.chunkThrottleMs,
      onChunk: (chunk) => {
        this.#aggregateTail.append(chunk);
        this.#cellTail.append(chunk);
        options.onChunk(this.#aggregateTail.text(), this.#cellTail.text());
      },
    });
  }

  push(text: string): void {
    this.#sink.push(text);
  }

  display(message: DisplayMessage): void {
    if (message.mimeType.startsWith("image/")) {
      if (
        this.#displayImages.length >= MAX_DISPLAY_IMAGES ||
        Buffer.byteLength(message.dataBase64, "utf8") >
          MAX_DISPLAY_PAYLOAD_BYTES
      ) {
        this.#displayImagesOmitted += 1;
        this.#sink.push(
          `display images omitted: ${this.#displayImagesOmitted}\n`
        );
        return;
      }
      this.#displayImages.push({
        type: "image",
        mimeType: message.mimeType,
        data: message.dataBase64,
      });
      return;
    }
    const text = Buffer.from(message.dataBase64, "base64").toString("utf8");
    if (message.mimeType === "application/json") {
      let value: unknown;
      try {
        value = JSON.parse(text);
      } catch (error) {
        if (error instanceof SyntaxError) {
          throw new DisplayPayloadError(message.mimeType, { cause: error });
        }
        throw error;
      }
      this.#jsonDisplayCount += 1;
      if (
        this.#jsonOutputs.length >= MAX_JSON_OUTPUTS ||
        Buffer.byteLength(text, "utf8") > MAX_DISPLAY_PAYLOAD_BYTES
      ) {
        this.#jsonOutputsOmitted += 1;
        this.#sink.push(
          `display outputs omitted: ${this.#jsonOutputsOmitted}\n`
        );
      } else {
        this.#jsonOutputs.push(value);
      }
      this.#sink.push(
        `display[${this.#jsonDisplayCount}]:\n${formatDisplayJson(value)}\n`
      );
      return;
    }
    if (message.mimeType === "text/markdown") {
      this.#hasMarkdown = true;
    }
    this.#sink.push(text.endsWith("\n") ? text : `${text}\n`);
  }

  aggregateText(): string {
    return this.#aggregateTail.text();
  }

  async finish(isError: boolean): Promise<EvalOutputResult> {
    await this.#processImages();
    const summary = await this.#finalSummary();
    if (isError && summary.outputBytes > FAILURE_OUTPUT_MAX_BYTES) {
      return await this.#failureResult(summary);
    }
    const meta = truncationMetaFromSummary(summary);
    const notice =
      summary.artifactId === undefined
        ? undefined
        : artifactNotice(summary.artifactId);
    return {
      output: summary.output.trimEnd(),
      images: [...this.#images],
      jsonOutputs: [...this.#jsonOutputs],
      hasMarkdown: this.#hasMarkdown,
      truncated: summary.truncated,
      ...(notice === undefined ? {} : { notice }),
      ...(meta === undefined ? {} : { meta }),
    };
  }

  async #failureResult(summary: OutputSummary): Promise<EvalOutputResult> {
    // Guarantee the full output lands in the artifact when one is configured
    // and the sink did not already spill one.
    let artifactId = summary.artifactId;
    if (artifactId === undefined && this.#options.artifactPath !== undefined) {
      await writeFile(
        this.#options.artifactPath,
        this.#aggregateTail.text(),
        "utf8"
      );
      artifactId = this.#options.artifactPath;
    }
    const visible = failureBudgetOutput(summary, this.#options.headBytes);
    const meta: TruncationMeta = {
      direction: "middle",
      truncatedBy: "middle",
      totalLines: summary.totalLines,
      totalBytes: summary.totalBytes,
      outputLines: visible.outputLines,
      outputBytes: visible.outputBytes,
      ...(visible.headRange === undefined
        ? {}
        : { headRange: visible.headRange }),
      ...(visible.tailRange === undefined
        ? {}
        : { tailRange: visible.tailRange }),
      elidedBytes: visible.elidedBytes,
      elidedLines: visible.elidedLines,
      ...(artifactId === undefined ? {} : { artifactId }),
    };
    return {
      output: visible.text.trimEnd(),
      images: [...this.#images],
      jsonOutputs: [...this.#jsonOutputs],
      hasMarkdown: this.#hasMarkdown,
      truncated: true,
      ...(artifactId === undefined ? {} : { notice: artifactNotice(artifactId) }),
      meta,
    };
  }

  async flush(): Promise<void> {
    await this.#sink.dump();
  }

  async #processImages(): Promise<void> {
    if (this.#imagesProcessed) {
      return;
    }
    this.#imagesProcessed = true;
    const resize = this.#options.imageResizer ?? resizeEvalImage;
    for (const source of this.#displayImages) {
      const resized = await resize(source, this.#options.model);
      this.#images.push(resized.image);
      const description =
        resized.dimensionNote ?? `[${resized.image.mimeType}]`;
      this.#sink.push(`display image ${this.#images.length}: ${description}\n`);
    }
  }

  async #finalSummary(): Promise<OutputSummary> {
    const summary = await this.#sink.dump();
    if (summary.truncated || summary.totalLines <= DEFAULT_MAX_LINES) {
      return summary;
    }
    const truncated = truncateTail(summary.output, {
      maxLines: DEFAULT_MAX_LINES,
      maxBytes: Number.MAX_SAFE_INTEGER,
    });
    let artifactId = summary.artifactId;
    if (artifactId === undefined && this.#options.artifactPath !== undefined) {
      await writeFile(
        this.#options.artifactPath,
        this.#aggregateTail.text(),
        "utf8"
      );
      artifactId = this.#options.artifactPath;
    }
    return {
      ...summary,
      output: truncated.content,
      truncated: true,
      outputLines: truncated.outputLines,
      outputBytes: truncated.outputBytes,
      ...(artifactId === undefined ? {} : { artifactId }),
    };
  }
}

export function webpExclusionForModel(model: WebpModel): true | undefined {
  if (model === undefined) {
    return;
  }
  return model.provider === "ollama" ||
    model.provider === "ollama-cloud" ||
    model.provider === "llama.cpp" ||
    model.provider === "lm-studio" ||
    model.provider === "local-server" ||
    model.api === "ollama-chat"
    ? true
    : undefined;
}

export const resizeEvalImage: EvalImageResizer = async (image, model) => {
  const excludeWebP = webpExclusionForModel(model);
  const forceWebpConversion =
    excludeWebP === true && image.mimeType === "image/webp";
  const resized = await resizeImage(
    Buffer.from(image.data, "base64"),
    image.mimeType,
    forceWebpConversion
      ? { maxBytes: Buffer.byteLength(image.data, "utf8") }
      : undefined
  );
  let output: EvalImageContent =
    resized === null
      ? image
      : { type: "image", data: resized.data, mimeType: resized.mimeType };
  if (excludeWebP === true && output.mimeType === "image/webp") {
    const converted = await convertToPng(output.data, output.mimeType);
    if (converted === null) {
      throw new TypeError(
        `Unable to convert ${output.mimeType} display output for the active model`
      );
    }
    output = {
      type: "image",
      data: converted.data,
      mimeType: converted.mimeType,
    };
  }
  const dimensionNote =
    resized === null ? undefined : formatDimensionNote(resized);
  return {
    image: output,
    ...(dimensionNote === undefined ? {} : { dimensionNote }),
  };
};

function formatDisplayJson(value: unknown): string {
  let text: string;
  try {
    text = JSON.stringify(value, null, 2) ?? String(value);
  } catch (error) {
    if (!(error instanceof TypeError)) {
      throw error;
    }
    text = String(value);
  }
  if (text.length <= MAX_DISPLAY_TEXT_BYTES) {
    return text;
  }
  return `${text.slice(0, MAX_DISPLAY_TEXT_BYTES)}\n[…${text.length - MAX_DISPLAY_TEXT_BYTES}ch elided…]`;
}

interface FailureBudgetVisible {
  readonly elidedBytes: number;
  readonly elidedLines: number;
  readonly headRange?: { readonly start: number; readonly end: number };
  readonly outputBytes: number;
  readonly outputLines: number;
  readonly tailRange?: { readonly start: number; readonly end: number };
  readonly text: string;
}

function countLinesOf(text: string): number {
  return text.length === 0 ? 0 : text.split("\n").length;
}

/**
 * Bounds failed-cell output to the failure budget: head + middle elision
 * marker + tail, mirroring how OutputSink marks omitted middles.
 */
function failureBudgetOutput(
  summary: OutputSummary,
  headBytes: number
): FailureBudgetVisible {
  const budget = FAILURE_OUTPUT_MAX_BYTES - FAILURE_MARKER_RESERVE_BYTES;
  const head = truncateHeadBytes(
    summary.output,
    Math.min(headBytes, budget)
  );
  const tail = truncateTail(summary.output.slice(head.text.length), {
    maxLines: DEFAULT_MAX_LINES,
    maxBytes: Math.max(0, budget - head.bytes),
  });
  const headLines = countLinesOf(head.text);
  const elidedBytes = Math.max(
    0,
    summary.outputBytes - head.bytes - tail.outputBytes
  );
  const elidedLines = Math.max(
    0,
    summary.totalLines - headLines - tail.outputLines
  );
  const headSeparator = head.text.endsWith("\n") ? "" : "\n";
  const tailSeparator =
    tail.content.length === 0 || tail.content.startsWith("\n") ? "" : "\n";
  const text = `${head.text}${headSeparator}${formatMiddleElisionMarker(elidedLines, elidedBytes)}${tailSeparator}${tail.content}`;
  return {
    text,
    elidedBytes,
    elidedLines,
    outputBytes: Buffer.byteLength(text, "utf8"),
    outputLines: countLinesOf(text),
    ...(headLines > 0 ? { headRange: { start: 1, end: headLines } } : {}),
    ...(tail.outputLines > 0
      ? {
          tailRange: {
            start: Math.max(1, summary.totalLines - tail.outputLines + 1),
            end: summary.totalLines,
          },
        }
      : {}),
  };
}

function truncationMetaFromSummary(
  summary: OutputSummary
): TruncationMeta | undefined {
  if (!summary.truncated) {
    return;
  }
  const artifact =
    summary.artifactId === undefined ? {} : { artifactId: summary.artifactId };
  if (summary.elidedBytes !== undefined && summary.elidedBytes > 0) {
    const elidedLines =
      summary.elidedLines ??
      Math.max(0, summary.totalLines - summary.outputLines);
    const keptLines = Math.max(0, summary.outputLines - 1);
    const headLines = Math.ceil(keptLines / 2);
    const tailLines = keptLines - headLines;
    return {
      direction: "middle",
      truncatedBy: "middle",
      totalLines: summary.totalLines,
      totalBytes: summary.totalBytes,
      outputLines: summary.outputLines,
      outputBytes: summary.outputBytes,
      ...(headLines > 0 ? { headRange: { start: 1, end: headLines } } : {}),
      ...(tailLines > 0
        ? {
            tailRange: {
              start: summary.totalLines - tailLines + 1,
              end: summary.totalLines,
            },
          }
        : {}),
      elidedBytes: summary.elidedBytes,
      elidedLines,
      ...artifact,
    };
  }
  return {
    direction: "tail",
    truncatedBy: summary.outputBytes < summary.totalBytes ? "bytes" : "lines",
    totalLines: summary.totalLines,
    totalBytes: summary.totalBytes,
    outputLines: summary.outputLines,
    outputBytes: summary.outputBytes,
    shownRange: {
      start: Math.max(1, summary.totalLines - summary.outputLines + 1),
      end: summary.totalLines,
    },
    ...artifact,
  };
}

export function marshalToolResult(result: AgentToolResult<unknown>) {
  const texts = result.content
    .filter((part) => part.type === "text")
    .map((part) => part.text);
  const images = result.content
    .filter((part) => part.type === "image")
    .map((part) => ({ mimeType: part.mimeType, dataBase64: part.data }));
  const details =
    typeof result.details === "object" &&
    result.details !== null &&
    !Array.isArray(result.details) &&
    Object.keys(result.details).length === 0
      ? undefined
      : result.details;
  const hasError = toolResultIsError(result);
  const text = texts.join("\n");
  return images.length === 0 && details === undefined && !hasError
    ? { text }
    : { text, details, images, hasError };
}

export function toolResultIsError(result: AgentToolResult<unknown>): boolean {
  const details = result.details;
  return (
    typeof details === "object" &&
    details !== null &&
    "isError" in details &&
    details.isError === true
  );
}