import { randomUUID } from "node:crypto";
import type { Api, Model } from "@earendil-works/pi-ai/compat";
import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { getConfig, type EffectiveBtwSettings } from "./config.ts";
import { type ScopeToken, type ThreadStore } from "./threads.ts";
import type { SummarySnapshot } from "./context-policy.ts";
import { runSummary, type SummaryRunResult } from "./summary.ts";

/** The optional model argument preserves the original four-argument runner seam. */
export type SummaryRunner = (ctx: ExtensionContext, input: string, maxTokens: number, signal: AbortSignal, model?: Model<Api>) => Promise<SummaryRunResult>;
export type SummaryCoordinatorOptions = {
  runner?: SummaryRunner;
  config?: () => EffectiveBtwSettings;
  id?: () => string;
  now?: () => number;
  shutdownGraceMs?: number;
};

type Job = { controller: AbortController; generation: number; token: ScopeToken; runId: string; key: string; promise: Promise<void> };
type KeyState = { dirty: boolean; running?: Job; ctx?: ExtensionContext; model?: Model<Api>; token?: ScopeToken; threadId?: string; generation?: number };
const keyFor = (token: ScopeToken, threadId: string) => `${token.scopeId}\u0000${threadId}`;

/**
 * Process-local, best-effort rolling-summary maintenance. Store CAS/metering is
 * still authoritative across processes. A process-local active registry prevents
 * reset from accidentally overlapping an abandoned run with a reused key.
 */
export class SummaryCoordinator {
  private readonly states = new Map<string, KeyState>();
  private readonly activeJobs = new Map<string, Job>();
  private generation = 0;
  private interactiveDepth = 0;
  private readonly leases = new Set<symbol>();
  private readonly pendingFlushes = new Set<Promise<void>>();
  private stopping = false;
  private finishGateOpen = true;
  private readonly runner: SummaryRunner;
  private readonly config: () => EffectiveBtwSettings;
  private readonly id: () => string;
  private readonly now: () => number;
  private readonly grace: number;

  constructor(private readonly threads: ThreadStore, options: SummaryCoordinatorOptions = {}) {
    this.runner = options.runner ?? ((ctx, input, maxTokens, signal, model) => runSummary(ctx, input, maxTokens, signal, { model }));
    this.config = options.config ?? getConfig;
    this.id = options.id ?? randomUUID;
    this.now = options.now ?? Date.now;
    this.grace = Math.max(0, options.shutdownGraceMs ?? 100);
  }

  /** Starting any foreground operation globally preempts maintenance. */
  beginInteractive(_ctx: ExtensionContext, _token: ScopeToken, _threadId: string): () => void {
    const lease = Symbol("btw-interactive");
    this.leases.add(lease);
    this.interactiveDepth++;
    for (const state of this.states.values()) {
      if (state.running) {
        state.dirty = true;
        state.running.controller.abort();
      }
    }
    let released = false;
    return () => {
      if (released) return;
      released = true;
      if (!this.leases.delete(lease)) return; // resetScope invalidated this stale release.
      this.interactiveDepth = Math.max(0, this.interactiveDepth - 1);
      if (!this.interactiveDepth && !this.stopping) this.pumpAll();
    };
  }

  /** Persistence is the barrier: a crash before it completes cannot schedule a summary. */
  afterSuccessfulAttempt(ctx: ExtensionContext, token: ScopeToken, threadId: string): void {
    const generation = this.generation;
    // Capture the actual object before the async persistence barrier.
    const model = ctx.model as Model<Api> | undefined;
    const persisted = Promise.resolve()
      .then(() => this.threads.flush())
      .then(() => {
        if (this.stopping || generation !== this.generation || !this.threads.isCurrent(token)) return;
        const key = keyFor(token, threadId);
        const state = this.states.get(key) ?? { dirty: false };
        state.dirty = true;
        state.ctx = ctx; state.model = model; state.token = token; state.threadId = threadId; state.generation = generation;
        this.states.set(key, state);
        this.pump(ctx, model, token, threadId, generation);
      })
      .catch(() => { /* persistence already reports its own one-time UI failure */ });
    this.pendingFlushes.add(persisted);
    void persisted.finally(() => this.pendingFlushes.delete(persisted));
  }

  /** Invalidates detached persistence callbacks and prevents old jobs from rescheduling. */
  resetScope(): void {
    this.generation++;
    this.interactiveDepth = 0;
    this.leases.clear();
    for (const job of this.activeJobs.values()) job.controller.abort();
    this.states.clear();
  }

  private pumpAll(): void {
    for (const state of this.states.values()) {
      if (state.dirty && state.ctx && state.token && state.threadId && state.generation !== undefined) {
        this.pump(state.ctx, state.model, state.token, state.threadId, state.generation);
      }
    }
  }

