import type { CountTokensParameters } from "@google/genai";

/**
 * Known token counting providers with provider-specific estimation ratios.
 */
export type TokenProvider =
	| "anthropic"
	| "openai"
	| "google"
	| "mistral"
	| "bedrock"
	| "unknown";

const CHARS_PER_TOKEN_BY_PROVIDER: Record<TokenProvider, number> = {
	anthropic: 3.5,
	openai: 4.0,
	google: 4.0,
	mistral: 3.8,
	bedrock: 3.5,
	unknown: 4.0,
};

/**
 * Interface for token encoding (typically tiktoken).
 */
interface TokenEncoder {
	encode(text: string): Uint32Array | number[];
}

let encoder: TokenEncoder | null = null;
let encoderFailed = false;

/**
 * Parsed credentials from Google Gemini CLI API key JSON.
 */
interface GeminiCliCredentials {
	token: string;
	projectId: string;
}

/**
 * Dependency injection interface for Google Gemini token counting.
 */
interface GeminiCountTokensDeps {
	buildServer(apiKeyRaw: string): Promise<{
		countTokens(
			request: CountTokensParameters,
		): Promise<{ totalTokens?: number | null }>;
	}>;
}

async function getEncoder(): Promise<TokenEncoder | null> {
	if (encoder) return encoder;
	if (encoderFailed) return null;
	try {
		// @ts-expect-error — tiktoken may not have type declarations in extensions tsconfig
		const tiktoken = await import("tiktoken");
		encoder = tiktoken.encoding_for_model("gpt-4o") as TokenEncoder;
		return encoder;
	} catch {
		encoderFailed = true;
		return null;
	}
}

/**
 * Count tokens in text using tiktoken if available, otherwise estimate.
 */
export async function countTokens(text: string): Promise<number> {
	const enc = await getEncoder();
	if (enc) {
		const tokens = enc.encode(text);
		return tokens.length;
	}
	return Math.ceil(text.length / 4);
}

/**
 * Synchronously count tokens (requires tiktoken to be pre-loaded).
 */
export function countTokensSync(text: string): number {
	if (encoder) {
		return encoder.encode(text).length;
	}
	return Math.ceil(text.length / 4);
}

/**
 * Initialize the token counter by loading tiktoken encoder.
 */
export async function initTokenCounter(): Promise<boolean> {
	const enc = await getEncoder();
	return enc !== null;
}

/**
 * Check if tiktoken encoder is loaded for accurate token counting.
 */
export function isAccurateCountingAvailable(): boolean {
	return encoder !== null;
}

/**
 * Get the provider-specific characters-per-token ratio for estimation.
 */
export function getCharsPerToken(provider: TokenProvider): number {
	return (
		CHARS_PER_TOKEN_BY_PROVIDER[provider] ?? CHARS_PER_TOKEN_BY_PROVIDER.unknown
	);
}

/**
 * Estimate token count for text using provider-specific ratio.
 */
export function estimateTokensForProvider(
	text: string,
	provider: TokenProvider,
): number {
	const ratio = getCharsPerToken(provider);
	return Math.ceil(text.length / ratio);
}

/**
 * Parse Google Gemini CLI API key JSON to extract token and project ID.
 */
export function parseGoogleGeminiCliApiKey(
	apiKeyRaw: string,
): GeminiCliCredentials | undefined {
	try {
		const parsed = JSON.parse(apiKeyRaw) as Partial<GeminiCliCredentials>;
		if (
			typeof parsed.token !== "string" ||
			typeof parsed.projectId !== "string"
		) {
			return undefined;
		}
		if (!parsed.token || !parsed.projectId) {
			return undefined;
		}
		return { token: parsed.token, projectId: parsed.projectId };
	} catch {
		return undefined;
	}
}

async function buildGoogleGeminiCliServer(apiKeyRaw: string) {
	const credentials = parseGoogleGeminiCliApiKey(apiKeyRaw);
	if (!credentials) {
		throw new Error("Invalid Google Gemini CLI credentials");
	}
	const [{ CodeAssistServer }, { OAuth2Client }] = await Promise.all([
		import("@google/gemini-cli-core"),
		import("google-auth-library"),
	]);
	const authClient = new OAuth2Client();
	authClient.setCredentials({ access_token: credentials.token });
	return new CodeAssistServer(
		authClient as unknown as ConstructorParameters<typeof CodeAssistServer>[0],
		credentials.projectId,
		{ headers: {} },
	);
}

/**
 * Type guard for Google Gemini token counting payload.
 */
export function isGoogleGeminiCountablePayload(
	payload: unknown,
): payload is CountTokensParameters {
	if (!payload || typeof payload !== "object") return false;
	const candidate = payload as Record<string, unknown>;
	return (
		typeof candidate.model === "string" && Array.isArray(candidate.contents)
	);
}

/**
 * Count tokens in a Google Gemini CLI request using their server API.
 */
export async function countGoogleGeminiCliTokens(
	payload: unknown,
	apiKeyRaw: string | undefined,
	deps: GeminiCountTokensDeps = { buildServer: buildGoogleGeminiCliServer },
): Promise<number | undefined> {
	if (!apiKeyRaw || !isGoogleGeminiCountablePayload(payload)) return undefined;
	const response = await (await deps.buildServer(apiKeyRaw)).countTokens(
		payload,
	);
	return typeof response.totalTokens === "number"
		? response.totalTokens
		: undefined;
}
