Codewalkers/src/agent/accounts/usage.ts

import { readFileSync, existsSync, writeFileSync, mkdirSync } from 'node:fs';
import { join, dirname } from 'node:path';
import type { Account } from '../../db/schema.js';
import type { AgentInfo } from '../types.js';
import type { AccountCredentialManager } from '../credentials/types.js';
import { createModuleLogger } from '../../logger/index.js';
import { getAccountConfigDir } from './paths.js';
import { setupAccountConfigDir } from './setup.js';

const log = createModuleLogger('account-usage');

const USAGE_API_URL = 'https://api.anthropic.com/api/oauth/usage';
const TOKEN_REFRESH_URL = 'https://console.anthropic.com/v1/oauth/token';
const OAUTH_CLIENT_ID = '9d1c250a-e61b-44d9-88ed-5944d1962f5e';
const TOKEN_REFRESH_BUFFER_MS = 300_000; // 5 minutes

export interface OAuthCredentials {
  accessToken: string;
  refreshToken: string | null;
  expiresAt: number | null; // ms epoch, null for setup tokens
  subscriptionType: string | null;
  rateLimitTier: string | null;
}

export interface UsageTier {
  utilization: number;
  resets_at: string | null;
}

export interface AccountUsage {
  five_hour: UsageTier | null;
  seven_day: UsageTier | null;
  seven_day_sonnet: UsageTier | null;
  seven_day_opus: UsageTier | null;
  extra_usage: {
    is_enabled: boolean;
    monthly_limit: number | null;
    used_credits: number | null;
    utilization: number | null;
  } | null;
}

export interface AccountHealthResult {
  id: string;
  email: string;
  provider: string;
  credentialsValid: boolean;
  tokenValid: boolean;
  tokenExpiresAt: string | null;
  subscriptionType: string | null;
  error: string | null;
  usage: AccountUsage | null;
  isExhausted: boolean;
  exhaustedUntil: string | null;
  lastUsedAt: string | null;
  agentCount: number;
  activeAgentCount: number;
}

function readCredentials(configDir: string): OAuthCredentials | null {
  try {
    const credPath = join(configDir, '.credentials.json');
    if (!existsSync(credPath)) return null;
    const raw = readFileSync(credPath, 'utf-8');
    const parsed = JSON.parse(raw);
    const oauth = parsed.claudeAiOauth;
    if (!oauth || !oauth.accessToken) return null;
    return {
      accessToken: oauth.accessToken,
      refreshToken: oauth.refreshToken ?? null,
      expiresAt: oauth.expiresAt ?? null,
      subscriptionType: oauth.subscriptionType ?? null,
      rateLimitTier: oauth.rateLimitTier ?? null,
    };
  } catch {
    return null;
  }
}

function isTokenExpired(credentials: OAuthCredentials): boolean {
  if (!credentials.expiresAt) return false; // Setup tokens without expiry are treated as non-expired
  return credentials.expiresAt < Date.now() + TOKEN_REFRESH_BUFFER_MS;
}

/**
 * Write credentials back to the config directory.
 * Matches ccswitch's update_credentials_with_token() behavior.
 */
function writeCredentials(
  configDir: string,
  accessToken: string,
  refreshToken: string,
  expiresIn: number,
): void {
  const credPath = join(configDir, '.credentials.json');

  // Read existing credentials to preserve other fields
  let existing: Record<string, unknown> = {};
  try {
    if (existsSync(credPath)) {
      existing = JSON.parse(readFileSync(credPath, 'utf-8'));
    }
  } catch {
    // Start fresh if can't read
  }

  // Calculate expiry in milliseconds (matching ccswitch behavior)
  const nowMs = Date.now();
  const expiresAt = nowMs + (expiresIn * 1000);

  // Update claudeAiOauth section
  const claudeAiOauth = (existing.claudeAiOauth as Record<string, unknown>) ?? {};
  claudeAiOauth.accessToken = accessToken;
  claudeAiOauth.refreshToken = refreshToken;
  claudeAiOauth.expiresAt = expiresAt;
  existing.claudeAiOauth = claudeAiOauth;

  // Ensure directory exists
  mkdirSync(dirname(credPath), { recursive: true });

  // Write back (compact JSON for consistency with ccswitch)
  writeFileSync(credPath, JSON.stringify(existing));
  log.debug({ configDir }, 'credentials written after token refresh');
}

async function refreshToken(
  refreshTokenStr: string,
): Promise<{ accessToken: string; refreshToken: string; expiresIn: number } | null> {
  try {
    const response = await fetch(TOKEN_REFRESH_URL, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        grant_type: 'refresh_token',
        refresh_token: refreshTokenStr,
        client_id: OAUTH_CLIENT_ID,
        scope: 'user:inference user:profile',
      }),
    });
    if (!response.ok) return null;
    const data = await response.json();
    return {
      accessToken: data.access_token,
      refreshToken: data.refresh_token,
      expiresIn: data.expires_in,
    };
  } catch {
    return null;
  }
}

async function fetchUsage(accessToken: string): Promise<AccountUsage | null> {
  try {
    const response = await fetch(USAGE_API_URL, {
      method: 'GET',
      headers: {
        Authorization: `Bearer ${accessToken}`,
        'anthropic-beta': 'oauth-2025-04-20',
        'Content-Type': 'application/json',
      },
    });
    if (!response.ok) return null;
    const data = await response.json();
    return {
      five_hour: data.five_hour ?? null,
      seven_day: data.seven_day ?? null,
      seven_day_sonnet: data.seven_day_sonnet ?? null,
      seven_day_opus: data.seven_day_opus ?? null,
      extra_usage: data.extra_usage ?? null,
    };
  } catch {
    return null;
  }
}

