/**
 * Usage Extension - Minimal API usage indicator for pi
 *
 * Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it
 * via two channels:
 *   • pi.events "usage:update" — for other extensions (e.g. footer-display)
 *   • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars
 *
 * Rendering / footer layout is handled by the separate footer-display extension.
 */

import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
import {
  Container,
  Input,
  Spacer,
  Text,
  getEditorKeybindings,
  type Focusable,
} from "@mariozechner/pi-tui";
import {
  canShowForProvider,
  clampPercent,
  colorForPercent,
  detectProvider,
  ensureFreshAuthForProviders,
  fetchAllUsages,
  fetchClaudeUsage,
  fetchCodexUsage,
  fetchGoogleUsage,
  fetchZaiUsage,
  providerToOAuthProviderId,
  readAuth,
  readUsageCache,
  resolveUsageEndpoints,
  writeUsageCache,
  type OAuthProviderId,
  type ProviderKey,
  type UsageByProvider,
  type UsageData,
} from "./core";

const CACHE_TTL_MS = 15 * 60 * 1000;
const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000;

const STATUS_KEY = "usage-bars";

// ---------------------------------------------------------------------------
// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿)
// ---------------------------------------------------------------------------
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
const BRAILLE_EMPTY = "\u28C0";
const BAR_WIDTH = 5;

function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string {
  const v = clampPercent(value);
  const levels = BRAILLE_GRADIENT.length - 1;
  const totalSteps = width * levels;
  const filledSteps = Math.round((v / 100) * totalSteps);
  const full = Math.floor(filledSteps / levels);
  const partial = filledSteps % levels;
  const empty = width - full - (partial ? 1 : 0);
  const color = colorForPercent(v);
  const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
  const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
  const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
  return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
}

function renderBrailleBarWide(theme: any, value: number): string {
  return renderBrailleBar(theme, value, 12);
}

const PROVIDER_LABELS: Record<ProviderKey, string> = {
  codex: "Codex",
  claude: "Claude",
  zai: "Z.AI",
  gemini: "Gemini",
  antigravity: "Antigravity",
};

// ---------------------------------------------------------------------------
// /usage command popup
// ---------------------------------------------------------------------------
interface SubscriptionItem {
  name: string;
  provider: ProviderKey;
  data: UsageData | null;
  isActive: boolean;
}

class UsageSelectorComponent extends Container implements Focusable {
  private searchInput: Input;
  private listContainer: Container;
  private hintText: Text;
  private tui: any;
  private theme: any;
  private onCancelCallback: () => void;
  private allItems: SubscriptionItem[] = [];
  private filteredItems: SubscriptionItem[] = [];
  private selectedIndex = 0;
  private loading = true;
  private activeProvider: ProviderKey | null;
  private fetchAllFn: () => Promise<UsageByProvider>;
  private _focused = false;

