/**
 * Test script for GitLab Duo - full streaming flow
 * Run: npx tsx test.ts [model-id]
 *
 * Examples:
 *   npx tsx test.ts                      # Test default (duo-chat-sonnet-4-5)
 *   npx tsx test.ts duo-chat-gpt-5-codex # Test GPT-5 Codex (Responses API)
 *   npx tsx test.ts duo-chat-gpt-5-1     # Test GPT-5.1 (Chat Completions API)
 */

import {
	type Api,
	type AssistantMessageEventStream,
	type Context,
	createAssistantMessageEventStream,
	type Model,
	registerApiProvider,
	type SimpleStreamOptions,
	streamSimple,
} from "@mariozechner/pi-ai";
import { readFileSync } from "fs";
import { homedir } from "os";
import { join } from "path";

// =============================================================================
// Constants (copied from index.ts)
// =============================================================================

const GITLAB_COM_URL = "https://gitlab.com";
const AI_GATEWAY_URL = "https://cloud.gitlab.com";
const ANTHROPIC_PROXY_URL = `${AI_GATEWAY_URL}/ai/v1/proxy/anthropic/`;
const OPENAI_PROXY_URL = `${AI_GATEWAY_URL}/ai/v1/proxy/openai/v1`;
const DIRECT_ACCESS_TTL = 25 * 60 * 1000;

type OpenAIApi = "openai-completions" | "openai-responses";
const MODEL_MAPPINGS: Record<string, { api: "anthropic-messages" | OpenAIApi; backendModel: string; baseUrl: string }> =
	{
		"duo-chat-opus-4-5": {
			api: "anthropic-messages",
			backendModel: "claude-opus-4-5-20251101",
			baseUrl: ANTHROPIC_PROXY_URL,
		},
		"duo-chat-sonnet-4-5": {
			api: "anthropic-messages",
			backendModel: "claude-sonnet-4-5-20250929",
			baseUrl: ANTHROPIC_PROXY_URL,
		},
		"duo-chat-haiku-4-5": {
			api: "anthropic-messages",
			backendModel: "claude-haiku-4-5-20251001",
			baseUrl: ANTHROPIC_PROXY_URL,
		},
		// All GPT models use Responses API for consistent tool call ID format across model switches
		"duo-chat-gpt-5-1": { api: "openai-responses", backendModel: "gpt-5.1-2025-11-13", baseUrl: OPENAI_PROXY_URL },
		"duo-chat-gpt-5-mini": {
			api: "openai-responses",
			backendModel: "gpt-5-mini-2025-08-07",
			baseUrl: OPENAI_PROXY_URL,
		},
		"duo-chat-gpt-5-codex": { api: "openai-responses", backendModel: "gpt-5-codex", baseUrl: OPENAI_PROXY_URL },
	};

// Model definitions for cost tracking
const MODEL_DEFS: Record<
	string,
	{
		name: string;
		cost: { input: number; output: number; cacheRead: number; cacheWrite: number };
		contextWindow: number;
		maxTokens: number;
	}
> = {
	"duo-chat-opus-4-5": {
		name: "GitLab Duo Claude Opus 4.5",
		cost: { input: 15, output: 75, cacheRead: 1.5, cacheWrite: 18.75 },
		contextWindow: 200000,
		maxTokens: 32000,
	},
	"duo-chat-sonnet-4-5": {
		name: "GitLab Duo Claude Sonnet 4.5",
		cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
		contextWindow: 200000,
		maxTokens: 16384,
	},
	"duo-chat-haiku-4-5": {
		name: "GitLab Duo Claude Haiku 4.5",
		cost: { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 1.25 },
		contextWindow: 200000,
		maxTokens: 8192,
	},
	"duo-chat-gpt-5-1": {
		name: "GitLab Duo GPT-5.1",
		cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
		contextWindow: 128000,
		maxTokens: 16384,
	},
	"duo-chat-gpt-5-mini": {
		name: "GitLab Duo GPT-5 Mini",
		cost: { input: 0.15, output: 0.6, cacheRead: 0, cacheWrite: 0 },
		contextWindow: 128000,
		maxTokens: 16384,
	},
	"duo-chat-gpt-5-codex": {
		name: "GitLab Duo GPT-5 Codex",
		cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
		contextWindow: 128000,
		maxTokens: 16384,
	},
};

