/**
 * Credential Pool — manages OAuth tokens and API keys for agent containers.
 *
 * Handles OAuth token refresh, LRU credential selection, and tenant pinning.
 * Credentials are stored in SQLite and selected per agent spawn.
 */
import fs from 'fs';
import os from 'os';
import path from 'path';
import { execFile } from 'child_process';
import { promisify } from 'util';

import {
  getActiveCredentials,
  getAllCredentials,
  getCredential,
  updateCredentialAfterUse,
  updateCredentialAuthJson,
  updateCredentialStatus,
  updateCredentialTokens,
} from './db.js';
import { logger } from './logger.js';
import {
  getCredentialProvider,
  getRuntimeCredentialEnvKeys,
  getTenantProvider,
  normalizeProvider,
  RuntimeProvider,
  getTenantPiBackend,
  getPiCredentialPoolProvider,
} from './runtime-provider.js';
import { Credential, Tenant } from './types.js';

const TOKEN_URL = 'https://platform.claude.com/v1/oauth/token';
const CLIENT_ID = '9d1c250a-e61b-44d9-88ed-5944d1962f5e';
const CODEX_BIN = fs.existsSync('/home/admin/.npm-global/bin/codex')
  ? '/home/admin/.npm-global/bin/codex'
  : 'codex';
const execFileAsync = promisify(execFile);
// Refresh when less than 10 minutes remain
const REFRESH_BUFFER_MS = 10 * 60 * 1000;

function decodeJwtExp(token?: string): number | undefined {
  if (!token) return undefined;
  const parts = token.split('.');
  if (parts.length < 2) return undefined;
  try {
    const payload = JSON.parse(
      Buffer.from(
        parts[1].replace(/-/g, '+').replace(/_/g, '/'),
        'base64',
      ).toString('utf8'),
    ) as { exp?: number };
    return typeof payload.exp === 'number' ? payload.exp * 1000 : undefined;
  } catch {
    return undefined;
  }
}

function getCodexCredentialExpiryMs(cred: Credential): number | undefined {
  if (cred.expires_at) return cred.expires_at;
  if (!cred.auth_json) return undefined;
  try {
    const parsed = JSON.parse(cred.auth_json) as {
      tokens?: { access_token?: string };
    };
    return decodeJwtExp(parsed.tokens?.access_token);
  } catch {
    return undefined;
  }
}

function parseClaudeCredentialsAuthJson(
  authJson: string,
): {
  accessToken?: string;
  refreshToken?: string;
  expiresAt?: number;
  subscriptionType?: string;
  rateLimitTier?: string;
} {
  try {
    const parsed = JSON.parse(authJson) as {
      claudeAiOauth?: {
        accessToken?: string;
        refreshToken?: string;
        expiresAt?: number;
        subscriptionType?: string;
        rateLimitTier?: string;
      };
    };
    const oauth = parsed?.claudeAiOauth;
    if (!oauth || typeof oauth !== 'object') return {};
    return {
      accessToken:
        typeof oauth.accessToken === 'string' ? oauth.accessToken : undefined,
      refreshToken:
        typeof oauth.refreshToken === 'string' ? oauth.refreshToken : undefined,
      expiresAt:
        typeof oauth.expiresAt === 'number' ? oauth.expiresAt : undefined,
      subscriptionType:
        typeof oauth.subscriptionType === 'string'
          ? oauth.subscriptionType
          : undefined,
      rateLimitTier:
        typeof oauth.rateLimitTier === 'string' ? oauth.rateLimitTier : undefined,
    };
  } catch {
    return {};
  }
}

function getClaudeCredentialExpiryMs(cred: Credential): number | undefined {
  if (cred.expires_at) return cred.expires_at;
  if (!cred.auth_json) return undefined;
  return parseClaudeCredentialsAuthJson(cred.auth_json).expiresAt;
}

