codexbar-mobile/hooks/useClaudeUsage.ts

import { useCallback, useEffect, useRef, useState } from "react";
import {
  fetchClaudeOrgs,
  fetchClaudeUsage,
  parseClaudeUsageResponse,
} from "@/lib/api/claudeApi";
import {
  isClaudeSessionKeyValid,
  validateClaudeSessionKey,
} from "@/lib/claudeCredentials";
import { RETRY_DELAY_MS } from "@/lib/constants";
import {
  clearClaudeLastActiveOrg,
  clearClaudeSessionKey,
  clearClaudeUsageCache,
  loadClaudeLastActiveOrg,
  loadClaudeSessionKey,
  loadClaudeUsageCache,
  saveClaudeLastActiveOrg,
  saveClaudeSessionKey,
  saveClaudeUsageCache,
} from "@/lib/storage";
import type { ClaudeAuth, ClaudeUsageResponse } from "@/types/claude";

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

function shouldRetry(error: unknown): boolean {
  const message = error instanceof Error ? error.message : "";
  return ![
    "TOKEN_EXPIRED",
    "NO_ORGS_FOUND",
    "INVALID_CLAUDE_RESPONSE",
    "INVALID_CLAUDE_ORGS_RESPONSE",
  ].includes(message);
}

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

async function requestClaudeUsage(auth: ClaudeAuth): Promise<ClaudeUsageResponse> {
  let orgUuid = auth.lastActiveOrg?.trim();
  if (!orgUuid) {
    const orgs = await fetchClaudeOrgs(auth);
    if (!orgs.length) throw new Error("NO_ORGS_FOUND");
    orgUuid = orgs[0].uuid;
    await saveClaudeLastActiveOrg(orgUuid);
  }
  return fetchClaudeUsage(auth, orgUuid);
}

export function useClaudeUsage() {
  const [auth, setAuth] = useState<ClaudeAuth | null>(null);
  const [pendingKey, setPendingKey] = useState("");
  const [pendingOrg, setPendingOrg] = useState("");
  const [usage, setUsage] = useState<ClaudeUsageResponse | 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: ClaudeAuth): Promise<void> => {
    const sanitized = {
      sessionKey: credentials.sessionKey.trim(),
      lastActiveOrg: credentials.lastActiveOrg?.trim() || undefined,
    };
    const requestKey = `${sanitized.lastActiveOrg ?? ""}:${sanitized.sessionKey}`;
    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: ClaudeUsageResponse;
        try {
          data = await requestClaudeUsage(sanitized);
        } catch (firstError) {
          if (!shouldRetry(firstError)) throw firstError;
          await wait(RETRY_DELAY_MS);
          if (version !== requestVersion.current) return;
          data = await requestClaudeUsage(sanitized);
        }

        if (version !== requestVersion.current) return;
        const fetchedAt = Date.now();
        await saveClaudeUsageCache({ 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 clearClaudeSessionKey();
          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([
      loadClaudeSessionKey(),
      loadClaudeLastActiveOrg(),
      loadClaudeUsageCache(),
    ]).then(async ([key, org, cache]) => {
      if (!active) return;

      const storedAuth = key
        ? { sessionKey: key.trim(), lastActiveOrg: org?.trim() || undefined }
        : null;

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

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

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

  const saveKey = useCallback(async () => {
    const trimmedKey = pendingKey.trim();
    const validationError = validateClaudeSessionKey(trimmedKey);
    if (validationError) {
      setError(validationError);
      return;
    }

    const trimmedOrg = pendingOrg.trim() || undefined;
    await saveClaudeSessionKey(trimmedKey);
    if (trimmedOrg) await saveClaudeLastActiveOrg(trimmedOrg);
    else await clearClaudeLastActiveOrg();

    const nextAuth = { sessionKey: trimmedKey, lastActiveOrg: trimmedOrg };
    setAuth(nextAuth);
    setPendingKey("");
    setPendingOrg("");
    setError(null);
    await fetchUsage(nextAuth);
  }, [fetchUsage, pendingKey, pendingOrg]);

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

  const reloadCredentials = useCallback(async () => {
    const [key, org] = await Promise.all([
      loadClaudeSessionKey(),
      loadClaudeLastActiveOrg(),
    ]);
    if (!key) {
      requestVersion.current += 1;
      setAuth(null);
      setUsage(null);
      setLastFetchedAt(null);
      setStatus("idle");
      return;
    }

    const stored = {
      sessionKey: key.trim(),
      lastActiveOrg: org?.trim() || undefined,
    };
    setAuth(stored);
    await fetchUsage(stored);
  }, [fetchUsage]);

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

  return {
    auth,
    pendingKey,
    setPendingKey,
    pendingOrg,
    setPendingOrg,
    usage,
    lastFetchedAt,
    status,
    error,
    keyValidationError: pendingKey
      ? validateClaudeSessionKey(pendingKey)
      : null,
    canSaveKey: isClaudeSessionKeyValid(pendingKey),
    saveKey,
    refresh,
    reloadCredentials,
    clearKey,
  };
}