  get focused(): boolean { return this._focused; }
  set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; }

  constructor(
    tui: any,
    theme: any,
    activeProvider: ProviderKey | null,
    fetchAll: () => Promise<UsageByProvider>,
    onCancel: () => void,
  ) {
    super();
    this.tui = tui;
    this.theme = theme;
    this.activeProvider = activeProvider;
    this.fetchAllFn = fetchAll;
    this.onCancelCallback = onCancel;

    this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
    this.addChild(new Spacer(1));
    this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
    this.addChild(this.hintText);
    this.addChild(new Spacer(1));
    this.searchInput = new Input();
    this.addChild(this.searchInput);
    this.addChild(new Spacer(1));
    this.listContainer = new Container();
    this.addChild(this.listContainer);
    this.addChild(new Spacer(1));
    this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));

    this.fetchAllFn()
      .then((results) => {
        this.loading = false;
        this.buildItems(results);
        this.updateList();
        this.hintText.setText(
          theme.fg("muted", "Only showing providers with credentials. ") +
            theme.fg("dim", "✓ = active provider"),
        );
        this.tui.requestRender();
      })
      .catch(() => {
        this.loading = false;
        this.hintText.setText(theme.fg("error", "Failed to fetch usage data"));
        this.tui.requestRender();
      });

    this.updateList();
  }

  private buildItems(results: UsageByProvider) {
    const providers: Array<{ key: ProviderKey; name: string }> = [
      { key: "codex", name: "Codex" },
      { key: "claude", name: "Claude" },
      { key: "zai", name: "Z.AI" },
      { key: "gemini", name: "Gemini" },
      { key: "antigravity", name: "Antigravity" },
    ];
    this.allItems = [];
    for (const p of providers) {
      if (results[p.key] !== null) {
        this.allItems.push({
          name: p.name,
          provider: p.key,
          data: results[p.key],
          isActive: this.activeProvider === p.key,
        });
      }
    }
    this.filteredItems = this.allItems;
    this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
  }

  private filterItems(query: string) {
    if (!query) {
      this.filteredItems = this.allItems;
    } else {
      const q = query.toLowerCase();
      this.filteredItems = this.allItems.filter(
        (item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q),
      );
    }
    this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
  }

  private renderItem(item: SubscriptionItem, isSelected: boolean) {
    const t = this.theme;
    const pointer = isSelected ? t.fg("accent", "→ ") : "  ";
    const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
    const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
    this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
    const indent = "    ";

    if (!item.data) {
      this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0));
    } else if (item.data.error) {
      this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0));
    } else {
      const session = clampPercent(item.data.session);
      const weekly = clampPercent(item.data.weekly);
      const sessionReset = item.data.sessionResetsIn
        ? t.fg("dim", `  resets in ${item.data.sessionResetsIn}`) : "";
      const weeklyReset = item.data.weeklyResetsIn
        ? t.fg("dim", `  resets in ${item.data.weeklyResetsIn}`) : "";

      this.listContainer.addChild(new Text(
        indent + t.fg("muted", "Session  ") +
          renderBrailleBarWide(t, session) + " " +
          t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset,
        0, 0,
      ));
      this.listContainer.addChild(new Text(
        indent + t.fg("muted", "Weekly   ") +
          renderBrailleBarWide(t, weekly) + " " +
          t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset,
        0, 0,
      ));

      if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
        this.listContainer.addChild(new Text(
          indent + t.fg("muted", "Extra    ") +
            t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
          0, 0,
        ));
      }
    }
    this.listContainer.addChild(new Spacer(1));
  }

  private updateList() {
    this.listContainer.clear();
    if (this.loading) {
      this.listContainer.addChild(new Text(this.theme.fg("muted", "  Loading…"), 0, 0));
      return;
    }
    if (this.filteredItems.length === 0) {
      this.listContainer.addChild(new Text(this.theme.fg("muted", "  No matching providers"), 0, 0));
      return;
    }
    for (let i = 0; i < this.filteredItems.length; i++) {
      this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
    }
  }

  handleInput(keyData: string): void {
    const kb = getEditorKeybindings();
    if (kb.matches(keyData, "selectUp")) {
      if (this.filteredItems.length === 0) return;
      this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
      this.updateList(); return;
    }
    if (kb.matches(keyData, "selectDown")) {
      if (this.filteredItems.length === 0) return;
      this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
      this.updateList(); return;
    }
    if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
      this.onCancelCallback(); return;
    }
    this.searchInput.handleInput(keyData);
    this.filterItems(this.searchInput.getValue());
    this.updateList();
  }
}

// ---------------------------------------------------------------------------
// Extension state
// ---------------------------------------------------------------------------
interface UsageState extends UsageByProvider {
  lastPoll: number;
  activeProvider: ProviderKey | null;
}

interface PollOptions {
  cacheTtl?: number;
  forceFresh?: boolean;
}

