codexbar-mobile/hooks/useCodexUsage.ts

import { useCallback, useEffect, useRef, useState } from "react";
import {
  fetchCodexUsage,
  parseCodexUsageResponse,
} from "@/lib/api/codexApi";
import { RETRY_DELAY_MS } from "@/lib/constants";
import { pickAndReadCodexAuth } from "@/lib/fileReader";
import {
  clearCodexAuth,
  clearCodexUsageCache,
  loadCodexAuth,
  loadCodexUsageCache,
  saveCodexAuth,
  saveCodexUsageCache,
} from "@/lib/storage";
import type { CodexAuth, CodexUsageResponse } from "@/types/codex";

type Status = "idle" | "loading" | "success" | "error";

function shouldRetry(error: unknown): boolean {
  const message = error instanceof Error ? error.message : "";
  return (
    message !== "TOKEN_EXPIRED" &&
    message !== "INVALID_CODEX_RESPONSE"
  );
}

function wait(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

export function useCodexUsage() {
  const [auth, setAuth] = useState<CodexAuth | null>(null);
  const [usage, setUsage] = useState<CodexUsageResponse | null>(null);
  const [lastFetchedAt, setLastFetchedAt] = useState<number | null>(null);
  const [status, setStatus] = useState<Status>("idle");
  const [error, setError] = useState<string | null>(null);
  const mounted = useRef(true);
  const inFlight = useRef<{
    key: string;
    promise: Promise<void>;
  } | null>(null);
  const requestVersion = useRef(0);

  useEffect(() => {
    mounted.current = true;
    return () => {
      mounted.current = false;
    };
  }, []);

  const fetchUsage = useCallback((credentials: CodexAuth): Promise<void> => {
    const requestKey = `${credentials.accountId ?? ""}:${credentials.accessToken}`;
    if (inFlight.current?.key === requestKey) return inFlight.current.promise;
    const version = ++requestVersion.current;

    const promise = (async () => {
      if (mounted.current) {
        setStatus("loading");
        setError(null);
      }

      try {
        let data: CodexUsageResponse;
        try {
          data = await fetchCodexUsage(credentials);
        } catch (firstError) {
          if (!shouldRetry(firstError)) throw firstError;
          await wait(RETRY_DELAY_MS);
          if (version !== requestVersion.current) return;
          data = await fetchCodexUsage(credentials);
        }

        if (version !== requestVersion.current) return;
        const fetchedAt = Date.now();
        await saveCodexUsageCache({ data, lastFetchedAt: fetchedAt });
        if (mounted.current && version === requestVersion.current) {
          setUsage(data);
          setLastFetchedAt(fetchedAt);
          setStatus("success");
        }
      } catch (caught: unknown) {
        if (version !== requestVersion.current) return;
        const message = caught instanceof Error ? caught.message : "UNKNOWN_ERROR";
        if (message === "TOKEN_EXPIRED") {
          await clearCodexAuth();
          if (mounted.current) {
            setAuth(null);
            setUsage(null);
            setLastFetchedAt(null);
          }
        }
        if (mounted.current) {
          setError(message);
          setStatus("error");
        }
      }
    })().finally(() => {
      if (inFlight.current?.promise === promise) inFlight.current = null;
    });

    inFlight.current = { key: requestKey, promise };
    return promise;
  }, []);

  useEffect(() => {
    let active = true;

    void Promise.all([loadCodexAuth(), loadCodexUsageCache()]).then(
      async ([storedAuth, cache]) => {
        if (!active) return;

        if (storedAuth && cache) {
          try {
            if (!Number.isFinite(cache.lastFetchedAt)) {
              throw new Error("INVALID_CACHE_TIMESTAMP");
            }
            const cachedUsage = parseCodexUsageResponse(cache.data);
            setUsage(cachedUsage);
            setLastFetchedAt(cache.lastFetchedAt);
            setStatus("success");
          } catch {
            await clearCodexUsageCache();
          }
        }

        setAuth(storedAuth);
        if (storedAuth) await fetchUsage(storedAuth);
      }
    );

    return () => {
      active = false;
    };
  }, [fetchUsage]);

  const importAuthFile = useCallback(async () => {
    try {
      const parsed = await pickAndReadCodexAuth();
      await saveCodexAuth(parsed);
      setAuth(parsed);
      setError(null);
      await fetchUsage(parsed);
    } catch (caught: unknown) {
      const message = caught instanceof Error ? caught.message : "UNKNOWN_ERROR";
      if (message !== "PICKER_CANCELLED") setError(message);
    }
  }, [fetchUsage]);

  const refresh = useCallback(async () => {
    if (auth) await fetchUsage(auth);
  }, [auth, fetchUsage]);

  const reloadCredentials = useCallback(async () => {
    const stored = await loadCodexAuth();
    if (!stored) {
      requestVersion.current += 1;
      setAuth(null);
      setUsage(null);
      setLastFetchedAt(null);
      setStatus("idle");
      return;
    }
    setAuth(stored);
    await fetchUsage(stored);
  }, [fetchUsage]);

  const clearAuth = useCallback(async () => {
    requestVersion.current += 1;
    await clearCodexAuth();
    setAuth(null);
    setUsage(null);
    setLastFetchedAt(null);
    setStatus("idle");
    setError(null);
  }, []);

  return {
    auth,
    usage,
    lastFetchedAt,
    status,
    error,
    importAuthFile,
    refresh,
    reloadCredentials,
    clearAuth,
  };
}