import type { AgentToolResult } from "@earendil-works/pi-coding-agent";
import { detachedNotificationSpillPath } from "./detached-cell-notification.ts";
import {
  currentDetachedResult,
  detachedErrorResult,
  snapshotDetachedCell,
} from "./detached-cell-snapshot.ts";
import {
  activeDetachedCellReuseError,
  allowsDetachedCellTransition,
  detachedCellIsActive,
} from "./detached-cell-state.ts";
import { DetachedNotificationQueue } from "./detached-notification-queue.ts";
import type {
  EvalCellStatus,
  EvalKernel,
  EvalLanguage,
  EvalToolDetails,
  EvalToolInput,
} from "./types.ts";

// Derived from the canonical lifecycle vocabulary (EvalCellStatus): a cell only
// reaches the detached manager once it is running, so "pending" never applies.
export type EvalDetachedCellState = Exclude<EvalCellStatus, "pending">;

/** Terminal cell states: no further output can arrive, wait calls return at once. */
export function isTerminalDetachedState(state: EvalDetachedCellState): boolean {
  return state === "complete" || state === "error" || state === "cancelled";
}

/** Poll interval for the wait tool's output loop. */
const WAIT_POLL_INTERVAL_MS = 250;

type LiveResultProvider = () => AgentToolResult<EvalToolDetails>;

interface ManagedCell {
  canDetach: boolean;
  readonly cellId: string;
  readonly input: EvalToolInput;
  kernel: EvalKernel | undefined;
  /** Wait-tool cursor: UTF-16 offset into outputTail already returned by wait calls. */
  lastWaitSeenChars: number;
  liveResult: LiveResultProvider | undefined;
  notificationQueued: boolean;
  readonly spillPath: string | undefined;
  readonly startedAtMs: number;
  state: EvalDetachedCellState;
  stateRetained: boolean | undefined;
  readonly terminal: PromiseWithResolvers<EvalDetachedCellSnapshot>;
  terminalResult: AgentToolResult<EvalToolDetails> | undefined;
  wasDetached: boolean;
}

export interface EvalDetachedCellSnapshot {
  readonly cellId: string;
  readonly language: EvalLanguage;
  readonly outputTail: string;
  readonly result: AgentToolResult<EvalToolDetails>;
  readonly state: EvalDetachedCellState;
  readonly stateRetained: boolean | undefined;
}

export interface EvalDetachedCellNotification {
  readonly cellId: string;
  readonly content: string;
}

export interface EvalDetachedCellNotifier {
  notify: (cells: readonly EvalDetachedCellNotification[]) => void;
}

export interface EvalDetachedCellStatusEntry {
  readonly cellId: string;
  readonly language: EvalLanguage;
  readonly startedAtMs: number;
  readonly title?: string;
}

export interface EvalDetachedCellManagerOptions {
  readonly artifactsDir?: string;
  readonly notifier?: EvalDetachedCellNotifier;
  readonly now?: () => number;
  readonly onStatusChange?: (
    entries: readonly EvalDetachedCellStatusEntry[]
  ) => void;
}

export interface WaitForOutputOptions {
  readonly maxChars: number;
  readonly yieldTimeMs: number;
}

export interface WaitForOutputResult {
  readonly newText: string;
  readonly snapshot: EvalDetachedCellSnapshot;
}

export class EvalDetachedCellManager {
  readonly #artifactsDir: string | undefined;
  readonly #onStatusChange:
    | ((entries: readonly EvalDetachedCellStatusEntry[]) => void)
    | undefined;
  readonly #cells = new Map<string, ManagedCell>();
  readonly #detachedByLanguage = new Map<EvalLanguage, ManagedCell>();
  readonly #notificationQueue: DetachedNotificationQueue;
  readonly #now: () => number;

