gemini-nexus/src/lib/antigravity/QuotaService.ts

import axios, { AxiosInstance } from 'axios';
import { QuotaData, LoadProjectResponse, QuotaApiResponse } from './types';
import { logger } from '../../utils/logger';

// Constants
const QUOTA_API_ENDPOINTS = [
  'https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:fetchAvailableModels',
  'https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels',
  'https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels',
] as const;
const CLOUD_CODE_BASE_URL = 'https://cloudcode-pa.googleapis.com';
const USER_AGENT = 'antigravity/1.11.3 Darwin/arm64'; // Keeping the same UA as source

// Service Class
export class QuotaService {
  private static createClient(timeoutSecs: number = 15): AxiosInstance {
    return axios.create({
      timeout: timeoutSecs * 1000,
      headers: {
        'User-Agent': USER_AGENT,
        'Content-Type': 'application/json',
      },
    });
  }

  /**
   * Fetch Project ID and Subscription Type
   */
  private static async fetchProjectId(
    accessToken: string,
    email: string,
  ): Promise<[string | undefined, string | undefined]> {
    const client = this.createClient();
    const meta = { metadata: { ideType: 'ANTIGRAVITY' } };

    try {
      const res = await client.post<LoadProjectResponse>(
        `${CLOUD_CODE_BASE_URL}/v1internal:loadCodeAssist`,
        meta,
        {
          headers: {
            Authorization: `Bearer ${accessToken}`,
            'User-Agent': 'antigravity/windows/amd64',
          },
        },
      );

      if (res.status >= 200 && res.status < 300) {
        const data = res.data;
        const projectId = data.cloudaicompanionProject;

        // Core logic: Preferentially get subscription ID from paid_tier
        const subscriptionTier = data.paidTier?.id || data.currentTier?.id;

        if (subscriptionTier) {
          logger.info(`📊 [${email}] Subscription Identified: ${subscriptionTier}`);
        }

        return [projectId, subscriptionTier];
      } else {
        logger.warn(`⚠️  [${email}] loadCodeAssist failed: Status: ${res.status}`);
      }
    } catch (error: any) {
      logger.error(`❌ [${email}] loadCodeAssist Network Error: ${error.message}`);
    }

    return [undefined, undefined];
  }

  /**
   * Unified entry point for querying account quota
   */
  public static async fetchQuota(accessToken: string, email: string) {
    return this.fetchQuotaInner(accessToken, email);
  }

  /**
   * Logic for querying account quota (Inner)
   */
  private static async fetchQuotaInner(
    accessToken: string,
    email: string,
  ): Promise<{ quotaData: QuotaData; projectId?: string }> {
    // 1. Get Project ID and Subscription Type
    const [projectId, subscriptionTier] = await this.fetchProjectId(accessToken, email);

    const finalProjectId = projectId;

    const client = this.createClient();
    const payload = finalProjectId ? { project: finalProjectId } : {};
    let lastError: Error | null = null;

    for (let endpointIndex = 0; endpointIndex < QUOTA_API_ENDPOINTS.length; endpointIndex++) {
      const endpoint = QUOTA_API_ENDPOINTS[endpointIndex];
      const hasNextEndpoint = endpointIndex + 1 < QUOTA_API_ENDPOINTS.length;
      logger.info(`Sending quota request to ${endpoint}`);

      try {
        const response = await client.post<QuotaApiResponse>(endpoint, payload, {
          headers: {
            Authorization: `Bearer ${accessToken}`,
            'User-Agent': USER_AGENT,
          },
        });

        const quotaResponse = response.data;
        const quotaData: QuotaData = {
          models: {},
          isForbidden: false,
          subscriptionTier: subscriptionTier,
        };

        logger.info(`Quota API returned ${Object.keys(quotaResponse.models || {}).length} models:`);

        if (quotaResponse.models) {
          for (const [name, info] of Object.entries(quotaResponse.models)) {
            logger.info(`   - ${name}`);
            if (info.quotaInfo) {
              const fraction = info.quotaInfo.remainingFraction ?? 0;
              const percentage = Math.floor(fraction * 100);
              const resetTime = info.quotaInfo.resetTime || '';

              // Only save models we care about, filtering out old versions (< 3.0)
              const isGemini = name.includes('gemini');
              const isClaude = name.includes('claude');
              const isOldGemini = /gemini-[12](\.|$|-)/.test(name);

              if ((isGemini || isClaude) && !isOldGemini) {
                quotaData.models[name] = { percentage, resetTime };
              }
            }
          }
        }

        if (endpointIndex > 0) {
          logger.info(`Quota API fallback succeeded at endpoint #${endpointIndex + 1}`);
        }

        return { quotaData, projectId };
      } catch (error: any) {
        let shouldFallback = true;

        if (axios.isAxiosError(error)) {
          const status = error.response?.status;
          let text = '';
          try {
            text = JSON.stringify(error.response?.data || '');
          } catch {
            text = '[Unable to serialize response data]';
          }

          // ✅ Handle 403 Forbidden specifically - return immediately, do not retry
          if (status === 403) {
            logger.warn(`Account no permission (403 Forbidden), marked as forbidden`);
            return {
              quotaData: {
                models: {},
                isForbidden: true,
                subscriptionTier: subscriptionTier,
              },
              projectId,
            };
          }

          if (hasNextEndpoint && (status === 429 || (typeof status === 'number' && status >= 500))) {
            logger.warn(
              `Quota API ${endpoint} returned ${status}, falling back to next endpoint`,
            );
            lastError = new Error(`HTTP ${status} - ${text}`);
            await new Promise((r) => setTimeout(r, 1000));
            continue;
          }

          logger.warn(`API Error: ${status} - ${text}`);
          lastError = new Error(`HTTP ${status} - ${text}`);
          shouldFallback = typeof status !== 'number';
        } else {
          logger.warn(`Request Failed at ${endpoint}: ${error.message}`);
          lastError = error instanceof Error ? error : new Error(String(error));
        }

        if (hasNextEndpoint && shouldFallback) {
          logger.warn(`Quota API request failed at ${endpoint}, falling back to next endpoint`);
          await new Promise((r) => setTimeout(r, 1000));
        } else {
          throw lastError ?? new Error(`Quota query failed: ${error.message}`);
        }
      }
    }

    throw lastError ?? new Error('Unknown error in fetchQuota');
  }
}