async function refreshCodexCredential(
  cred: Credential,
): Promise<Credential | undefined> {
  if (!cred.auth_json) {
    updateCredentialStatus(cred.id, 'error', 'No Codex auth bundle');
    return undefined;
  }

  const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'codex-refresh-'));
  const codexHome = path.join(tmpDir, '.codex');
  const workDir = path.join(tmpDir, 'work');
  try {
    fs.mkdirSync(codexHome, { recursive: true });
    fs.mkdirSync(workDir, { recursive: true });
    fs.writeFileSync(path.join(codexHome, 'auth.json'), cred.auth_json, {
      mode: 0o600,
    });
    fs.writeFileSync(path.join(codexHome, 'config.toml'), 'model = "gpt-5.2"\n', {
      mode: 0o600,
    });
    fs.writeFileSync(path.join(workDir, '.gitkeep'), '');

    await execFileAsync('git', ['init', '-q'], { cwd: workDir, timeout: 10000 });
    await execFileAsync(
      CODEX_BIN,
      ['exec', '--skip-git-repo-check', 'Say only: auth refresh ok'],
      {
        cwd: workDir,
        env: {
          ...process.env,
          HOME: tmpDir,
        },
        timeout: 45000,
        maxBuffer: 2 * 1024 * 1024,
      },
    );

    const refreshedAuthJson = fs.readFileSync(
      path.join(codexHome, 'auth.json'),
      'utf8',
    );
    const parsed = JSON.parse(refreshedAuthJson) as {
      tokens?: { access_token?: string };
    };
    const expiresAt = decodeJwtExp(parsed.tokens?.access_token);
    updateCredentialAuthJson(cred.id, refreshedAuthJson, expiresAt);
    return {
      ...cred,
      auth_json: refreshedAuthJson,
      expires_at: expiresAt,
      status: 'active',
      error_message: undefined,
    };
  } catch (err) {
    logger.error({ credId: cred.id, err }, 'Codex credential refresh error');
    updateCredentialStatus(cred.id, 'error', `Refresh error: ${err}`);
    return undefined;
  } finally {
    fs.rmSync(tmpDir, { recursive: true, force: true });
  }
}

/**
 * Refresh an OAuth credential's access token using its refresh token.
 * Updates the DB row with new tokens on success, marks expired on failure.
 * Returns the updated credential or undefined on failure.
 */
export async function refreshOAuthCredential(
  cred: Credential,
): Promise<Credential | undefined> {
  let refreshToken = cred.refresh_token;
  if (!refreshToken && cred.auth_json) {
    refreshToken = parseClaudeCredentialsAuthJson(cred.auth_json).refreshToken;
  }
  if (!refreshToken) {
    logger.error({ credId: cred.id }, 'OAuth credential has no refresh token');
    updateCredentialStatus(cred.id, 'error', 'No refresh token');
    return undefined;
  }

  try {
    const res = await fetch(TOKEN_URL, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        grant_type: 'refresh_token',
        refresh_token: refreshToken,
        client_id: CLIENT_ID,
      }),
      signal: AbortSignal.timeout(15000),
    });

    if (!res.ok) {
      const body = await res.text();
      logger.error(
        { credId: cred.id, status: res.status, body },
        'OAuth token refresh failed',
      );
      if (res.status === 401 || res.status === 403) {
        updateCredentialStatus(cred.id, 'expired', `Refresh failed: ${res.status}`);
      } else {
        updateCredentialStatus(cred.id, 'error', `Refresh failed: ${res.status} ${body.substring(0, 200)}`);
      }
      return undefined;
    }

    const data = (await res.json()) as {
      access_token: string;
      refresh_token?: string;
      expires_in: number;
    };

    const expiresAt = Date.now() + data.expires_in * 1000;
    const newRefreshToken = data.refresh_token ?? refreshToken;

    updateCredentialTokens(cred.id, data.access_token, newRefreshToken, expiresAt);

    let updatedAuthJson = cred.auth_json;
    if (cred.auth_json) {
      try {
        const parsed = JSON.parse(cred.auth_json) as Record<string, unknown>;
        const oauth = (parsed as any).claudeAiOauth;
        if (oauth && typeof oauth === 'object') {
          (oauth as any).accessToken = data.access_token;
          (oauth as any).refreshToken = newRefreshToken;
          (oauth as any).expiresAt = expiresAt;
          updatedAuthJson = JSON.stringify(parsed);
          updateCredentialAuthJson(cred.id, updatedAuthJson, expiresAt);
        }
      } catch {
        // ignore
      }
    }

    logger.info(
      { credId: cred.id, expiresIn: data.expires_in },
      'OAuth credential refreshed',
    );

    return {
      ...cred,
      access_token: data.access_token,
      refresh_token: newRefreshToken,
      auth_json: updatedAuthJson,
      expires_at: expiresAt,
      status: 'active',
      error_message: undefined,
    };
  } catch (err) {
    logger.error(
      { credId: cred.id, err },
      'OAuth token refresh error',
    );
    updateCredentialStatus(cred.id, 'error', `Refresh error: ${err}`);
    return undefined;
  }
}