export default function (pi: ExtensionAPI) {
  const endpoints = resolveUsageEndpoints();
  const state: UsageState = {
    codex: null, claude: null, zai: null, gemini: null, antigravity: null,
    lastPoll: 0, activeProvider: null,
  };

  let pollInFlight: Promise<void> | null = null;
  let pollQueued = false;
  let pollStartedAt = 0;
  let streamingTimer: ReturnType<typeof setInterval> | null = null;
  let ctx: any = null;

  // ---------------------------------------------------------------------------
  // Status update
  // ---------------------------------------------------------------------------
  function updateStatus() {
    const active = state.activeProvider;
    const data = active ? state[active] : null;

    // Always emit event for other extensions (e.g. footer-display)
    if (data && !data.error) {
      pi.events.emit("usage:update", {
        session: data.session,
        weekly: data.weekly,
        sessionResetsIn: data.sessionResetsIn,
        sessionResetsAt: data.sessionResetsAt,
        weeklyResetsIn: data.weeklyResetsIn,
      });
    }

    if (!ctx?.hasUI) return;

    const theme = ctx.ui.theme;

    if (!active) {
      ctx.ui.setStatus(STATUS_KEY, undefined);
      return;
    }

    const auth = readAuth();
    if (!canShowForProvider(active, auth, endpoints)) {
      ctx.ui.setStatus(STATUS_KEY, undefined);
      return;
    }

    if (!data) {
      ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026"));
      return;
    }

    if (data.error) {
      const cache = readUsageCache();
      const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
      const note = blockedUntil > Date.now()
        ? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : "";
      ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`));
      return;
    }

    const session = clampPercent(data.session);
    const weekly = clampPercent(data.weekly);

    let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
    if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn);

    let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
    if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`);

    ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w);
  }

  function updateProviderFrom(modelLike: any): boolean {
    const previous = state.activeProvider;
    state.activeProvider = detectProvider(modelLike);
    if (previous !== state.activeProvider) { updateStatus(); return true; }
    return false;
  }

  // ---------------------------------------------------------------------------
  // Polling
  // ---------------------------------------------------------------------------
  async function runPollInner(options: PollOptions = {}) {
    const auth = readAuth();
    const active = state.activeProvider;

    if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
      state.lastPoll = Date.now(); updateStatus(); return;
    }

    const cache = readUsageCache();
    const now = Date.now();
    const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;

    const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
    if (now < blockedUntil) {
      if (cache?.data?.[active]) {
        state[active] = cache.data[active]!;
      } else {
        // Rate-limited but no cached data — show a meaningful status instead
        // of leaving state null (which shows eternal "loading…").
        const retryMin = Math.ceil((blockedUntil - now) / 60000);
        state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` };
      }
      state.lastPoll = now; updateStatus(); return;
    }

    if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
      state[active] = cache.data[active]!;
      state.lastPoll = now; updateStatus(); return;
    }

    const oauthId = providerToOAuthProviderId(active);
    let effectiveAuth = auth;
    if (oauthId && active !== "zai") {
      const creds = auth[oauthId as keyof typeof auth] as
        | { access?: string; refresh?: string; expires?: number } | undefined;
      const expires = typeof creds?.expires === "number" ? creds.expires : 0;
      const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
      if (tokenExpiredOrMissing && creds?.refresh) {
        try {
          const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true });
          const timeoutPromise = new Promise<never>((_, reject) =>
            setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000),
          );
          const refreshed = await Promise.race([refreshPromise, timeoutPromise]);
          if (refreshed.auth) effectiveAuth = refreshed.auth;
        } catch {}
      }
    }

    let result: UsageData;
    if (active === "codex") {
      const access = effectiveAuth["openai-codex"]?.access;
      result = access ? await fetchCodexUsage(access)
        : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
    } else if (active === "claude") {
      const access = effectiveAuth.anthropic?.access;
      result = access ? await fetchClaudeUsage(access)
        : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
    } else if (active === "zai") {
      const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
      result = token ? await fetchZaiUsage(token, { endpoints })
        : { session: 0, weekly: 0, error: "missing token (try /login again)" };
    } else if (active === "gemini") {
      const creds = effectiveAuth["google-gemini-cli"];
      result = creds?.access
        ? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
        : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
    } else {
      const creds = effectiveAuth["google-antigravity"];
      result = creds?.access
        ? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
        : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
    }

    state[active] = result;

    if (result.error) {
      if (result.error === "HTTP 429") {
        const nextCache: import("./core").UsageCache = {
          timestamp: cache?.timestamp ?? now,
          data: { ...(cache?.data ?? {}) },
          rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS },
        };
        writeUsageCache(nextCache);
      }
    } else {
      const nextCache: import("./core").UsageCache = {
        timestamp: now,
        data: { ...(cache?.data ?? {}), [active]: result },
        rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
      };
      delete nextCache.rateLimitedUntil![active];
      writeUsageCache(nextCache);
    }

    state.lastPoll = now;
    updateStatus();
  }

  async function runPoll(options: PollOptions = {}): Promise<void> {
    const timeout = new Promise<never>((_, reject) =>
      setTimeout(() => reject(new Error("runPoll timeout")), 25_000),
    );
    await Promise.race([runPollInner(options), timeout]);
  }

  const POLL_TIMEOUT_MS = 30_000;

  async function poll(options: PollOptions = {}) {
    // If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it
    // so we don't queue forever behind a stuck request.
    if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) {
      pollInFlight = null;
      pollQueued = false;
      const active = state.activeProvider;
      if (active && !state[active]) {
        state[active] = { session: 0, weekly: 0, error: "poll timeout" };
        updateStatus();
      }
    }

    if (pollInFlight) { pollQueued = true; await pollInFlight; return; }
    do {
      pollQueued = false;
      pollStartedAt = Date.now();
      pollInFlight = runPoll(options).catch(() => {
        // If runPoll threw, ensure we don't leave status stuck at "loading…"
        const active = state.activeProvider;
        if (active && !state[active]) {
          state[active] = { session: 0, weekly: 0, error: "poll failed" };
          updateStatus();
        }
      }).finally(() => { pollInFlight = null; pollStartedAt = 0; });
      await pollInFlight;
    } while (pollQueued);
  }

  function startStreamingTimer() {
    if (streamingTimer !== null) return;
    streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS);
  }

  function stopStreamingTimer() {
    if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; }
  }

  // ── Lifecycle ────────────────────────────────────────────────────────────

  pi.on("session_start", async (_event, _ctx) => {
    ctx = _ctx;
    updateProviderFrom(_ctx.model);
    await poll();
  });

  pi.on("session_shutdown", async (_event, _ctx) => {
    stopStreamingTimer();
    if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined);
  });

  pi.on("model_select", async (event, _ctx) => {
    ctx = _ctx;
    const changed = updateProviderFrom(event.model ?? _ctx.model);
    if (changed) await poll();
  });

  pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); });

  pi.on("before_agent_start", async (_event, _ctx) => {
    ctx = _ctx;
    await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
  });

  pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); });

  pi.on("agent_end", async (_event, _ctx) => {
    ctx = _ctx;
    stopStreamingTimer();
    await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
  });

  pi.events.on("claude-account:switched", () => {
    const cache = readUsageCache();
    if (cache?.data?.claude) {
      const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } };
      delete nextCache.data.claude;
      writeUsageCache(nextCache);
    }
    void poll({ forceFresh: true });
  });

  // ── /usage command ───────────────────────────────────────────────────────

  pi.registerCommand("usage", {
    description: "Show API usage for all subscriptions",
    handler: async (_args, _ctx) => {
      ctx = _ctx;
      updateProviderFrom(_ctx.model);
      try {
        if (_ctx?.hasUI) {
          await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
            return new UsageSelectorComponent(
              tui, theme, state.activeProvider,
              () => fetchAllUsages({ endpoints }),
              () => done(),
            );
          });
        }
      } finally {
        await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
      }
    },
  });
}