  constructor(options: EvalDetachedCellManagerOptions = {}) {
    this.#artifactsDir = options.artifactsDir;
    this.#onStatusChange = options.onStatusChange;
    this.#notificationQueue = new DetachedNotificationQueue(
      options.notifier,
      options.artifactsDir
    );
    this.#now = options.now ?? Date.now;
  }

  create(cellId: string, input: EvalToolInput): ManagedCell {
    const existing = this.#cells.get(cellId);
    if (existing !== undefined) {
      if (detachedCellIsActive(existing.state)) {
        throw activeDetachedCellReuseError(existing);
      }
      this.#cells.delete(cellId);
    }
    const cell: ManagedCell = {
      cellId,
      input,
      spillPath: detachedNotificationSpillPath(this.#artifactsDir, cellId),
      startedAtMs: this.#now(),
      state: "running",
      canDetach: false,
      wasDetached: false,
      kernel: undefined,
      stateRetained: undefined,
      liveResult: undefined,
      terminalResult: undefined,
      notificationQueued: false,
      lastWaitSeenChars: 0,
      terminal: Promise.withResolvers<EvalDetachedCellSnapshot>(),
    };
    this.#cells.set(cellId, cell);
    return cell;
  }

  markRunning(
    cell: ManagedCell,
    kernel: EvalKernel,
    liveResult: LiveResultProvider
  ): void {
    if (cell.state !== "running") {
      return;
    }
    cell.kernel = kernel;
    cell.liveResult = liveResult;
    cell.canDetach = true;
  }

  detach(cell: ManagedCell): boolean {
    if (
      !(cell.canDetach && allowsDetachedCellTransition(cell.state, "detached"))
    ) {
      return false;
    }
    cell.state = "detached";
    cell.wasDetached = true;
    this.#detachedByLanguage.set(cell.input.language, cell);
    this.#emitStatus();
    return true;
  }

  complete(
    cell: ManagedCell,
    result: AgentToolResult<EvalToolDetails>
  ): boolean {
    return this.#settle(
      cell,
      result.details.isError === true ? "error" : "complete",
      result
    );
  }

  fail(cell: ManagedCell, error: Error): boolean {
    return this.#settle(cell, "error", detachedErrorResult(cell, error));
  }

  async stop(
    cellId: string,
    reason = "Stopped detached eval cell"
  ): Promise<EvalDetachedCellSnapshot> {
    const cell = this.#get(cellId);
    if (cell.state === "detached") {
      // Interrupt first so the notification snapshot enqueued by #settle
      // sees the real stateRetained outcome; settle unconditionally so a
      // failing interrupt never prevents the terminal state.
      if (cell.kernel !== undefined) {
        try {
          const handle = await cell.kernel.interrupt(reason);
          cell.stateRetained = await handle.stateRetained;
        } catch {
          cell.stateRetained = undefined;
        }
      }
      this.#settle(cell, "cancelled", currentDetachedResult(cell));
    }
    return this.#snapshot(cell);
  }

  peek(cellId: string): EvalDetachedCellSnapshot {
    return this.#snapshot(this.#get(cellId));
  }

  busyFor(language: EvalLanguage): EvalDetachedCellSnapshot | undefined {
    const cell = this.#detachedByLanguage.get(language);
    return cell === undefined ? undefined : this.#snapshot(cell);
  }

  async waitForTerminal(cellId: string): Promise<EvalDetachedCellSnapshot> {
    return await this.#get(cellId).terminal.promise;
  }

  /**
   * Wait until the cell settles, produces new output, or the yield budget
   * elapses. Returns the current snapshot plus only the output not yet
   * returned by a previous wait call, capped at maxChars code points.
   */
  async waitForOutput(
    cellId: string,
    options: WaitForOutputOptions
  ): Promise<WaitForOutputResult> {
    const cell = this.#get(cellId);
    const deadline = this.#now() + options.yieldTimeMs;
    await this.#pollForOutput(cell, deadline);
    const snapshot = this.#snapshot(cell);
    return {
      snapshot,
      newText: this.#consumeWaitOutput(
        cell,
        snapshot.outputTail,
        options.maxChars
      ),
    };
  }

  async #pollForOutput(cell: ManagedCell, deadlineMs: number): Promise<void> {
    while (true) {
      const snapshot = this.#snapshot(cell);
      if (isTerminalDetachedState(snapshot.state)) {
        return;
      }
      if (snapshot.outputTail.length > cell.lastWaitSeenChars) {
        return;
      }
      if (this.#now() >= deadlineMs) {
        return;
      }
      // Sleep one poll interval; wake early when the cell settles so a
      // terminal snapshot is returned without waiting out the interval.
      await new Promise<void>((resolve) => {
        const timer = setTimeout(resolve, WAIT_POLL_INTERVAL_MS);
        void cell.terminal.promise.then(() => {
          clearTimeout(timer);
          resolve();
        });
      });
    }
  }

  #consumeWaitOutput(
    cell: ManagedCell,
    outputTail: string,
    maxChars: number
  ): string {
    // A shrunken tail means the kernel truncated its buffer; restart the cursor.
    if (outputTail.length < cell.lastWaitSeenChars) {
      cell.lastWaitSeenChars = 0;
    }
    let newText = outputTail.slice(cell.lastWaitSeenChars);
    if (Array.from(newText).length > maxChars) {
      newText = Array.from(newText).slice(0, maxChars).join("");
    }
    cell.lastWaitSeenChars += newText.length;
    return newText;
  }

  async dispose(): Promise<void> {
    const detached = [...this.#detachedByLanguage.values()];
    await Promise.allSettled(
      detached.map(
        async (cell) =>
          await this.stop(
            cell.cellId,
            "Session ended; detached eval cell cancelled"
          )
      )
    );
    await this.#notificationQueue.flush();
  }

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

  #settle(
    cell: ManagedCell,
    state: "complete" | "error" | "cancelled",
    result: AgentToolResult<EvalToolDetails>
  ): boolean {
    if (!allowsDetachedCellTransition(cell.state, state)) {
      return false;
    }
    cell.state = state;
    cell.terminalResult = result;
    cell.liveResult = undefined;
    cell.terminal.resolve(this.#snapshot(cell));
    if (cell.wasDetached) {
      if (this.#detachedByLanguage.get(cell.input.language) === cell) {
        this.#detachedByLanguage.delete(cell.input.language);
      }
      this.#emitStatus();
      if (!cell.notificationQueued) {
        cell.notificationQueued = true;
        this.#notificationQueue.enqueue({
          snapshot: () => this.#snapshot(cell),
          spillPath: cell.spillPath,
        });
      }
    }
    return true;
  }

  #emitStatus(): void {
    this.#onStatusChange?.(
      [...this.#detachedByLanguage.values()].map((cell) => ({
        cellId: cell.cellId,
        language: cell.input.language,
        startedAtMs: cell.startedAtMs,
        ...(cell.input.title === undefined ? {} : { title: cell.input.title }),
      }))
    );
  }

  #snapshot(cell: ManagedCell): EvalDetachedCellSnapshot {
    return snapshotDetachedCell(cell, this.#now());
  }

  #get(cellId: string): ManagedCell {
    const cell = this.#cells.get(cellId);
    if (cell === undefined) {
      throw new Error(`Unknown detached eval cell "${cellId}"`);
    }
    return cell;
  }
}