import { createWriteStream, mkdirSync, type WriteStream } from "node:fs";
import { dirname } from "node:path";
import {
  DEFAULT_MAX_BYTES,
  DEFAULT_MAX_LINES,
  truncateTail,
} from "../host-sdk.ts";
import { formatMiddleElisionMarker } from "./output-meta.ts";

export {
  artifactNotice,
  formatMiddleElisionMarker,
  resolveSessionArtifactsDir,
} from "./output-meta.ts";
export { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, truncateTail };

export interface OutputSummary {
  readonly artifactId?: string;
  readonly columnDroppedBytes?: number;
  readonly columnTruncatedLines?: number;
  readonly elidedBytes?: number;
  readonly elidedLines?: number;
  readonly output: string;
  readonly outputBytes: number;
  readonly outputLines: number;
  readonly totalBytes: number;
  readonly totalLines: number;
  readonly truncated: boolean;
}

export interface OutputSinkOptions {
  readonly artifactPath?: string;
  readonly chunkThrottleMs?: number;
  readonly headBytes?: number;
  readonly maxColumns?: number;
  readonly onChunk?: (chunk: string) => void;
  readonly spillThreshold?: number;
}

interface ByteSlice {
  readonly bytes: number;
  readonly text: string;
}

function countNewlines(text: string): number {
  let count = 0;
  let cursor = text.indexOf("\n");
  while (cursor !== -1) {
    count += 1;
    cursor = text.indexOf("\n", cursor + 1);
  }
  return count;
}

function lineCount(text: string): number {
  return text.length === 0 ? 0 : countNewlines(text) + 1;
}

export function truncateHeadBytes(text: string, maxBytes: number): ByteSlice {
  if (maxBytes <= 0) {
    return { text: "", bytes: 0 };
  }
  const buffer = Buffer.from(text, "utf8");
  if (buffer.length <= maxBytes) {
    return { text, bytes: buffer.length };
  }
  let end = maxBytes;
  while (end > 0 && (buffer[end] & 0xc0) === 0x80) {
    end -= 1;
  }
  const slice = buffer.subarray(0, end);
  return { text: slice.toString("utf8"), bytes: slice.length };
}

function truncateTailBytes(text: string, maxBytes: number): ByteSlice {
  if (maxBytes <= 0) {
    return { text: "", bytes: 0 };
  }
  const buffer = Buffer.from(text, "utf8");
  if (buffer.length <= maxBytes) {
    return { text, bytes: buffer.length };
  }
  let start = buffer.length - maxBytes;
  while (start < buffer.length && (buffer[start] & 0xc0) === 0x80) {
    start += 1;
  }
  const slice = buffer.subarray(start);
  return { text: slice.toString("utf8"), bytes: slice.length };
}

export class TailBuffer {
  readonly #maxBytes: number;
  #chunks: Buffer[] = [];
  #frontSkip = 0;
  #bytes = 0;

  constructor(maxBytes: number) {
    this.#maxBytes = Math.max(0, Math.floor(maxBytes));
  }