export async function checkAccountHealth(
  account: Account,
  agents: AgentInfo[],
  credentialManager?: AccountCredentialManager,
  workspaceRoot?: string,
): Promise<AccountHealthResult> {
  const configDir = workspaceRoot ? getAccountConfigDir(workspaceRoot, account.id) : null;

  const accountAgents = agents.filter((a) => a.accountId === account.id);
  const activeAgents = accountAgents.filter(
    (a) => a.status === 'running' || a.status === 'waiting_for_input',
  );

  const base: AccountHealthResult = {
    id: account.id,
    email: account.email,
    provider: account.provider,
    credentialsValid: false,
    tokenValid: false,
    tokenExpiresAt: null,
    subscriptionType: null,
    error: null,
    usage: null,
    isExhausted: account.isExhausted,
    exhaustedUntil: account.exhaustedUntil?.toISOString() ?? null,
    lastUsedAt: account.lastUsedAt?.toISOString() ?? null,
    agentCount: accountAgents.length,
    activeAgentCount: activeAgents.length,
  };

  if (!configDir) {
    return { ...base, error: 'Cannot derive config dir: workspaceRoot not provided' };
  }

  // Ensure DB credentials are written to disk so file-based checks can find them
  if (account.configJson && account.credentials) {
    try {
      setupAccountConfigDir(configDir, {
        configJson: JSON.parse(account.configJson),
        credentials: account.credentials,
      });
    } catch (err) {
      log.warn({ accountId: account.id, err: err instanceof Error ? err.message : String(err) }, 'failed to sync DB credentials to disk');
    }
  }

  try {
    // Use credential manager if provided, otherwise fall back to direct functions
    let accessToken: string;
    let currentExpiresAt: number;
    let subscriptionType: string | null = null;

    if (credentialManager) {
      const result = await credentialManager.ensureValid(configDir, account.id);
      if (!result.valid || !result.credentials) {
        return {
          ...base,
          credentialsValid: result.credentials !== null,
          error: result.error ?? 'Credentials validation failed',
        };
      }
      accessToken = result.credentials.accessToken;
      currentExpiresAt = result.credentials.expiresAt;
      subscriptionType = result.credentials.subscriptionType;
    } else {
      // Legacy path: direct function calls
      const credentials = readCredentials(configDir);
      if (!credentials) {
        return {
          ...base,
          error: 'Credentials file not found or unreadable',
        };
      }

      accessToken = credentials.accessToken;
      currentExpiresAt = credentials.expiresAt;
      subscriptionType = credentials.subscriptionType;

      if (isTokenExpired(credentials)) {
        if (!credentials.refreshToken) {
          log.warn({ accountId: account.id }, 'setup token expired, no refresh token');
          return {
            ...base,
            credentialsValid: true,
            error: 'Setup token expired, no refresh token available',
          };
        }
        log.info({ accountId: account.id, email: account.email }, 'token expired, refreshing');
        const refreshed = await refreshToken(credentials.refreshToken);
        if (!refreshed) {
          log.warn({ accountId: account.id }, 'token refresh failed');
          return {
            ...base,
            credentialsValid: true,
            error: 'Token expired and refresh failed',
          };
        }
        accessToken = refreshed.accessToken;

        // Persist the refreshed credentials back to disk
        const newRefreshToken = refreshed.refreshToken || credentials.refreshToken;
        writeCredentials(configDir, accessToken, newRefreshToken, refreshed.expiresIn);
        currentExpiresAt = Date.now() + (refreshed.expiresIn * 1000);
        log.info({ accountId: account.id, expiresIn: refreshed.expiresIn }, 'token refreshed and persisted');
      }
    }

    const usage = await fetchUsage(accessToken);
    if (!usage) {
      return {
        ...base,
        credentialsValid: true,
        error: 'Usage API request failed',
      };
    }

    return {
      ...base,
      credentialsValid: true,
      tokenValid: true,
      tokenExpiresAt: new Date(currentExpiresAt).toISOString(),
      subscriptionType,
      usage,
    };
  } catch (err) {
    return {
      ...base,
      error: err instanceof Error ? err.message : String(err),
    };
  }
}

/**
 * Ensure account credentials are valid and refreshed if needed.
 * Call this before spawning an agent to ensure the credentials file
 * has fresh tokens that the agent subprocess can use.
 *
 * Returns true if credentials are valid (or were successfully refreshed).
 * Returns false if credentials are missing or refresh failed.
 *
 * @deprecated Use AccountCredentialManager.ensureValid() instead for event emission support.
 */
export async function ensureAccountCredentials(configDir: string): Promise<boolean> {
  const credentials = readCredentials(configDir);
  if (!credentials) {
    log.warn({ configDir }, 'no credentials found');
    return false;
  }

  if (!isTokenExpired(credentials)) {
    log.debug({ configDir }, 'credentials valid, no refresh needed');
    return true;
  }

  log.info({ configDir }, 'credentials expired, refreshing before spawn');
  const refreshed = await refreshToken(credentials.refreshToken);
  if (!refreshed) {
    log.error({ configDir }, 'failed to refresh credentials');
    return false;
  }

  const newRefreshToken = refreshed.refreshToken || credentials.refreshToken;
  writeCredentials(configDir, refreshed.accessToken, newRefreshToken, refreshed.expiresIn);
  log.info({ configDir, expiresIn: refreshed.expiresIn }, 'credentials refreshed before spawn');
  return true;
}