import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import {
  PROVIDER_IDS,
  PROVIDER_LABELS,
  createFileAuthStorage,
  fetchAllQuotas,
  fetchQuota,
} from "./quotas.js";
import type { ProviderId, QuotaWindow } from "./types.js";
import type { AuthStorage } from "./quotas.js";

const STATUS_KEY = "pi-quotas",
  REFRESH_MS = 60_000,
  MS_PER_DAY = 86_400_000,
  MS_PER_HOUR = 3_600_000,
  MS_PER_MINUTE = 60_000,
  BAR_WIDTH = 20,
  BAR_STEP = 5,
  AMBER_THRESHOLD = 80;

/** Map Pi provider name to our ProviderId. */
function toProviderId(provider: string | undefined): ProviderId | undefined {
  if (!provider) {
    return undefined;
  }
  if (PROVIDER_IDS.includes(provider as ProviderId)) {
    return provider as ProviderId;
  }
  // Aliases
  if (provider === "google-antigravity" || provider === "google" || provider === "code-assist") {
    return "antigravity";
  }
  if (provider === "openai" || provider === "codex") {
    return "openai-codex";
  }
  if (provider === "xai" || provider === "grok") {
    return "xai";
  }
  if (provider === "kimi") {
    return "kimi-coding";
  }
  if (provider === "opencode" || provider === "opencode-go") {
    return "opencode-go";
  }
  return undefined;
}

function authStorage(ctx: ExtensionContext): AuthStorage {
  const reg = ctx.modelRegistry as any;
  if (typeof reg?.authStorage?.getApiKey === "function") {
    return reg.authStorage;
  }
  return createFileAuthStorage();
}

function formatResetTime(resetsAt: string): string {
  const ms = new Date(resetsAt).getTime() - Date.now();
  if (ms <= 0) {
    return "";
  }
  const days = Math.floor(ms / MS_PER_DAY),
    hours = Math.floor((ms % MS_PER_DAY) / MS_PER_HOUR),
    minutes = Math.floor((ms % MS_PER_HOUR) / MS_PER_MINUTE);
  if (days > 0) {
    return ` ${days}d ${hours}h`;
  }
  return ` ${hours > 0 ? `${hours}h ` : ""}${minutes}m`;
}

/** Shorten standard labels for compact footer display. */
function shortenFooterLabel(label: string): string {
  return label
    .replace("5-Hour", "5h")
    .replace("Weekly", "Wk")
    .replace("Monthly", "Mo")
    .replace(/Balance\s*\([^)]+\)/, "Bal");
}

/** Filter Antigravity windows relevant to current active model. */
function filterWindowsForModel(
  provider: ProviderId,
  modelId: string | undefined,
  windows: QuotaWindow[],
): QuotaWindow[] {
  if (provider !== "antigravity" || !modelId) {
    return windows;
  }
  const isClaudeOrGpt =
    modelId.includes("claude") || modelId.includes("gpt") || modelId.includes("3p");
  if (isClaudeOrGpt) {
    const matched = windows.filter((w) => w.label.startsWith("Claude"));
    return matched.length > 0 ? matched : windows;
  }
  const geminiMatched = windows.filter((w) => w.label.startsWith("Gemini"));
  return geminiMatched.length > 0 ? geminiMatched : windows;
}

// ── Footer status ────────────────────────────────────────
function formatFooter(
  ctx: ExtensionContext,
  result: Awaited<ReturnType<typeof fetchQuota>>,
): string | undefined {
  if (!result.ok) {
    return undefined;
  }
  const { theme } = ctx.ui;
  const targetWindows = filterWindowsForModel(result.provider, ctx.model?.id, result.windows);

  return targetWindows
    .map((w) => {
      const pct = Math.round(w.usedPercent),
        color = winIsLimited(w) ? "error" : pct >= AMBER_THRESHOLD ? "warning" : "success";
      const val = winIsCurrency(w)
          ? `${theme.fg(color as any, `${w.usedValue}`)}`
          : `${theme.fg(color as any, `${pct}%`)}`,
        label = theme.fg("dim", shortenFooterLabel(w.label)),
        reset =
          w.resetsAt && (winIsCurrency(w) || pct > 0)
            ? theme.fg("dim", formatResetTime(w.resetsAt))
            : "";
      return `${label} ${val}${reset}`;
    })
    .join(" ");
}

function winIsCurrency(w: QuotaWindow): boolean {
  return w.isCurrency === true;
}

function winIsLimited(w: QuotaWindow): boolean {
  return w.limited === true;
}

export default function piQuotas(pi: ExtensionAPI) {
  let timer: ReturnType<typeof setInterval> | undefined, ctx: ExtensionContext | undefined;

  async function refresh() {
    if (!ctx?.hasUI) {
      return;
    }
    const provider = toProviderId(ctx.model?.provider);
    if (!provider) {
      ctx.ui.setStatus(STATUS_KEY, undefined);
      return;
    }
    try {
      const result = await fetchQuota(authStorage(ctx), provider);
      if (!ctx?.hasUI) {
        return;
      }
      const text = formatFooter(ctx, result);
      ctx.ui.setStatus(STATUS_KEY, text);
    } catch {
      // Swallow stale context or network errors
    }
  }

  pi.on("session_start", (_ev, newCtx) => {
    ctx = newCtx;
    if (timer) {
      clearInterval(timer);
    }
    timer = setInterval(() => void refresh(), REFRESH_MS);
    timer.unref?.();
    void refresh();
  });

  pi.on("turn_end", (_ev, newCtx) => {
    ctx = newCtx;
    void refresh();
  });

  pi.on("model_select", (_ev, newCtx) => {
    ctx = newCtx;
    void refresh();
  });

  pi.on("session_shutdown", (_ev, shutdownCtx) => {
    ctx = undefined;
    if (timer) {
      clearInterval(timer);
    }
    timer = undefined;
    try {
      shutdownCtx.ui.setStatus(STATUS_KEY, undefined);
    } catch {
      /* Stale */
    }
  });

  // ── /quotas & /usage command ──────────────────────────────
  const handleQuotasCommand = async (_args: string, cmdCtx: any) => {
    const auth = authStorage(cmdCtx),
      results = await fetchAllQuotas(auth, { force: true }),
      lines: string[] = [];

    for (const result of results) {
      const label = PROVIDER_LABELS[result.provider];
      if (!result.ok) {
        lines.push(`${label}: ${result.error}`);
        continue;
      }
      if (result.windows.length === 0) {
        continue;
      }
      const summary = result.windows
        .map((w) => {
          if (winIsCurrency(w)) {
            return `${w.label}: ${w.usedValue}`;
          }
          const pct = Math.round(w.usedPercent),
            filled = Math.floor(pct / BAR_STEP),
            bar = "\u2588".repeat(filled) + "\u2591".repeat(BAR_WIDTH - filled),
            reset = w.resetsAt && (winIsCurrency(w) || pct > 0) ? formatResetTime(w.resetsAt) : "",
            ltd = winIsLimited(w) ? " (LIMIT REACHED)" : "";
          return `  ${w.label}: ${bar} ${pct}%${reset}${ltd}`;
        })
        .join("\n");
      lines.push(`${label}:\n${summary}`);
    }

    cmdCtx.ui.notify(lines.join("\n\n") || "No quota data available", "info");
  };

  pi.registerCommand("quotas", {
    description: "Show remaining quotas for all providers",
    handler: handleQuotasCommand,
  });

  pi.registerCommand("usage", {
    description: "Show remaining quotas for all providers (alias for /quotas)",
    handler: handleQuotasCommand,
  });
}