import { FooterComponent, type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
import { truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";

const STATUS_KEY = "pi-throttle";
const FOOTER_STATE_KEY = Symbol.for("pi-throttle.footer-state");
const MAX_TIMER_MS = 2_147_483_647;

type FooterState = {
  patched?: boolean;
  statusText?: string;
};

const footerState = ((globalThis as unknown as Record<symbol, FooterState>)[FOOTER_STATE_KEY] ??= {});

function installFooterPatch(): void {
  if (footerState.patched) return;

  const originalRender = FooterComponent.prototype.render;
  FooterComponent.prototype.render = function renderWithThrottleStatus(width: number): string[] {
    const lines = originalRender.call(this, width);
    const statusText = footerState.statusText;
    if (!statusText || lines.length === 0) return lines;

    const statusWidth = visibleWidth(statusText);
    const availableForLeft = width - statusWidth - 2;
    if (availableForLeft < 1) return lines;

    const left = truncateToWidth(lines[0], availableForLeft, availableForLeft >= 3 ? "..." : "");
    const padding = " ".repeat(Math.max(2, width - visibleWidth(left) - statusWidth));

    return [`${left}${padding}${statusText}`, ...lines.slice(1)];
  };

  footerState.patched = true;
}

function formatSeconds(seconds: number): string {
  return `${seconds.toString()}s`;
}

export default function piThrottle(pi: ExtensionAPI) {
  installFooterPatch();
  let throttleSeconds = 0;
  let configurationVersion = 0;
  let lastProviderCallAt: number | undefined;
  const pendingDelays = new Set<AbortController>();

  const cancelPendingDelays = () => {
    for (const controller of pendingDelays) controller.abort();
    pendingDelays.clear();
  };

  const updateStatus = (ctx: ExtensionContext) => {
    footerState.statusText =
      throttleSeconds > 0 ? `throttle: ${formatSeconds(throttleSeconds)}` : undefined;
    // Clearing even an absent status requests a render without adding Pi's status row.
    ctx.ui.setStatus(STATUS_KEY, undefined);
  };

  const delay = (milliseconds: number, signal?: AbortSignal): Promise<void> => {
    const controller = new AbortController();
    pendingDelays.add(controller);

    return new Promise((resolve) => {
      let remaining = milliseconds;
      let timer: ReturnType<typeof setTimeout> | undefined;

      const finish = () => {
        if (timer !== undefined) clearTimeout(timer);
        signal?.removeEventListener("abort", finish);
        controller.signal.removeEventListener("abort", finish);
        pendingDelays.delete(controller);
        resolve();
      };

      const scheduleNextChunk = () => {
        if (remaining <= 0) {
          finish();
          return;
        }

        const chunk = Math.min(remaining, MAX_TIMER_MS);
        remaining -= chunk;
        timer = setTimeout(scheduleNextChunk, chunk);
      };

      if (signal?.aborted) {
        finish();
        return;
      }

      signal?.addEventListener("abort", finish, { once: true });
      controller.signal.addEventListener("abort", finish, { once: true });
      scheduleNextChunk();
    });
  };

  pi.registerCommand("throttle", {
    description: "Set the minimum interval between provider calls in seconds (0 disables)",
    handler: async (args, ctx) => {
      const rawSeconds = args.trim();
      const seconds = Number(rawSeconds);

      if (
        rawSeconds.length === 0 ||
        !Number.isFinite(seconds) ||
        seconds < 0 ||
        seconds * 1000 > Number.MAX_SAFE_INTEGER
      ) {
        ctx.ui.notify("Usage: /throttle <non-negative number of seconds>", "error");
        return;
      }

      cancelPendingDelays();
      throttleSeconds = seconds;
      configurationVersion += 1;
      updateStatus(ctx);

      if (seconds === 0) {
        ctx.ui.notify("Provider-call throttling disabled", "info");
      } else {
        ctx.ui.notify(`Provider calls will start at least ${formatSeconds(seconds)} apart`, "info");
      }
    },
  });

  pi.on("before_provider_request", async (_event, ctx) => {
    let showingWaiting = false;

    try {
      while (!ctx.signal?.aborted) {
        const version = configurationVersion;
        const intervalMilliseconds = throttleSeconds * 1000;
        const now = Date.now();
        const elapsed = lastProviderCallAt === undefined ? Infinity : now - lastProviderCallAt;
        const remaining = Math.max(0, intervalMilliseconds - elapsed);

        if (remaining > 0) {
          if (!showingWaiting) {
            ctx.ui.setWorkingMessage("...Waiting");
            showingWaiting = true;
          }
          await delay(remaining, ctx.signal);
        }
        if (ctx.signal?.aborted) return;

        // A /throttle command cancels pending delays. Recalculate using its new value.
        if (version !== configurationVersion) continue;

        lastProviderCallAt = Date.now();
        return;
      }
    } finally {
      if (showingWaiting) ctx.ui.setWorkingMessage();
    }
  });

  pi.on("message_end", (event) => {
    if (event.message.role === "assistant") lastProviderCallAt = Date.now();
  });

  pi.on("session_start", (_event, ctx) => {
    updateStatus(ctx);
  });

  pi.on("session_shutdown", () => {
    cancelPendingDelays();
    footerState.statusText = undefined;
  });
}