import type { Command, CompactionCap, ExtensionApi, JsonValue } from "@berry/api";

/** Blocking compaction at this threshold unless `/context-cap` changes it. */
const DEFAULT_MAX_CONTEXT_TOKENS = 2 ** 18;
/** Background compaction this far below the blocking threshold by default. */
const DEFAULT_BACKGROUND_TOKENS = 32_768;

const K = 1024;
const M = 1024 * 1024;

// server-wide, so every conversation compacts at one threshold; a reload restores the defaults
let enabled = false;
let maxContextTokens = DEFAULT_MAX_CONTEXT_TOKENS;
let backgroundTokens = DEFAULT_BACKGROUND_TOKENS;

/** A token count as an integer with an optional `k` or `m` suffix, so 2^18 reads `256k`. */
function parseTokens(text: string): number | undefined {
  const match = /^(\d+(?:\.\d+)?)\s*([kKmM])?$/.exec(text.trim());
  if (!match) return undefined;
  const scale = match[2] === undefined ? 1 : match[2].toLowerCase() === "m" ? M : K;
  const value = Math.floor(Number(match[1]) * scale);
  return Number.isFinite(value) && value > 0 ? value : undefined;
}

function status(): string {
  if (!enabled) return "context cap off: compaction follows each model's own window";
  const backgroundAt = maxContextTokens - backgroundTokens;
  return `context cap: blocking compaction at ${maxContextTokens} tokens, ` +
    `background at ${backgroundAt}; models with smaller windows keep their own policy`;
}

/** Cap effective context windows for this load; tune it with the context-cap command. */
export function registerContextCap<S extends JsonValue, E extends JsonValue>(
  ext: ExtensionApi<S, E>,
): void {
  ext.compactionCap((): CompactionCap | undefined =>
    enabled ? { maxContextTokens, backgroundTokens } : undefined
  );
}

/** Show or tune the context cap: `/context-cap [tokens] [background]`, `on`, or `off`. */
export function contextCapCommand(): Command<JsonValue, JsonValue> {
  return {
    name: "context-cap",
    description:
      "show or tune the context cap: `/context-cap [tokens] [background]`, `on`, or `off`",
    execute: (args) => {
      const parts = args.trim().split(/\s+/).filter((part) => part !== "");
      if (parts.length > 2) throw new Error("expected at most a threshold and a background offset");
      const [first = "", second = ""] = parts;
      if (first === "") return status();
      if (first === "off") {
        enabled = false;
        return status();
      }
      if (first === "on") {
        enabled = true;
        return status();
      }
      const max = parseTokens(first);
      if (max === undefined) {
        throw new Error(`expected a token count like 262144 or 256k, got ${JSON.stringify(first)}`);
      }
      let background = backgroundTokens;
      if (second !== "") {
        const parsed = parseTokens(second);
        if (parsed === undefined) {
          throw new Error(
            `expected a token count like 32768 or 32k, got ${JSON.stringify(second)}`,
          );
        }
        background = parsed;
      }
      if (background >= max) {
        throw new Error(
          `background offset ${background} must stay below the ${max}-token threshold`,
        );
      }
      maxContextTokens = max;
      backgroundTokens = background;
      enabled = true;
      return status();
    },
  };
}