  append(text: string): void {
    if (text.length === 0) {
      return;
    }
    if (this.#maxBytes === 0) {
      this.#chunks = [];
      this.#frontSkip = 0;
      this.#bytes = 0;
      return;
    }
    const incomingBytes = Buffer.byteLength(text, "utf8");
    if (incomingBytes >= this.#maxBytes) {
      // A single chunk that alone exceeds the cap: keep only its tail.
      this.#chunks = [
        Buffer.from(truncateTailBytes(text, this.#maxBytes).text, "utf8"),
      ];
      this.#frontSkip = 0;
      this.#bytes = this.#chunks[0].length;
      return;
    }
    this.#chunks.push(Buffer.from(text, "utf8"));
    this.#bytes += incomingBytes;
    while (this.#bytes > this.#maxBytes && this.#chunks.length > 1) {
      const front = this.#chunks[0];
      if (front === undefined) {
        break;
      }
      const frontRemaining = front.length - this.#frontSkip;
      if (this.#bytes - frontRemaining >= this.#maxBytes) {
        // The whole front chunk falls outside the retained window: drop it.
        this.#chunks.shift();
        this.#frontSkip = 0;
        this.#bytes -= frontRemaining;
        continue;
      }
      // The cutoff falls inside the front chunk: trim it at the UTF-8
      // boundary, byte-identical to truncateTailBytes on the concatenation.
      let newSkip = this.#frontSkip + (this.#bytes - this.#maxBytes);
      while (newSkip < front.length && (front[newSkip] & 0xc0) === 0x80) {
        newSkip += 1;
      }
      if (newSkip >= front.length) {
        this.#chunks.shift();
        this.#frontSkip = 0;
        this.#bytes -= frontRemaining;
        continue;
      }
      this.#bytes -= newSkip - this.#frontSkip;
      this.#frontSkip = newSkip;
      break;
    }
    if (this.#bytes > this.#maxBytes) {
      // Loop stopped at a single over-cap chunk: trim it in place.
      const front = this.#chunks[0];
      if (front !== undefined) {
        let newSkip = this.#frontSkip + (this.#bytes - this.#maxBytes);
        while (newSkip < front.length && (front[newSkip] & 0xc0) === 0x80) {
          newSkip += 1;
        }
        this.#frontSkip = Math.min(newSkip, front.length);
        this.#bytes = front.length - this.#frontSkip;
      }
    }
  }

  text(): string {
    if (this.#chunks.length === 0) {
      return "";
    }
    if (this.#chunks.length === 1) {
      return this.#chunks[0].subarray(this.#frontSkip).toString("utf8");
    }
    const parts = [this.#chunks[0].subarray(this.#frontSkip).toString("utf8")];
    for (let i = 1; i < this.#chunks.length; i += 1) {
      const chunk = this.#chunks[i];
      if (chunk === undefined) {
        break;
      }
      parts.push(chunk.toString("utf8"));
    }
    return parts.join("");
  }

  bytes(): number {
    return this.#bytes;
  }
}

export class OutputSink {
  readonly #artifactPath: string | undefined;
  readonly #spillThreshold: number;
  readonly #headLimit: number;
  readonly #maxColumns: number;
  readonly #onChunk: ((chunk: string) => void) | undefined;
  readonly #chunkThrottleMs: number;
  readonly #tail: TailBuffer;
  #head = "";
  #headBytes = 0;
  #totalNewlines = 0;
  #totalBytes = 0;
  #sawData = false;
  #truncated = false;
  #currentLineBytes = 0;
  #columnCapped = false;
  #columnDroppedBytes = 0;
  #columnTruncatedLines = 0;
  #lastChunkTime = 0;
  #pendingChunk = "";
  #beforeSpill = "";
  #file: WriteStream | undefined;
  #fileError: Error | undefined;
  #dumpPromise: Promise<OutputSummary> | undefined;

  constructor(options: OutputSinkOptions = {}) {
    this.#artifactPath = options.artifactPath;
    this.#spillThreshold = Math.max(
      0,
      Math.floor(options.spillThreshold ?? DEFAULT_MAX_BYTES)
    );
    this.#headLimit = Math.max(0, Math.floor(options.headBytes ?? 0));
    this.#maxColumns = Math.max(0, Math.floor(options.maxColumns ?? 0));
    this.#onChunk = options.onChunk;
    this.#chunkThrottleMs = Math.max(
      0,
      Math.floor(options.chunkThrottleMs ?? 0)
    );
    this.#tail = new TailBuffer(this.#spillThreshold);
  }

  push(chunk: string): void {
    if (chunk.length === 0) {
      return;
    }
    this.#emitPreview(chunk);
    const rawBytes = Buffer.byteLength(chunk, "utf8");
    this.#totalBytes += rawBytes;
    this.#totalNewlines += countNewlines(chunk);
    this.#sawData = true;
    this.#mirrorRaw(chunk);
    this.#retain(this.#maxColumns > 0 ? this.#clampColumns(chunk) : chunk);
  }

  dump(notice?: string): Promise<OutputSummary> {
    this.#dumpPromise ??= this.#finishDump(notice);
    return this.#dumpPromise;
  }

