/**
 * Usage Extension - Minimal API usage indicator for pi
 *
 * Shows Codex (OpenAI), Anthropic (Claude), Z.AI, and optionally
 * Google Gemini CLI / Antigravity usage as color-coded percentages
 * in the footer status bar.
 */

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;           // reuse cached data for 15 min
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; // 1 hour back-off after 429
const STATUS_KEY = "usage-bars";

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

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 renderBar(pct: number, width = 16): string {
    const value = clampPercent(pct);
    const filled = Math.round((value / 100) * width);
    const color = colorForPercent(value);
    const full = "█".repeat(Math.max(0, filled));
    const empty = "░".repeat(Math.max(0, width - filled));
    return this.theme.fg(color, full) + this.theme.fg("dim", empty);
  }

  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  ") +
            this.renderBar(session) +
            " " +
            t.fg(colorForPercent(session), `${session}%`.padStart(4)) +
            sessionReset,
          0,
          0,
        ),
      );

      this.listContainer.addChild(
        new Text(
          indent +
            t.fg("muted", "Weekly   ") +
            this.renderBar(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();
  }
}

interface UsageState extends UsageByProvider {
  lastPoll: number;
  activeProvider: ProviderKey | null;
}

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 ctx: any = null;

  function renderPercent(theme: any, value: number): string {
    const v = clampPercent(value);
    return theme.fg(colorForPercent(v), `${v}%`);
  }

  function renderBar(theme: any, value: number): string {
    const v = clampPercent(value);
    const width = 8;
    const filled = Math.round((v / 100) * width);
    const full = "█".repeat(Math.max(0, Math.min(width, filled)));
    const empty = "░".repeat(Math.max(0, width - filled));
    return theme.fg(colorForPercent(v), full) + theme.fg("dim", empty);
  }

  function pickDataForProvider(provider: ProviderKey | null): UsageData | null {
    if (!provider) return null;
    return state[provider];
  }

  function updateStatus() {
    const active = state.activeProvider;
    const data = pickDataForProvider(active);

    if (data && !data.error) {
      pi.events.emit("usage:update", {
        session: data.session,
        weekly: data.weekly,
        sessionResetsIn: data.sessionResetsIn,
        weeklyResetsIn: data.weeklyResetsIn,
      });
    }

    if (!ctx?.hasUI) return;

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

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

    const theme = ctx.ui.theme;
    const label = PROVIDER_LABELS[active];

    if (!data) {
      ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", `${label} usage: loading…`));
      return;
    }

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

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

    const sessionReset = data.sessionResetsIn ? theme.fg("dim", ` ⟳ ${data.sessionResetsIn}`) : "";
    const weeklyReset = data.weeklyResetsIn ? theme.fg("dim", ` ⟳ ${data.weeklyResetsIn}`) : "";

    const status =
      theme.fg("dim", `${label} `) +
      theme.fg("muted", "S ") +
      renderBar(theme, session) +
      " " +
      renderPercent(theme, session) +
      sessionReset +
      theme.fg("muted", " W ") +
      renderBar(theme, weekly) +
      " " +
      renderPercent(theme, weekly) +
      weeklyReset;

    ctx.ui.setStatus(STATUS_KEY, status);
  }

  function updateProviderFrom(modelLike: any): boolean {
    const previous = state.activeProvider;
    state.activeProvider = detectProvider(modelLike);

    if (previous !== state.activeProvider) {
      updateStatus();
      return true;
    }

    return false;
  }

  async function runPoll() {
    const auth = readAuth();
    const active = state.activeProvider;

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

    // --- Shared disk cache check ---
    // All pi sessions read and write the same cache file so that only one
    // process hits the API per CACHE_TTL_MS window, no matter how many
    // sessions are open at once.
    const cache = readUsageCache();
    const now = Date.now();

    // Respect cross-session rate-limit back-off written by any session.
    const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
    if (now < blockedUntil) {
      // Use whatever data is in the cache so the bar still shows something.
      if (cache?.data?.[active]) state[active] = cache.data[active]!;
      state.lastPoll = now;
      updateStatus();
      return;
    }

    // If another session already polled recently, use their result.
    if (cache && now - cache.timestamp < CACHE_TTL_MS && cache.data?.[active]) {
      state[active] = cache.data[active]!;
      state.lastPoll = now;
      updateStatus();
      return;
    }

    // --- Proactive token refresh ---
    // Before hitting the API, check whether the stored access token is expired.
    // This is the main cause of HTTP 401 errors: switching accounts via
    // /switch-claude restores a profile whose access token has since expired
    // (the refresh token is still valid). We use pi's own OAuth resolver so
    // the new tokens are written back to auth.json and the profile stays in
    // sync. This is safe at turn_start / session_start because pi hasn't made
    // any Claude API calls yet, so there's no parallel refresh to conflict with.
    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 refreshed = await ensureFreshAuthForProviders([oauthId as OAuthProviderId], {
            auth,
            persist: true,
          });
          if (refreshed.auth) effectiveAuth = refreshed.auth;
        } catch {
          // Ignore refresh errors — fall through with existing auth
        }
      }
    }

    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;

    // Write result + rate-limit state to shared cache so other sessions
    // don't need to re-hit the API within CACHE_TTL_MS.
    //
    // Error results (other than 429) are NOT cached: they should be retried
    // on the next input instead of being replayed from cache for 15 minutes.
    // The most common error is HTTP 401 (expired token after an account switch)
    // which resolves on the very next poll once the token is refreshed above.
    if (result.error) {
      if (result.error === "HTTP 429") {
        // Write rate-limit backoff but preserve the last good data in cache.
        const nextCache: import("./core").UsageCache = {
          timestamp: cache?.timestamp ?? now,
          data: { ...(cache?.data ?? {}) },
          rateLimitedUntil: {
            ...(cache?.rateLimitedUntil ?? {}),
            [active]: now + RATE_LIMITED_BACKOFF_MS,
          },
        };
        writeUsageCache(nextCache);
      }
      // All other errors: don't update cache — next turn will retry from scratch.
    } 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 poll() {
    if (pollInFlight) {
      pollQueued = true;
      await pollInFlight;
      return;
    }

    do {
      pollQueued = false;
      pollInFlight = runPoll()
        .catch(() => {
          // Never crash extension event handlers on transient polling errors.
        })
        .finally(() => {
          pollInFlight = null;
        });

      await pollInFlight;
    } while (pollQueued);
  }

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

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

  // Refresh usage on every turn (like claude-pulse's UserPromptSubmit hook).
  // The disk cache means the API is only hit at most once per CACHE_TTL_MS
  // regardless of how many turns or sessions are active.
  pi.on("turn_start", async (_event, _ctx) => {
    ctx = _ctx;
    updateProviderFrom(_ctx.model);
    await poll();
  });

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

  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) => {
            const selector = new UsageSelectorComponent(
              tui,
              theme,
              state.activeProvider,
              () => fetchAllUsages({ endpoints }),
              () => done(),
            );
            return selector;
          });
        }
      } finally {
        await poll();
      }
    },
  });

  // Re-poll immediately when the Claude account is switched via /switch-claude
  pi.events.on("claude-account:switched", () => {
    void poll();
  });
}