// Path to pi-ai oauth module (same node_modules as the agent-runner)
const PI_AI_OAUTH_PATH = '/home/admin/treehouse/nanoclaw/container/agent-runner/node_modules/@mariozechner/pi-ai/dist/utils/oauth/index.js';

/**
 * Refresh an OAuth pi credential by calling pi-ai's refreshToken() for the
 * specific backend. This is the ONLY place token refresh for pi happens —
 * containers receive fresh tokens via PI_AUTH_JSON injection (never self-refresh).
 */
export async function refreshPiCredential(
  cred: Credential,
): Promise<Credential | undefined> {
  if (!cred.auth_json || !cred.pi_backend) {
    logger.warn({ credId: cred.id }, 'Pi credential missing auth_json or pi_backend');
    return undefined;
  }
  try {
    const piAuthData = JSON.parse(cred.auth_json) as Record<string, unknown>;
    const credSection = piAuthData[cred.pi_backend];
    if (!credSection || typeof credSection !== 'object') {
      logger.warn({ credId: cred.id, piBackend: cred.pi_backend }, 'No credential section found in pi auth_json');
      return undefined;
    }

    // Dynamically import pi-ai OAuth module (avoids adding it as a compile-time dep)
    const piOauth = await import(PI_AI_OAUTH_PATH) as { getOAuthProvider: (id: string) => any };
    const provider = piOauth.getOAuthProvider(cred.pi_backend);
    if (!provider?.refreshToken) {
      logger.warn({ piBackend: cred.pi_backend }, 'No OAuth provider or refreshToken method found for pi backend');
      return undefined;
    }

    const refreshed = await provider.refreshToken(credSection);
    const newSection = { type: 'oauth', ...refreshed };
    const newAuthJson = JSON.stringify({ ...piAuthData, [cred.pi_backend]: newSection });
    const expiresAt = typeof refreshed.expires === 'number' ? refreshed.expires : undefined;

    updateCredentialAuthJson(cred.id, newAuthJson, expiresAt ?? null);
    logger.info(
      { credId: cred.id, piBackend: cred.pi_backend, expiresIn: expiresAt ? Math.round((expiresAt - Date.now()) / 1000) : 'unknown' },
      'Pi OAuth credential refreshed',
    );

    return { ...cred, auth_json: newAuthJson, expires_at: expiresAt ?? null, status: 'active', error_message: undefined };
  } catch (err) {
    logger.error({ credId: cred.id, piBackend: (cred as any).pi_backend, err }, 'Pi OAuth token refresh error');
    updateCredentialStatus(cred.id, 'error', `Pi refresh error: ${err}`);
    return undefined;
  }
}

// Round-robin counter per provider: tracks the next index to use
const rrCounters = new Map<string, number>();

/**
 * Select a credential for a tenant's agent spawn.
 *
 * Priority:
 * 1. Tenant has own anthropic_key → use that (no pool credential needed)
 * 2. Tenant is pinned to a credential_id → use that specific credential
 * 3. Otherwise → pick from pool using round-robin (cycles through credentials in stable order)
 *
 * For OAuth credentials, refreshes the token if it's close to expiry.
 * Returns the secrets dict to pass to the container, or undefined if no credential available.
 */