  #emitPreview(chunk: string): void {
    if (this.#onChunk === undefined) {
      return;
    }
    const now = Date.now();
    if (now - this.#lastChunkTime >= this.#chunkThrottleMs) {
      this.#lastChunkTime = now;
      this.#onChunk(this.#pendingChunk + chunk);
      this.#pendingChunk = "";
      return;
    }
    this.#pendingChunk += chunk;
  }

  #mirrorRaw(chunk: string): void {
    if (this.#artifactPath === undefined) {
      return;
    }
    if (this.#file !== undefined) {
      this.#file.write(chunk);
      return;
    }
    if (this.#totalBytes <= this.#spillThreshold) {
      this.#beforeSpill += chunk;
      return;
    }
    mkdirSync(dirname(this.#artifactPath), { recursive: true });
    const stream = createWriteStream(this.#artifactPath, { encoding: "utf8" });
    stream.on("error", (error) => {
      this.#fileError = error;
    });
    this.#file = stream;
    if (this.#beforeSpill.length > 0) {
      stream.write(this.#beforeSpill);
    }
    this.#beforeSpill = "";
    stream.write(chunk);
  }

  #retain(text: string): void {
    let tailText = text;
    if (this.#headBytes < this.#headLimit) {
      const head = truncateHeadBytes(text, this.#headLimit - this.#headBytes);
      this.#head += head.text;
      this.#headBytes += head.bytes;
      tailText = text.slice(head.text.length);
    }
    this.#tail.append(tailText);
    const effectiveBytes = this.#totalBytes - this.#columnDroppedBytes;
    if (effectiveBytes > this.#headBytes + this.#tail.bytes()) {
      this.#truncated = true;
    }
  }

  #clampColumns(chunk: string): string {
    const output: string[] = [];
    let cursor = 0;
    while (cursor < chunk.length) {
      const newline = chunk.indexOf("\n", cursor);
      const end = newline === -1 ? chunk.length : newline;
      const segment = chunk.slice(cursor, end);
      if (segment.length > 0) {
        const segmentBytes = Buffer.byteLength(segment, "utf8");
        if (this.#columnCapped) {
          this.#columnDroppedBytes += segmentBytes;
        } else {
          const remaining = Math.max(
            0,
            this.#maxColumns - this.#currentLineBytes
          );
          const kept = truncateHeadBytes(segment, remaining);
          output.push(kept.text);
          this.#currentLineBytes += kept.bytes;
          if (kept.bytes < segmentBytes) {
            output.push("…");
            this.#columnDroppedBytes += segmentBytes - kept.bytes;
            this.#columnTruncatedLines += 1;
            this.#columnCapped = true;
            this.#truncated = true;
          }
        }
      }
      if (newline === -1) {
        break;
      }
      output.push("\n");
      this.#currentLineBytes = 0;
      this.#columnCapped = false;
      cursor = newline + 1;
    }
    return output.join("");
  }

  async #finishDump(notice: string | undefined): Promise<OutputSummary> {
    if (this.#onChunk !== undefined && this.#pendingChunk.length > 0) {
      this.#onChunk(this.#pendingChunk);
      this.#pendingChunk = "";
    }
    await this.#closeFile();
    let tail = this.#tail.text();
    if (lineCount(tail) > DEFAULT_MAX_LINES) {
      tail = truncateTail(tail, {
        maxLines: DEFAULT_MAX_LINES,
        maxBytes: Number.MAX_SAFE_INTEGER,
      }).content;
      this.#truncated = true;
    }
    const totalLines = this.#sawData ? this.#totalNewlines + 1 : 0;
    const tailBytes = Buffer.byteLength(tail, "utf8");
    const effectiveBytes = Math.max(
      0,
      this.#totalBytes - this.#columnDroppedBytes
    );
    let body = this.#head + tail;
    let elidedBytes: number | undefined;
    let elidedLines: number | undefined;
    if (this.#headBytes > 0 && effectiveBytes > this.#headBytes + tailBytes) {
      elidedBytes = effectiveBytes - this.#headBytes - tailBytes;
      elidedLines = Math.max(
        0,
        totalLines - lineCount(this.#head) - lineCount(tail)
      );
      const headSeparator = this.#head.endsWith("\n") ? "" : "\n";
      const tailSeparator =
        tail.length === 0 || tail.startsWith("\n") ? "" : "\n";
      body = `${this.#head}${headSeparator}${formatMiddleElisionMarker(elidedLines, elidedBytes)}${tailSeparator}${tail}`;
      this.#truncated = true;
    }
    return {
      output: notice === undefined ? body : `[${notice}]\n${body}`,
      truncated: this.#truncated,
      totalLines,
      totalBytes: this.#totalBytes,
      outputLines: lineCount(body),
      outputBytes: Buffer.byteLength(body, "utf8"),
      elidedBytes,
      elidedLines,
      columnDroppedBytes:
        this.#columnDroppedBytes > 0 ? this.#columnDroppedBytes : undefined,
      columnTruncatedLines:
        this.#columnTruncatedLines > 0 ? this.#columnTruncatedLines : undefined,
      artifactId: this.#file === undefined ? undefined : this.#artifactPath,
    };
  }

  async #closeFile(): Promise<void> {
    const stream = this.#file;
    if (stream === undefined) {
      return;
    }
    if (this.#fileError !== undefined) {
      throw this.#fileError;
    }
    await new Promise<void>((resolve, reject) => {
      const onError = (error: Error) => {
        stream.off("finish", onFinish);
        reject(error);
      };
      const onFinish = () => {
        stream.off("error", onError);
        resolve();
      };
      stream.once("error", onError);
      stream.once("finish", onFinish);
      stream.end();
    });
  }
}