  private pump(ctx: ExtensionContext, model: Model<Api> | undefined, token: ScopeToken, threadId: string, generation: number): void {
    if (this.stopping || generation !== this.generation || this.interactiveDepth || !this.threads.isCurrent(token)) return;
    const settings = this.config();
    if (!settings.summaryEnabled || !model) return;
    const key = keyFor(token, threadId);
    const state = this.states.get(key);
    if (!state || !state.dirty || state.running || this.activeJobs.has(key)) return;
    state.ctx = ctx; state.model = model; state.token = token; state.threadId = threadId; state.generation = generation;
    const inputMaxTokens = Math.max(0, (model.contextWindow ?? 0) - settings.summaryMaxTokens - 1024);
    const snapshot = this.threads.getSummarySnapshotIfCurrent(token, threadId, {
      triggerTokens: settings.summaryTriggerTokens,
      retainTokens: settings.summaryRetainTokens,
      inputMaxTokens,
    });
    if (!snapshot) return;
    state.dirty = false; // consume exactly one persisted successful-attempt trigger.
    const controller = new AbortController();
    const job: Job = { controller, generation, token, runId: this.id(), key, promise: Promise.resolve() };
    state.running = job;
    this.activeJobs.set(key, job);
    job.promise = this.execute(ctx, model, token, threadId, snapshot, model.id, settings.summaryMaxTokens, job);
  }

  private async execute(ctx: ExtensionContext, model: Model<Api>, token: ScopeToken, threadId: string, snapshot: SummarySnapshot, modelId: string, maxTokens: number, job: Job): Promise<void> {
    let result: SummaryRunResult;
    try {
      result = await this.runner(ctx, snapshot.input, maxTokens, job.controller.signal, model);
    } catch (error) {
      // Custom test runners are not allowed to break a detached coordinator.
      result = { kind: job.controller.signal.aborted ? "aborted" : "failed", issued: true, ...(job.controller.signal.aborted ? {} : { error: error instanceof Error ? error.message : String(error) }) } as SummaryRunResult;
    }
    // A runner may ignore abort and return success. Such a result is real usage,
    // but never a commit after reset, preemption, or shutdown.
    if (job.controller.signal.aborted && result.issued) result = { kind: "aborted", issued: true, ...(result.usage ? { usage: result.usage } : {}) };
    if (result.issued && this.finishGateOpen) {
      const outcome = result.kind === "success"
        ? { summary: { text: result.text.trim(), throughEntryId: snapshot.throughEntryId, source: snapshot.source, sourceHash: snapshot.sourceHash, createdAt: safeStamp(this.now), model: modelId } }
        : result.kind === "aborted" ? "aborted" as const : "failed" as const;
      try {
        this.threads.finishSummary(token.scopeId, threadId, job.runId, outcome, result.usage);
        await this.threads.flush();
      } catch { /* routine maintenance never blocks or notifies foreground work */ }
    }

    if (this.activeJobs.get(job.key) === job) this.activeJobs.delete(job.key);
    const state = this.states.get(job.key);
    if (!state) return;
    if (state.running !== job) {
      // resetScope may have replaced state for this same process-local key.
      // The old job was the registry blocker; its settlement now lets current
      // dirty work proceed, without letting the old result affect that state.
      if (!this.stopping && !this.interactiveDepth && state.dirty && state.ctx && state.token && state.threadId && state.generation === this.generation && this.threads.isCurrent(state.token)) {
        this.pump(state.ctx, state.model, state.token, state.threadId, state.generation);
      }
      return;
    }
    state.running = undefined;
    if (job.generation !== this.generation || this.stopping || !this.threads.isCurrent(token)) return;
    if (result.kind === "aborted" && job.controller.signal.aborted) state.dirty = true;
    // Normal provider failures do not retry; a later persisted answer must dirty it.
    if (state.dirty && !this.interactiveDepth) this.pump(state.ctx!, state.model, token, threadId, job.generation);
  }

  /** A settlement seam that includes jobs detached by reset as well as current state. */
  async drain(): Promise<void> {
    while (true) {
      const jobs = [...this.activeJobs.values()].map((job) => job.promise);
      const pending = [...this.pendingFlushes];
      if (!jobs.length && !pending.length) return;
      await Promise.allSettled([...jobs, ...pending]);
    }
  }

  async shutdown(): Promise<void> {
    this.stopping = true;
    for (const job of this.activeJobs.values()) job.controller.abort();
    const settled = this.drain();
    let timer: ReturnType<typeof setTimeout> | undefined;
    const timeout = new Promise<void>((resolve) => {
      timer = setTimeout(resolve, this.grace);
      (timer as unknown as { unref?: () => void }).unref?.();
    });
    await Promise.race([settled, timeout]);
    if (timer) clearTimeout(timer);
    // No late/non-cooperative settlement may mutate or flush after shutdown returns.
    this.finishGateOpen = false;
    await this.threads.flush().catch(() => {});
  }
}

function safeStamp(now: () => number): string {
  try { return new Date(now()).toISOString(); } catch { return ""; }
}