export function selectCredentialForTenant(
  tenant?: Tenant,
): { secrets: Record<string, string>; credentialId?: string } | undefined {
  const provider = getTenantProvider(tenant);
  // pi uses Anthropic credentials -- fall back to claude pool for lookups
  const piBackend = provider === 'pi' ? getTenantPiBackend(tenant) : 'anthropic';
  const credPoolProvider: RuntimeProvider = provider === 'pi' ? getPiCredentialPoolProvider(piBackend) : provider;

  // Tenant has own API key — no pool credential needed
  if ((provider === 'claude' || (provider === 'pi' && piBackend === 'anthropic')) && tenant?.anthropic_key) {
    return {
      secrets: { ANTHROPIC_API_KEY: tenant.anthropic_key },
    };
  }

  let cred: Credential | undefined;

  // Tenant pinned to specific credential
  if (tenant?.credential_id) {
    cred = getCredential(tenant.credential_id);
    if (!cred) {
      logger.error(
        { tenantId: tenant.id, credentialId: tenant.credential_id },
        'Pinned credential not found',
      );
      return undefined;
    }
    if (cred.status !== 'active') {
      logger.warn(
        { tenantId: tenant.id, credentialId: cred.id, status: cred.status },
        'Pinned credential is not active',
      );
      return undefined;
    }
    if (getCredentialProvider(cred) !== provider) {
      logger.error(
        { tenantId: tenant.id, credentialId: cred.id, provider },
        'Pinned credential provider does not match tenant provider',
      );
      return undefined;
    }
  } else {
    // Round-robin selection from pool
    const active = getActiveCredentials()
      .filter((entry) => {
        if (getCredentialProvider(entry) !== credPoolProvider) return false;
        // For pi credentials, also match on pi_backend
        if (credPoolProvider === 'pi') {
          const credPiBackend = (entry as any).pi_backend;
          return !credPiBackend || credPiBackend === piBackend;
        }
        return true;
      })
      .sort((a, b) => a.id.localeCompare(b.id)); // stable order by id
    if (active.length === 0) {
      const allActive = getActiveCredentials();
      logger.warn({ provider, credPoolProvider, piBackend, totalActive: allActive.length, providers: allActive.map(c => `${c.id}:${c.provider}`) }, 'No active credentials in pool');
      return undefined;
    }
    const counter = rrCounters.get(credPoolProvider) ?? 0;
    cred = active[counter % active.length];
    rrCounters.set(credPoolProvider, counter + 1);
  }

  // Build secrets dict
  const secrets: Record<string, string> = {};
  if (credPoolProvider === 'pi' && cred.auth_json) {
    // pi credential: inject PI_AUTH_JSON for the agent-runner to write to ~/.pi/agent/auth.json
    secrets['PI_AUTH_JSON'] = cred.auth_json;
  } else if (credPoolProvider === 'codex' && cred.type === 'oauth' && cred.auth_json) {
    secrets['CODEX_AUTH_JSON'] = cred.auth_json;
    // For pi+openai: also extract the access_token so the pi agent-runner
    // can inject it via authStorage.setRuntimeApiKey('openai', ...)
    if (provider === 'pi') {
      try {
        const parsed = JSON.parse(cred.auth_json);
        if (parsed.tokens?.access_token) {
          secrets['OPENAI_API_KEY'] = parsed.tokens.access_token;
        }
      } catch { /* ignore parse errors */ }
    }
  } else if (credPoolProvider === 'codex' && cred.type === 'apikey' && cred.api_key) {
    secrets['OPENAI_API_KEY'] = cred.api_key;
  } else if ((provider === 'claude' || provider === 'pi') && cred.type === 'oauth' && cred.auth_json) {
    const parsed = parseClaudeCredentialsAuthJson(cred.auth_json);
    if (parsed.accessToken) {
      secrets['CLAUDE_CODE_OAUTH_TOKEN'] = parsed.accessToken;
    } else if (cred.access_token) {
      secrets['CLAUDE_CODE_OAUTH_TOKEN'] = cred.access_token;
    } else {
      logger.error({ credId: cred.id }, 'Claude auth_json missing accessToken');
      return undefined;
    }
  } else if (cred.type === 'oauth' && cred.access_token) {
    secrets['CLAUDE_CODE_OAUTH_TOKEN'] = cred.access_token;
  } else if (cred.type === 'apikey' && cred.api_key) {
    secrets['ANTHROPIC_API_KEY'] = cred.api_key;
  } else {
    logger.error({ credId: cred.id, type: cred.type }, 'Credential has no usable token');
    return undefined;
  }

  // Mark as used
  updateCredentialAfterUse(cred.id);

  return { secrets, credentialId: cred.id };
}

// --- Token file management for long-running containers ---

