import type {
  AgentToolResult,
  AgentToolUpdateCallback,
  ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import type { KernelToHostMessage } from "../bridge/protocol.ts";
import {
  type EvalImageResizer,
  EvalOutputCollector,
  type EvalOutputResult,
} from "./image.ts";
import type {
  EvalCellStatus,
  EvalStatusEvent,
  EvalToolDetails,
  EvalToolInput,
} from "./types.ts";

type KernelResult = Extract<KernelToHostMessage, { type: "result" }>;
type DisplayMessage = Extract<KernelToHostMessage, { type: "display" }>;
type ToolCall = EvalToolDetails["toolCalls"] extends readonly (infer Item)[]
  ? Item
  : never;

export interface CellState {
  active: boolean;
  durationMs: number;
  readonly input: EvalToolInput;
  readonly onUpdate: AgentToolUpdateCallback<EvalToolDetails> | undefined;
  output: string;
  readonly pendingBridgeCalls: Promise<void>[];
  phase: string | undefined;
  readonly signal: AbortSignal;
  // Subset of EvalCellStatus: a foreground cell never detaches or gets cancelled.
  status: Exclude<EvalCellStatus, "detached" | "cancelled">;
  readonly statusEvents: EvalStatusEvent[];
  readonly toolCalls: ToolCall[];
}

export interface CellResultBuilderOptions {
  readonly artifactPath?: string;
  /** Optional sink-level preview batching (ms); unset keeps per-chunk onChunk delivery. */
  readonly chunkThrottleMs?: number;
  readonly headBytes: number;
  readonly imageResizer?: EvalImageResizer;
  readonly maxColumns: number;
  readonly model: ExtensionContext["model"];
  readonly state: CellState;
}

// Plan 020: a live cell emits one update per output chunk, and bursty kernels
// push many chunks per tick — each update re-renders the whole row. Chunks that
// arrive within this window are folded into one trailing update (see
// #coalesceOutputUpdate); pending chunks are flushed synchronously at
// finalize/liveResult so the final state is never delayed.
const OUTPUT_UPDATE_COALESCE_MS = 50;

export class CellResultBuilder {
  readonly #output: EvalOutputCollector;
  readonly #state: CellState;
  #updatePending = false;
  #lastUpdateEmittedAt = 0;

  constructor(options: CellResultBuilderOptions) {
    this.#state = options.state;
    this.#output = new EvalOutputCollector({
      headBytes: options.headBytes,
      maxColumns: options.maxColumns,
      model: options.model,
      chunkThrottleMs: options.chunkThrottleMs,
      ...(options.artifactPath === undefined
        ? {}
        : { artifactPath: options.artifactPath }),
      ...(options.imageResizer === undefined
        ? {}
        : { imageResizer: options.imageResizer }),
      onChunk: (_aggregate, cell) => {
        options.state.output = cell;
        this.#coalesceOutputUpdate();
      },
    });
    options.state.status = "running";
    this.emitUpdate(false);
  }

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

  display(message: DisplayMessage): void {
    this.#output.display(message);
  }

  setPhase(title: string): void {
    this.#state.phase = title;
    this.emitUpdate(false);
  }

  async finalize(
    result: KernelResult
  ): Promise<AgentToolResult<EvalToolDetails>> {
    this.#state.durationMs = result.durationMs;
    // Flush any coalesced preview chunks while the cell is still running so
    // the trailing live update carries the exact burst.
    this.#flushPendingUpdate(false);
    if (result.ok) {
      if (result.valueRepr) {
        this.#output.push(`${result.valueRepr}\n`);
      }
      this.#state.status = "complete";
    } else {
      this.#output.push(`${result.error.message}\n`);
      this.#state.status = "error";
    }
    return await this.#finish(!result.ok);
  }

  async finalizeCancellation(
    error: Error
  ): Promise<AgentToolResult<EvalToolDetails>> {
    this.#flushPendingUpdate(false);
    this.#output.push(`${error.message}\n`);
    this.#state.status = "error";
    return await this.#finish(true);
  }

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

  liveResult(): AgentToolResult<EvalToolDetails> {
    this.#flushPendingUpdate(this.#state.status === "error");
    return {
      content: [{ type: "text", text: this.#liveUpdateText() }],
      details: this.#details(undefined, this.#state.status === "error"),
    };
  }

  emitUpdate(isError: boolean): void {
    if (!this.#state.active) {
      return;
    }
    this.#state.onUpdate?.({
      content: [{ type: "text", text: this.#liveUpdateText() }],
      details: this.#details(undefined, isError),
    });
  }

  /**
   * Leading+trailing coalescer for chunk-driven updates: the first chunk of a
   * burst emits immediately (so the live tail grows monotonically), chunks that
   * arrive within the coalescing window are folded into the next emit, and
   * finalize/liveResult flush any pending chunk so the last update always
   * reflects the full state. Status/tool-call updates call emitUpdate directly
   * and stay synchronous.
   */
  #coalesceOutputUpdate(): void {
    const now = Date.now();
    if (
      this.#updatePending &&
      now - this.#lastUpdateEmittedAt < OUTPUT_UPDATE_COALESCE_MS
    ) {
      return;
    }
    this.#updatePending = true;
    this.#lastUpdateEmittedAt = now;
    this.emitUpdate(false);
  }

  #flushPendingUpdate(isError: boolean): void {
    if (!this.#updatePending) {
      return;
    }
    this.#updatePending = false;
    this.emitUpdate(isError);
  }

  async #finish(isError: boolean): Promise<AgentToolResult<EvalToolDetails>> {
    const output = await this.#output.finish(isError);
    this.#state.output = output.output;
    const details = this.#details(output, isError);
    this.#flushPendingUpdate(isError);
    this.emitUpdate(isError);
    const text =
      output.output ||
      (output.images.length > 0
        ? `(displayed ${output.images.length} image${output.images.length === 1 ? "" : "s"}; no text output)`
        : "(no output)");
    return { content: [{ type: "text", text }, ...output.images], details };
  }

  #details(
    output: EvalOutputResult | undefined,
    isError: boolean
  ): EvalToolDetails {
    const statusEvents =
      this.#state.statusEvents.length > 0
        ? [...this.#state.statusEvents]
        : undefined;
    return {
      language: this.#state.input.language,
      languages: [this.#state.input.language],
      ...(this.#state.input.title === undefined
        ? {}
        : { title: this.#state.input.title }),
      durationMs: this.#state.durationMs,
      toolCalls: [...this.#state.toolCalls],
      truncated: output?.truncated ?? false,
      ...(isError ? { isError: true } : {}),
      ...(this.#state.phase === undefined ? {} : { phase: this.#state.phase }),
      cells: [
        {
          index: 0,
          ...(this.#state.input.title === undefined
            ? {}
            : { title: this.#state.input.title }),
          code: this.#state.input.code,
          language: this.#state.input.language,
          output: this.#state.output,
          status: this.#state.status,
          durationMs: this.#state.durationMs,
          ...(statusEvents === undefined ? {} : { statusEvents }),
          ...(output?.hasMarkdown ? { hasMarkdown: true } : {}),
        },
      ],
      ...(statusEvents === undefined ? {} : { statusEvents }),
      ...(output === undefined || output.jsonOutputs.length === 0
        ? {}
        : { jsonOutputs: output.jsonOutputs }),
      ...(output?.notice === undefined ? {} : { notice: output.notice }),
      ...(output?.meta === undefined ? {} : { meta: output.meta }),
    };
  }

  #liveUpdateText(): string {
    const title =
      this.#state.input.title === undefined
        ? ""
        : ` ${this.#state.input.title}`;
    const aggregateOutput = this.#output.aggregateText();
    const hasTrailingNewline = aggregateOutput.endsWith("\n");
    const regionEnd = hasTrailingNewline
      ? aggregateOutput.length - 1
      : aggregateOutput.length;
    // O(last 8 lines): locate the 8th-last newline by reverse scan instead of
    // splitting the whole retained tail on every chunk.
    let start = 0;
    let searchFrom = regionEnd;
    for (let remaining = 8; remaining > 0; remaining -= 1) {
      const newlineIndex = aggregateOutput.lastIndexOf("\n", searchFrom - 1);
      if (newlineIndex === -1) {
        break;
      }
      if (remaining === 1) {
        start = newlineIndex + 1;
      }
      searchFrom = newlineIndex;
    }
    const output = `${aggregateOutput.slice(start, regionEnd)}${hasTrailingNewline ? "\n" : ""}`;
    return `1/1 cells ${this.#state.status}\n[1] ${this.#state.input.language}${title} ${this.#state.status}${output.length === 0 ? "" : `\n${output}`}`;
  }
}