// =============================================================================
// Direct Access Token
// =============================================================================

interface DirectAccessToken {
	token: string;
	headers: Record<string, string>;
	expiresAt: number;
}

let cachedDirectAccess: DirectAccessToken | null = null;

async function getDirectAccessToken(gitlabAccessToken: string): Promise<DirectAccessToken> {
	const now = Date.now();
	if (cachedDirectAccess && cachedDirectAccess.expiresAt > now) {
		return cachedDirectAccess;
	}

	const url = `${GITLAB_COM_URL}/api/v4/ai/third_party_agents/direct_access`;
	console.log("Fetching direct access token from:", url);

	const response = await fetch(url, {
		method: "POST",
		headers: {
			Authorization: `Bearer ${gitlabAccessToken}`,
			"Content-Type": "application/json",
		},
		body: JSON.stringify({ feature_flags: { DuoAgentPlatformNext: true } }),
	});

	if (!response.ok) {
		const errorText = await response.text();
		throw new Error(`Failed to get direct access token: ${response.status} ${errorText}`);
	}

	const data = (await response.json()) as { token: string; headers: Record<string, string> };
	console.log("Got direct access token");
	cachedDirectAccess = {
		token: data.token,
		headers: data.headers,
		expiresAt: now + DIRECT_ACCESS_TTL,
	};
	return cachedDirectAccess;
}

// =============================================================================
// Stream Function (copied from index.ts)
// =============================================================================

function streamGitLabDuo(
	model: Model<Api>,
	context: Context,
	options?: SimpleStreamOptions,
): AssistantMessageEventStream {
	const stream = createAssistantMessageEventStream();

	(async () => {
		try {
			const gitlabAccessToken = options?.apiKey;
			console.log("streamGitLabDuo called");
			console.log("  model.id:", model.id);
			console.log("  options.apiKey present:", !!gitlabAccessToken);

			if (!gitlabAccessToken) {
				throw new Error("No GitLab access token provided in options.apiKey");
			}

			const mapping = MODEL_MAPPINGS[model.id];
			if (!mapping) throw new Error(`Unknown model: ${model.id}`);

			// Get direct access token (cached)
			const directAccess = await getDirectAccessToken(gitlabAccessToken);

			// Create a proxy model that uses the backend API
			const proxyModel: Model<typeof mapping.api> = {
				...model,
				id: mapping.backendModel,
				api: mapping.api,
				baseUrl: mapping.baseUrl,
			};

			// Merge GitLab headers with Authorization bearer token
			const headers = {
				...directAccess.headers,
				Authorization: `Bearer ${directAccess.token}`,
			};

			console.log("Calling streamSimple with proxy model:");
			console.log("  proxyModel.id:", proxyModel.id);
			console.log("  proxyModel.api:", proxyModel.api);
			console.log("  proxyModel.baseUrl:", proxyModel.baseUrl);

			// Delegate to pi-ai's built-in streaming
			const innerStream = streamSimple(proxyModel, context, {
				...options,
				apiKey: "gitlab-duo", // Dummy value to pass validation
				headers,
			});

			// Forward all events
			for await (const event of innerStream) {
				// Patch the model info back to gitlab-duo
				if ("partial" in event && event.partial) {
					event.partial.api = model.api;
					event.partial.provider = model.provider;
					event.partial.model = model.id;
				}
				if ("message" in event && event.message) {
					event.message.api = model.api;
					event.message.provider = model.provider;
					event.message.model = model.id;
				}
				if ("error" in event && event.error) {
					event.error.api = model.api;
					event.error.provider = model.provider;
					event.error.model = model.id;
				}
				stream.push(event);
			}
			stream.end();
		} catch (error) {
			console.error("Stream error:", error);
			stream.push({
				type: "error",
				reason: "error",
				error: {
					role: "assistant",
					content: [],
					api: model.api,
					provider: model.provider,
					model: model.id,
					usage: {
						input: 0,
						output: 0,
						cacheRead: 0,
						cacheWrite: 0,
						totalTokens: 0,
						cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
					},
					stopReason: "error",
					errorMessage: error instanceof Error ? error.message : String(error),
					timestamp: Date.now(),
				},
			});
			stream.end();
		}
	})();

	return stream;
}

