import { existsSync, readFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import {
  DEFAULT_COMPACTION_SETTINGS,
  type CompactionEntry,
  type ExtensionAPI,
  type ExtensionContext,
  type SessionEntry,
} from "@earendil-works/pi-coding-agent";
import {
  CONTINUATION_MESSAGE_TYPE,
  REMINDER_MESSAGE_TYPE,
  ROLLOVER_INSTRUCTIONS_PREFIX,
  type EpochConfig,
  type EpochTrigger,
} from "./constants.ts";
import { scanForConflictingPackages } from "./detect.ts";
import {
  buildContinuationText,
  buildEpochCompaction,
  getEpochState,
  type EpochBoundaryDetails,
} from "./epoch.ts";
import { evaluatePressure, reminderText, type PressureBand } from "./pressure.ts";
import { replayNotes } from "./notes.ts";

export interface ArmResult {
  armed: boolean;
  /** User/model-facing explanation when armed === false. */
  reason?: string;
  epochIndex?: number;
}

export interface RolloverOutcome {
  performed: boolean;
  reason?: string;
  details?: EpochBoundaryDetails;
}

interface PendingRollover {
  generation: number;
  trigger: EpochTrigger;
  sessionFile: string | null;
  /** epochIndex observed when armed - used to detect a boundary that raced us. */
  epochIndexAtArm: number;
}

interface RolloverExpectation {
  generation: number;
  trigger: EpochTrigger;
  sessionFile: string | null;
}

const ROLLOVER_TIMEOUT_MS = 30_000;
const CONTINUATION_IDLE_WAIT_MS = 300_000;

/**
 * Owns all session-scoped state for pi-context-epochs.
 *
 * Design notes (verified against pi 0.85.0 sources):
 * - ctx.compact() must never be called from inside a running agent turn
 *   (AgentSession.compact() awaits abort() -> waitForIdle() -> deadlock).
 *   Model-triggered rollovers are therefore armed by the tool and executed in
 *   the agent_settled handler, when the session is idle.
 * - The durable boundary itself is created through Pi's own compaction path:
 *   session_before_compact returns a package-owned CompactionResult, Pi
 *   appends the CompactionEntry, rebuilds agent state from
 *   buildSessionContext(), and emits session_compact. No LLM call happens.
 * - firstKeptEntryId is a value that can never match a real entry id, so
 *   buildContextEntries() projects the new epoch as: [epoch marker] +
 *   everything appended after the boundary. Raw history stays in the JSONL.
 */
export class EpochController {
  private readonly pi: ExtensionAPI;
  config: EpochConfig;

  private generation = 0;
  private pending: PendingRollover | null = null;
  private expectation: RolloverExpectation | null = null;
  private continuationDue: { generation: number; details: EpochBoundaryDetails } | null = null;
  private lastRollover: { details: EpochBoundaryDetails; at: number } | null = null;
  private sentReminders = new Set<PressureBand>();
  private conflictWarned = false;
  private conflictNames = new Set<string>();
  private lastAgentStopReason: string | null = null;
  private compactTimer: ReturnType<typeof setTimeout> | null = null;

  constructor(pi: ExtensionAPI, config: EpochConfig) {
    this.pi = pi;
    this.config = config;
  }

  // =========================================================================
  // Session lifecycle
  // =========================================================================

  onSessionStart(ctx: ExtensionContext): void {
    this.generation += 1;
    this.pending = null;
    this.expectation = null;
    this.continuationDue = null;
    this.sentReminders = new Set();
    this.lastAgentStopReason = null;

    if (!this.conflictWarned) {
      this.conflictWarned = true;
      const scan = scanForConflictingPackages(ctx.cwd);
      this.conflictNames = new Set(scan.conflicts);
      if (scan.conflicts.length > 0 && ctx.hasUI) {
        ctx.ui.notify(
          `pi-context-epochs: conflicting extension(s) detected (${scan.conflicts.join(", ")}). ` +
            "Both may try to replace compaction results; the last session_before_compact handler wins. " +
            "pi-context-epochs will fail closed if its epoch boundary is overwritten.",
          "warning",
        );
      }
    }

    this.updateStatus(ctx);
  }

  onSessionShutdown(): void {
    this.generation += 1;
    this.pending = null;
    this.expectation = null;
    this.continuationDue = null;
    if (this.compactTimer) {
      clearTimeout(this.compactTimer);
      this.compactTimer = null;
    }
  }

  onSessionTree(ctx: ExtensionContext): void {
    // Branch projection changed: reset per-epoch reminder state so bands are
    // re-evaluated for the branch that is now active.
    this.sentReminders = new Set();
    this.updateStatus(ctx);
  }

  /**
   * Context-window changes shift the pressure bands. If the (possibly larger)
   * window puts the session back below the soft band, re-arm reminders;
   * bands are always re-derived from the live model's real contextWindow.
   */
  onModelSelect(ctx: ExtensionContext): void {
    const evaluation = evaluatePressure(ctx.getContextUsage(), this.config);
    if (evaluation.band === "normal" || evaluation.band === "unknown") {
      this.sentReminders = new Set();
    }
    this.updateStatus(ctx);
  }

  // =========================================================================
  // The durable primitive: session_before_compact
  // =========================================================================

  /**
   * Convert compaction events into package-owned epoch boundaries.
   * Returns undefined for user-run /compact (native summarization preserved)
   * and whenever autoRollover is disabled.
   */
  onBeforeCompact(event: {
    reason: "manual" | "threshold" | "overflow";
    customInstructions: string | undefined;
    preparation: { tokensBefore: number };
    branchEntries: SessionEntry[];
  }, ctx: ExtensionContext): { compaction: ReturnType<typeof buildEpochCompaction> } | undefined {
    if (!this.config.enabled) return undefined;

    let trigger: EpochTrigger | null = null;
    if (event.customInstructions?.startsWith(ROLLOVER_INSTRUCTIONS_PREFIX)) {
      const requested = event.customInstructions.slice(ROLLOVER_INSTRUCTIONS_PREFIX.length);
      trigger = requested === "model" || requested === "manual" ? requested : "manual";
    } else if ((event.reason === "threshold" || event.reason === "overflow") && this.config.autoRollover) {
      trigger = event.reason;
    }
    if (!trigger) return undefined;

    const notes = replayNotes(ctx.sessionManager.getBranch());
    const compaction = buildEpochCompaction({
      branch: event.branchEntries,
      tokensBefore: event.preparation.tokensBefore,
      trigger,
      sessionId: ctx.sessionManager.getSessionId() ?? null,
      notesKeys: [...notes.keys()],
      anchorMaxChars: this.config.continuationAnchorMaxChars,
    });

    this.expectation = {
      generation: this.generation,
      trigger,
      sessionFile: ctx.sessionManager.getSessionFile() ?? null,
    };
    return { compaction };
  }

  // =========================================================================
  // Verification + continuation: session_compact
  // =========================================================================

  onSessionCompact(event: { compactionEntry: CompactionEntry; fromExtension: boolean; reason: string }, ctx: ExtensionContext): void {
    const details = event.compactionEntry.details as EpochBoundaryDetails | undefined;
    const isOurs = details?.kind === "pi-context-epochs/epoch" && details.version === 1;

    if (isOurs) {
      const boundary = details as EpochBoundaryDetails;
      this.lastRollover = { details: boundary, at: Date.now() };
      this.sentReminders = new Set();
      if (this.expectation && this.expectation.generation === this.generation) {
        // Only model/manual rollovers need a continuation nudge; threshold and
        // overflow boundaries are continued by pi itself.
        if (boundary.trigger === "model" || boundary.trigger === "manual") {
          this.continuationDue = { generation: this.generation, details: boundary };
        }
      }
      this.expectation = null;
      this.updateStatus(ctx);
      return;
    }

    // Not ours: if we expected our boundary, another extension overwrote it.
    if (this.expectation) {
      this.expectation = null;
      this.conflictNames.add("(runtime) competing compaction handler");
      if (ctx.hasUI) {
        ctx.ui.notify(
          "pi-context-epochs: the compaction result was replaced by another extension; epoch boundary NOT established. No continuation was sent.",
          "warning",
        );
      }
    }
  }

  onSessionCompactFailed(ctx: ExtensionContext): void {
    this.expectation = null;
    if (this.pending) {
      const wasModel = this.pending.trigger === "model";
      this.pending = null;
      if (wasModel && ctx.hasUI) {
        ctx.ui.notify("pi-context-epochs: epoch rollover failed; staying in the current context.", "warning");
      }
    }
  }

  // =========================================================================
  // Agent run lifecycle
  // =========================================================================

  onAgentEnd(event: { messages: Array<{ stopReason?: string }> }): void {
    const last = event.messages[event.messages.length - 1];
    this.lastAgentStopReason = last?.stopReason ?? null;
  }

  /**
   * Execute an armed model-triggered rollover. Runs only when the session is
   * idle (agent_settled), because AgentSession.compact() aborts any active run.
   */
  async onAgentSettled(ctx: ExtensionContext): Promise<void> {
    if (!this.pending || this.pending.generation !== this.generation) return;
    const pending = this.pending;
    this.pending = null;

    const sessionFile = ctx.sessionManager.getSessionFile() ?? null;
    if (sessionFile !== pending.sessionFile) return;

    if (this.lastAgentStopReason === "aborted") {
      if (ctx.hasUI) {
        ctx.ui.notify("pi-context-epochs: epoch rollover cancelled (run aborted).", "info");
      }
      return;
    }

    // A threshold/overflow rollover may have raced us during the same run.
    const state = getEpochState(ctx.sessionManager.getBranch());
    if (state.epochIndex > pending.epochIndexAtArm) {
      return; // boundary already established; Pi continues on its own
    }

    await this.performRollover(ctx, pending.trigger, pending.generation);
  }

  /** Shared execution path for model-triggered and /epoch new rollovers. */
  private async performRollover(ctx: ExtensionContext, trigger: EpochTrigger, generation: number): Promise<RolloverOutcome> {
    const marker = `${ROLLOVER_INSTRUCTIONS_PREFIX}${trigger}`;
    const result = await new Promise<{ ok: true } | { ok: false; error: Error }>((resolve) => {
      let settled = false;
      this.compactTimer = setTimeout(() => {
        if (!settled) {
          settled = true;
          this.compactTimer = null;
          resolve({ ok: false, error: new Error("epoch rollover timed out") });
        }
      }, ROLLOVER_TIMEOUT_MS);
      try {
        ctx.compact({
          customInstructions: marker,
          onComplete: () => {
            if (!settled) {
              settled = true;
              if (this.compactTimer) {
                clearTimeout(this.compactTimer);
                this.compactTimer = null;
              }
              resolve({ ok: true });
            }
          },
          onError: (error) => {
            if (!settled) {
              settled = true;
              if (this.compactTimer) {
                clearTimeout(this.compactTimer);
                this.compactTimer = null;
              }
              resolve({ ok: false, error });
            }
          },
        });
      } catch (error) {
        settled = true;
        if (this.compactTimer) {
          clearTimeout(this.compactTimer);
          this.compactTimer = null;
        }
        resolve({ ok: false, error: error instanceof Error ? error : new Error(String(error)) });
      }
    });

    if (!result.ok) {
      if (ctx.hasUI) {
        ctx.ui.notify(`pi-context-epochs: rollover failed: ${result.error.message}`, "warning");
      }
      return { performed: false, reason: result.error.message };
    }

    if (generation !== this.generation) {
      return { performed: false, reason: "stale generation" };
    }

    const due = this.continuationDue;
    if (!due || due.generation !== generation) {
      return { performed: false, reason: "boundary not owned by pi-context-epochs" };
    }
    this.continuationDue = null;
    await this.sendContinuation(ctx, due.details);
    return { performed: true, details: due.details };
  }

  /**
   * Fire the continuation nudge that starts the first turn of the new epoch.
   * In headless modes we wait for the follow-up run to settle so callers like
   * `pi -p` observe a deterministic final state.
   */
  private async sendContinuation(ctx: ExtensionContext, details: EpochBoundaryDetails): Promise<void> {
    this.pi.sendMessage(
      {
        customType: CONTINUATION_MESSAGE_TYPE,
        content: buildContinuationText(details),
        display: false,
        details: { epochId: details.epochId, epochIndex: details.epochIndex, trigger: details.trigger },
      },
      { triggerTurn: true },
    );
    if (ctx.mode === "print" || ctx.mode === "json") {
      await this.waitForIdleBounded(ctx, CONTINUATION_IDLE_WAIT_MS);
    }
  }

  private async waitForIdleBounded(ctx: ExtensionContext, timeoutMs: number): Promise<void> {
    const deadline = Date.now() + timeoutMs;
    // The continuation turn starts synchronously inside pi.sendMessage; wait
    // until the session is idle again (or the deadline passes).
    await new Promise<void>((resolve) => {
      const tick = () => {
        if (Date.now() > deadline || ctx.isIdle()) {
          resolve();
          return;
        }
        setTimeout(tick, 50);
      };
      setTimeout(tick, 50);
    });
  }

  // =========================================================================
  // Arming (new_context tool, /epoch new)
  // =========================================================================

  /** Pre-checks + arm a model-triggered rollover. Never compacts here. */
  armModelRollover(ctx: ExtensionContext): ArmResult {
    if (!this.config.enabled) {
      return { armed: false, reason: "pi-context-epochs is disabled by configuration." };
    }
    if (this.pending && this.pending.generation === this.generation) {
      return {
        armed: false,
        reason:
          "A context epoch boundary is already armed and takes effect when your current response finishes. Do not call new_context again.",
      };
    }

    const branch = ctx.sessionManager.getBranch();
    const state = getEpochState(branch);
    const lastEntry = branch[branch.length - 1];
    if (state.epochIndex > 0 && lastEntry?.type === "compaction") {
      return {
        armed: false,
        reason: `Session is already at epoch ${state.epochIndex} boundary with nothing after it; no rollover needed.`,
        epochIndex: state.epochIndex,
      };
    }

    // A compaction needs something before the cut point. Pi's cut-point logic
    // (keepRecentTokens, default 20000) refuses to compact when the active
    // context is smaller than that budget, so mirror that heuristic.
    const keepRecent = this.readKeepRecentTokens(ctx);
    const activeTokens = this.estimateActiveTokens(ctx);
    if (activeTokens < keepRecent) {
      return {
        armed: false,
        reason:
          `Active context is small (~${Math.round(activeTokens / 1000)}k tokens, below the ~${Math.round(keepRecent / 1000)}k rollover threshold). ` +
          "Retiring it would not reclaim meaningful space; continue working.",
        epochIndex: state.epochIndex,
      };
    }

    this.pending = {
      generation: this.generation,
      trigger: "model",
      sessionFile: ctx.sessionManager.getSessionFile() ?? null,
      epochIndexAtArm: state.epochIndex,
    };
    return { armed: true, epochIndex: state.epochIndex + 1 };
  }

  /** Manual rollover from /epoch new. Caller must ensure idle (waitForIdle). */
  async manualRollover(ctx: ExtensionContext): Promise<RolloverOutcome> {
    if (!this.config.enabled) {
      return { performed: false, reason: "pi-context-epochs is disabled by configuration." };
    }
    const branch = ctx.sessionManager.getBranch();
    const state = getEpochState(branch);
    const lastEntry = branch[branch.length - 1];
    if (state.epochIndex > 0 && lastEntry?.type === "compaction") {
      return { performed: false, reason: `already at epoch ${state.epochIndex} boundary` };
    }
    if (this.estimateActiveTokens(ctx) < this.readKeepRecentTokens(ctx)) {
      return { performed: false, reason: "active context too small to retire" };
    }
    return this.performRollover(ctx, "manual", this.generation);
  }

  // =========================================================================
  // Context pressure
  // =========================================================================

  async onTurnEnd(ctx: ExtensionContext): Promise<void> {
    this.updateStatus(ctx);
    if (!this.config.enabled) return;

    const usage = ctx.getContextUsage();
    const evaluation = evaluatePressure(usage, this.config);
    if (evaluation.band === "soft" && !this.sentReminders.has("soft")) {
      this.sentReminders.add("soft");
      this.sendReminder(ctx, "soft", evaluation);
    } else if (evaluation.band === "critical" && !this.sentReminders.has("critical")) {
      this.sentReminders.add("critical");
      this.sentReminders.add("soft");
      // With autoRollover the critical zone is owned by pi's threshold/overflow
      // machinery, which we convert into a fresh epoch - a steered reminder
      // here would land stale in the new epoch's context.
      if (!this.config.autoRollover) {
        this.sendReminder(ctx, "critical", evaluation);
      }
    }
  }

  private sendReminder(
    ctx: ExtensionContext,
    band: "soft" | "critical",
    evaluation: ReturnType<typeof evaluatePressure>,
  ): void {
    this.pi.sendMessage(
      {
        customType: REMINDER_MESSAGE_TYPE,
        content: reminderText(band, evaluation),
        display: true,
        details: { band },
      },
      { deliverAs: "steer" },
    );
  }

  // =========================================================================
  // Helpers
  // =========================================================================

  private readKeepRecentTokens(ctx: ExtensionContext): number {
    // Pi does not expose compaction settings to extensions; mirror the default
    // and honor explicit settings.json overrides (global, then project).
    let value = DEFAULT_COMPACTION_SETTINGS.keepRecentTokens;
    try {
      const paths = [
        join(homedir(), ".pi", "agent", "settings.json"),
        join(ctx.cwd, ".pi", "settings.json"),
      ];
      for (const path of paths) {
        if (!existsSync(path)) continue;
        const settings = JSON.parse(readFileSync(path, "utf8")) as {
          compaction?: { keepRecentTokens?: number };
        };
        const v = settings?.compaction?.keepRecentTokens;
        if (typeof v === "number" && Number.isFinite(v) && v > 0) value = v;
      }
    } catch {
      // keep default
    }
    return value;
  }

  /** Estimated tokens of the active (compaction-projected) model context. */
  estimateActiveTokens(ctx: ExtensionContext): number {
    try {
      const entries = ctx.sessionManager.buildContextEntries();
      let total = 0;
      for (const entry of entries) {
        total += estimateEntryTokens(entry);
      }
      return total;
    } catch {
      return 0;
    }
  }

  getLastRollover(): { details: EpochBoundaryDetails; at: number } | null {
    return this.lastRollover;
  }

  getEpochIndex(ctx: ExtensionContext): number {
    return getEpochState(ctx.sessionManager.getBranch()).epochIndex;
  }

  isPending(): boolean {
    return this.pending !== null && this.pending.generation === this.generation;
  }

  getConflicts(): string[] {
    return [...this.conflictNames];
  }

  updateStatus(ctx: ExtensionContext): void {
    if (ctx.mode !== "tui") return;
    try {
      const state = getEpochState(ctx.sessionManager.getBranch());
      const usage = ctx.getContextUsage();
      const pct = usage && usage.percent !== null ? ` · ${Math.round(usage.percent)}%` : "";
      ctx.ui.setStatus("ctx-epochs", `E${state.epochIndex}${pct}`);
    } catch {
      // status line is best-effort
    }
  }
}

/** chars/4 heuristic mirroring pi's estimateTokens, applied to session entries. */
function estimateEntryTokens(entry: SessionEntry): number {
  let chars = 0;
  if (entry.type === "message") {
    const message = entry.message as { role: string; content: unknown };
    if (message.role === "assistant") {
      for (const block of (message.content as Array<Record<string, any>>) ?? []) {
        if (block.type === "text") chars += (block.text as string).length;
        else if (block.type === "thinking") chars += (block.thinking as string).length;
        else if (block.type === "toolCall") {
          chars += (block.name as string).length + JSON.stringify(block.arguments ?? {}).length;
        }
      }
    } else if (typeof message.content === "string") {
      chars = message.content.length;
    } else if (Array.isArray(message.content)) {
      for (const block of message.content as Array<Record<string, any>>) {
        if (block.type === "text") chars += (block.text as string).length;
        else if (block.type === "image") chars += 4800;
      }
    }
  } else if (entry.type === "custom_message") {
    chars = typeof entry.content === "string" ? entry.content.length : JSON.stringify(entry.content ?? []).length;
  } else if (entry.type === "compaction") {
    chars = entry.summary.length;
  } else if (entry.type === "branch_summary") {
    chars = entry.summary.length;
  }
  return Math.ceil(chars / 4);
}