const TOKEN_FILE_NAME = '.token.json';

/**
 * Write a token file to a group's IPC directory so long-running containers
 * can pick up refreshed tokens without restarting.
 * Uses atomic write (write to temp, rename) to avoid partial reads.
 */
export function writeTokenFile(
  ipcDir: string,
  secrets: Record<string, string>,
  credentialId?: string,
): void {
  const tokenPath = path.join(ipcDir, TOKEN_FILE_NAME);
  const data = JSON.stringify({
    ...secrets,
    _credentialId: credentialId,
    _updatedAt: Date.now(),
  });
  const tmpPath = tokenPath + `.${process.pid}.tmp`;
  try {
    fs.writeFileSync(tmpPath, data, { mode: 0o600 });
    fs.renameSync(tmpPath, tokenPath);
  } catch (err) {
    logger.warn({ ipcDir, err }, 'Failed to write token file');
    try { fs.unlinkSync(tmpPath); } catch { /* ignore */ }
  }
}

/**
 * Remove a token file from a group's IPC directory (cleanup on container exit).
 */
export function removeTokenFile(ipcDir: string): void {
  try {
    fs.unlinkSync(path.join(ipcDir, TOKEN_FILE_NAME));
  } catch { /* ignore if not exists */ }
}

// Track which credential each active container group uses
// Map<groupFolder, credentialId>
const activeContainerCredentials = new Map<string, string>();

export function trackContainerCredential(groupFolder: string, credentialId: string): void {
  activeContainerCredentials.set(groupFolder, credentialId);
}

export function untrackContainerCredential(groupFolder: string): void {
  activeContainerCredentials.delete(groupFolder);
}

export function getActiveContainerCredentials(): Map<string, string> {
  return activeContainerCredentials;
}

export function syncCodexCredentialFromFile(
  credentialId: string,
  authFilePath: string,
): void {
  try {
    if (!fs.existsSync(authFilePath)) return;
    const authJson = fs.readFileSync(authFilePath, 'utf8');
    const parsed = JSON.parse(authJson) as { tokens?: { access_token?: string } };
    const expiresAt = decodeJwtExp(parsed.tokens?.access_token);
    updateCredentialAuthJson(credentialId, authJson, expiresAt);
  } catch (err) {
    logger.warn({ credentialId, authFilePath, err }, 'Failed to sync Codex auth bundle from runtime file');
  }
}

// Refresh interval: check every 5 minutes
const PROACTIVE_REFRESH_INTERVAL_MS = 5 * 60 * 1000;
// Max random jitter between individual credential refreshes (2 seconds)
const REFRESH_JITTER_MAX_MS = 2000;

let refreshTimer: ReturnType<typeof setInterval> | undefined;

/**
 * Start a proactive token refresh timer.
 * Periodically checks all OAuth credentials and refreshes those expiring soon.
 * After refresh, updates token files for all active containers using that credential.
 *
 * @param resolveIpcDir - function to resolve a group folder to its IPC directory path
 * @param getSecrets - function to build the full secrets dict for a credential
 */