// =============================================================================
// Main Test
// =============================================================================

interface AuthData {
	[provider: string]: {
		type: "oauth" | "api_key";
		refresh?: string;
		access?: string;
		expires?: number;
		key?: string;
	};
}

async function main() {
	const modelId = process.argv[2] || "duo-chat-sonnet-4-5";

	if (!MODEL_MAPPINGS[modelId]) {
		console.error(`Unknown model: ${modelId}`);
		console.error("Available models:", Object.keys(MODEL_MAPPINGS).join(", "));
		process.exit(1);
	}

	// Read auth.json
	const authPath = join(homedir(), ".pi", "agent", "auth.json");
	console.log("Reading auth from:", authPath);

	let authData: AuthData;
	try {
		authData = JSON.parse(readFileSync(authPath, "utf-8"));
	} catch (e) {
		console.error("Failed to read auth.json:", e);
		process.exit(1);
	}

	const gitlabCred = authData["gitlab-duo"];
	if (!gitlabCred || gitlabCred.type !== "oauth" || !gitlabCred.access) {
		console.error("No gitlab-duo OAuth credentials found. Run /login gitlab-duo first.");
		process.exit(1);
	}

	console.log("Found gitlab-duo OAuth credentials");
	const gitlabAccessToken = gitlabCred.access;

	// Register our custom API provider
	console.log("\nRegistering gitlab-duo-api provider...");
	registerApiProvider({
		api: "gitlab-duo-api" as Api,
		stream: (model, context, options) => streamGitLabDuo(model, context, options as SimpleStreamOptions),
		streamSimple: streamGitLabDuo,
	});

	// Create a test model
	const modelDef = MODEL_DEFS[modelId];
	const mapping = MODEL_MAPPINGS[modelId];
	// Enable reasoning for Anthropic models
	const supportsReasoning = mapping.api === "anthropic-messages";
	const testModel: Model<Api> = {
		id: modelId,
		name: modelDef.name,
		api: "gitlab-duo-api" as Api,
		provider: "gitlab-duo",
		baseUrl: AI_GATEWAY_URL,
		reasoning: supportsReasoning,
		input: ["text"],
		cost: modelDef.cost,
		contextWindow: modelDef.contextWindow,
		maxTokens: modelDef.maxTokens,
	};

	// Create test context
	const context: Context = {
		messages: [{ role: "user", content: "Say hello in exactly 3 words.", timestamp: Date.now() }],
	};

	// Check for --thinking flag
	const useThinking = process.argv.includes("--thinking");

	console.log("\nStarting stream test...");
	console.log("Model:", testModel.id);
	console.log("Backend:", MODEL_MAPPINGS[modelId].backendModel);
	console.log("API:", MODEL_MAPPINGS[modelId].api);
	console.log(
		"Reasoning:",
		supportsReasoning ? (useThinking ? "enabled" : "supported but not enabled") : "not supported",
	);
	console.log("Prompt:", context.messages[0].content);
	console.log("");

	// Call streamSimple
	const stream = streamSimple(testModel, context, {
		apiKey: gitlabAccessToken,
		maxTokens: 100,
		reasoning: useThinking && supportsReasoning ? "low" : undefined,
	});

	// Consume the stream
	let inThinking = false;
	for await (const event of stream) {
		if (event.type === "thinking_start") {
			inThinking = true;
			console.log("[Thinking]");
		} else if (event.type === "thinking_delta") {
			process.stdout.write(event.delta);
		} else if (event.type === "thinking_end") {
			inThinking = false;
			console.log("\n[/Thinking]\n");
		} else if (event.type === "text_delta") {
			process.stdout.write(event.delta);
		} else if (event.type === "error") {
			console.error("\nError:", event.error.errorMessage);
		} else if (event.type === "done") {
			console.log("\n\nDone! Stop reason:", event.reason);
			console.log("Usage:", event.message.usage);
		}
	}
}

main().catch(console.error);