export function startTokenRefreshTimer(
  resolveIpcDir: (groupFolder: string) => string,
): void {
  if (refreshTimer) return; // already running

  refreshTimer = setInterval(async () => {
    try {
      const allCreds = getAllCredentials();
      const now = Date.now();

      for (const cred of allCreds) {
        const provider = normalizeProvider(cred.provider);
        if (cred.type !== 'oauth' || cred.status !== 'active') continue;
        // For pi: check expiry from auth_json[pi_backend].expires, refresh via pi-ai
        if (provider === 'pi') {
          const piBackendKey = (cred as any).pi_backend as string | undefined;
          if (!piBackendKey || !cred.auth_json) continue;
          let piExpires = 0;
          try {
            const piAuthData = JSON.parse(cred.auth_json) as Record<string, any>;
            piExpires = typeof piAuthData[piBackendKey]?.expires === 'number' ? piAuthData[piBackendKey].expires : 0;
          } catch { continue; }
          if (piExpires - now >= REFRESH_BUFFER_MS) continue; // not expiring soon
          logger.info(
            { credId: cred.id, piBackend: piBackendKey, expiresIn: Math.round((piExpires - now) / 1000) },
            'Proactive token refresh: pi credential expiring soon',
          );
          await new Promise(resolve => setTimeout(resolve, Math.random() * REFRESH_JITTER_MAX_MS));
          const refreshed = await refreshPiCredential(cred);
          if (!refreshed) continue;
          for (const [groupFolder, credId] of activeContainerCredentials) {
            if (credId !== cred.id) continue;
            try {
              const ipcDir = resolveIpcDir(groupFolder);
              writeTokenFile(ipcDir, { PI_AUTH_JSON: refreshed.auth_json! }, refreshed.id);
              logger.debug({ groupFolder, credId: refreshed.id }, 'Updated PI_AUTH_JSON token file for active container');
            } catch (err) {
              logger.warn({ groupFolder, err }, 'Failed to update PI_AUTH_JSON token file for active container');
            }
          }
          continue;
        }

        const expiresAt = provider === 'codex'
          ? (getCodexCredentialExpiryMs(cred) ?? 0)
          : provider === 'claude'
            ? (getClaudeCredentialExpiryMs(cred) ?? (cred.expires_at ?? 0))
            : (cred.expires_at ?? 0);
        if (expiresAt - now >= REFRESH_BUFFER_MS) continue; // not expiring soon

        logger.info(
          { credId: cred.id, provider, expiresIn: Math.round((expiresAt - now) / 1000) },
          'Proactive token refresh: credential expiring soon',
        );

        await new Promise(resolve => setTimeout(resolve, Math.random() * REFRESH_JITTER_MAX_MS));

        const refreshed = provider === 'codex'
          ? await refreshCodexCredential(cred)
          : await refreshOAuthCredential(cred);
        if (!refreshed) continue;

        // Update token files for all active containers using this credential
        for (const [groupFolder, credId] of activeContainerCredentials) {
          if (credId !== cred.id) continue;
          try {
            const ipcDir = resolveIpcDir(groupFolder);
            const secrets: Record<string, string> = provider === 'codex'
              ? { CODEX_AUTH_JSON: refreshed.auth_json! }
              : { CLAUDE_CODE_OAUTH_TOKEN: refreshed.access_token! };
            writeTokenFile(ipcDir, secrets, refreshed.id);
            logger.debug(
              { groupFolder, credId: refreshed.id },
              'Updated token file for active container',
            );
          } catch (err) {
            logger.warn(
              { groupFolder, err },
              'Failed to update token file for active container',
            );
          }
        }
      }
    } catch (err) {
      logger.error({ err }, 'Proactive token refresh timer error');
    }
  }, PROACTIVE_REFRESH_INTERVAL_MS);

  // Don't prevent process exit
  refreshTimer.unref();
  logger.info('Started proactive token refresh timer (5 min interval)');
}

export function stopTokenRefreshTimer(): void {
  if (refreshTimer) {
    clearInterval(refreshTimer);
    refreshTimer = undefined;
  }
}


/**
 * Ensure the credential for a tenant is fresh before launching a task container.
 * If the credential is expired or expiring within REFRESH_BUFFER_MS, refreshes it now.
 */
export async function ensureCredentialFreshForTenant(tenant: Tenant | undefined): Promise<void> {
  if (!tenant) return;
  const provider = getTenantProvider(tenant);
  if (provider !== 'claude') return;
  if (tenant.anthropic_key) return;

  let cred: Credential | undefined;
  if (tenant.credential_id) {
    cred = getCredential(tenant.credential_id);
  } else {
    const active = getActiveCredentials().filter(
      (c) => getCredentialProvider(c) === 'claude' && c.type === 'oauth',
    );
    if (active.length === 0) return;
    cred = active.sort((a, b) => a.id.localeCompare(b.id))[0];
  }

  if (!cred || cred.type !== 'oauth' || cred.status !== 'active') return;

  const expiresAt = getClaudeCredentialExpiryMs(cred) ?? (cred.expires_at ?? 0);
  const now = Date.now();
  if (expiresAt - now >= REFRESH_BUFFER_MS) return;

  logger.info(
    { credId: cred.id, tenantId: tenant.id, expiresIn: Math.round((expiresAt - now) / 1000) },
    'Pre-task credential refresh: token expired or expiring soon',
  );
  await refreshOAuthCredential(cred);
}
