import {
	createAssistantMessageEventStream,
	type Api,
	type AssistantMessage,
	type Context,
	type Model,
	type ProviderHeaders,
	type SimpleStreamOptions,
} from "@earendil-works/pi-ai";
import { createGrammarToolInputProperties } from "./constrained-sampling.js";
import type { ExtensionAPI, ModelRegistry } from "@earendil-works/pi-coding-agent";
import type { ResponseCreateParamsStreaming } from "openai/resources/responses/responses.js";
import type { CodexConversionConfig } from "../adapter/activation/config.ts";
import type { ExecutionMode } from "../adapter/activation/execution-mode.ts";
import { resolveCodexRuntimePlan, resolveCodexRuntimePlanForState } from "../adapter/activation/runtime-plan.ts";
import { buildRequestBody } from "./openai-codex/request-body.ts";
import { applyResponsesLiteRequest, isResponsesLiteRequest, namespaceExistingResponsesLiteRequest, prepareResponsesLiteRequestImages, RESPONSES_LITE_HEADER } from "./openai-codex/responses-lite.ts";
import { assertSuccessfulCodexOutput, processCodexResponsesStream } from "./openai-codex/stream-events.ts";
import type { OpenAICodexStreamOptions, ResponsesBody, StreamEventShape } from "./openai-codex/types.ts";
import {
	hasContextNamespaceRouters,
	routeContextNamespaceToolStream,
} from "../context-management/namespace-tools.ts";

function initialAssistantMessage<TApi extends Api>(model: Model<TApi>): AssistantMessage {
	return {
		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: "pending",
		timestamp: Date.now(),
	};
}

function mergeHeaders(...groups: Array<ProviderHeaders | undefined>): ProviderHeaders {
	const headers = new Map<string, { name: string; value: string | null }>();
	for (const group of groups) {
		for (const [name, value] of Object.entries(group ?? {})) {
			headers.set(name.toLowerCase(), { name, value });
		}
	}
	return Object.fromEntries([...headers.values()].map(({ name, value }) => [name, value]));
}

function hasHeader(headers: ProviderHeaders | undefined, name: string): boolean {
	const expected = name.toLowerCase();
	return Object.entries(headers ?? {}).some(
		([key, value]) => key.toLowerCase() === expected && value !== null && value.trim() !== "",
	);
}

function clientAuth(provider: string, apiKey: string | undefined, headers: ProviderHeaders): { apiKey: string; headers: ProviderHeaders } {
	if (apiKey) return { apiKey, headers };
	if (hasHeader(headers, "authorization")) return { apiKey: "unused", headers };
	if (hasHeader(headers, "cf-aig-authorization")) {
		return { apiKey: "unused", headers: mergeHeaders(headers, { Authorization: null }) };
	}
	throw new Error(`No API key for provider: ${provider}`);
}

async function reportErrorResponse<TApi extends Api>(
	error: unknown,
	options: SimpleStreamOptions | undefined,
	model: Model<TApi>,
	APIError: typeof import("openai").APIError,
): Promise<void> {
	if (!(error instanceof APIError) || error.status === undefined || !error.headers) return;
	await options?.onResponse?.({
		status: error.status,
		headers: Object.fromEntries(error.headers.entries()),
	}, model);
}

export function streamCodeModeResponsesProxy<TApi extends Api>(
	model: Model<TApi>,
	context: Context,
	options?: SimpleStreamOptions,
) {
	const stream = createAssistantMessageEventStream();
	const output = initialAssistantMessage(model);

	void (async () => {
		try {
			const { default: OpenAI, APIError } = await import("openai");
			const grammarToolInputProperties = createGrammarToolInputProperties(context.tools, true);
			const effectiveOptions = { ...options, grammarToolInputProperties };
			let headers = mergeHeaders(model.headers, options?.headers);
			let body: ResponsesBody = buildRequestBody(model, context, effectiveOptions);
			const rewritten = await options?.onPayload?.(body, model);
			if (rewritten !== undefined) body = rewritten as ResponsesBody;
			body = isResponsesLiteRequest(body)
				? namespaceExistingResponsesLiteRequest({ ...body, parallel_tool_calls: false })
				: applyResponsesLiteRequest(body);
			body = await prepareResponsesLiteRequestImages(body);
			headers = mergeHeaders(headers, { [RESPONSES_LITE_HEADER]: "true" });

			const auth = clientAuth(model.provider, options?.apiKey, headers);
			const client = new OpenAI({
				apiKey: auth.apiKey,
				baseURL: model.baseUrl,
				defaultHeaders: auth.headers,
			});
			let response;
			try {
				response = await client.responses.create(
					body as unknown as ResponseCreateParamsStreaming,
					{
						...(options?.signal ? { signal: options.signal } : {}),
						...(options?.timeoutMs !== undefined ? { timeout: options.timeoutMs } : {}),
						maxRetries: options?.maxRetries ?? 0,
					},
				).withResponse();
			} catch (error) {
				await reportErrorResponse(error, options, model, APIError);
				throw error;
			}
			await options?.onResponse?.({
				status: response.response.status,
				headers: Object.fromEntries(response.response.headers.entries()),
			}, model);

			stream.push({ type: "start", partial: output });
			await processCodexResponsesStream(
				response.data as unknown as AsyncIterable<StreamEventShape>,
				output,
				stream,
				model,
				effectiveOptions as OpenAICodexStreamOptions,
			);
			if (options?.signal?.aborted) throw new Error("Request was aborted");
			assertSuccessfulCodexOutput(output);
			stream.push({ type: "done", reason: output.stopReason, message: output });
			stream.end();
		} catch (error) {
			for (const block of output.content) {
				if (typeof block === "object" && block !== null) delete (block as { partialJson?: unknown }).partialJson;
			}
			output.stopReason = options?.signal?.aborted ? "aborted" : "error";
			output.errorMessage = error instanceof Error ? error.message : String(error);
			stream.push({ type: "error", reason: output.stopReason, error: output });
			stream.end();
		}
	})();

	return stream;
}

export interface CodeModeProxyProviderRegistration {
	applyConfig(config: CodexConversionConfig, modelRegistry: CodeModeModelRegistry): void;
	shutdown(): void;
}

type CodeModeModelRegistry = Pick<ModelRegistry, "getAll" | "getProvider" | "getRegisteredProviderConfig">;
type RegisteredProviderConfig = Parameters<ExtensionAPI["registerProvider"]>[1];
type ResponsesApi = "openai-responses" | "openai-codex-responses";

function resolveProviderApis(
	config: CodexConversionConfig,
	executionMode: ExecutionMode | undefined,
	modelRegistry: CodeModeModelRegistry,
): Map<string, ResponsesApi> {
	const resolved = new Map<string, ResponsesApi>();
	for (const model of modelRegistry.getAll()) {
		if (
			model.api !== "openai-responses" &&
			model.api !== "openai-codex-responses"
		)
			continue;
		const api: ResponsesApi = model.api === "openai-responses"
			? "openai-responses"
			: "openai-codex-responses";
		const plan = resolveCodexRuntimePlan({ model }, config, executionMode);
		const mode = executionMode ?? config.executionMode;
		const configuredResponsesLite =
			model.api === "openai-responses" &&
			!config.voiceFeaturesOnly &&
			(mode === "code" || mode === "notebook") &&
			config.openai.proxyResponsesLite &&
			config.scope.additionalProviders.includes(
				model.provider.trim().toLowerCase(),
			);
		if (configuredResponsesLite || plan.contextManagement)
			resolved.set(model.provider, api);
	}
	return resolved;
}

export function registerCodeModeProxyProvider(
	pi: ExtensionAPI,
	getConfig: () => CodexConversionConfig,
	getExecutionMode: () => ExecutionMode | undefined = () => undefined,
	getAvailableToolNames: () => string[] | undefined = () => undefined,
): CodeModeProxyProviderRegistration {
	const registeredProviders = new Map<string, {
		previous: RegisteredProviderConfig | undefined;
		overlayStream: NonNullable<RegisteredProviderConfig["streamSimple"]>;
		modelRegistry: CodeModeModelRegistry;
		api: ResponsesApi;
	}>();
	const restoreProvider = (provider: string, registration: NonNullable<ReturnType<typeof registeredProviders.get>>) => {
		const current = registration.modelRegistry.getRegisteredProviderConfig?.(provider) as RegisteredProviderConfig | undefined;
		if (!current || current.streamSimple !== registration.overlayStream) return;
		const restored = { ...current } as RegisteredProviderConfig;
		if (registration.previous?.streamSimple) restored.streamSimple = registration.previous.streamSimple;
		else delete restored.streamSimple;
		if (registration.previous?.api) restored.api = registration.previous.api;
		else if (!registration.previous?.streamSimple && current.api === registration.api) delete restored.api;
		pi.unregisterProvider(provider);
		if (Object.keys(restored).length > 0) pi.registerProvider(provider, restored);
	};
	const shutdown = () => {
		for (const [provider, registration] of registeredProviders) restoreProvider(provider, registration);
		registeredProviders.clear();
	};
	const applyConfig = (config: CodexConversionConfig, modelRegistry: CodeModeModelRegistry) => {
		const desiredProviders = resolveProviderApis(
			config,
			getExecutionMode(),
			modelRegistry,
		);
		for (const [provider, api] of desiredProviders) {
			const existing = registeredProviders.get(provider);
			if (existing?.api === api) continue;
			if (existing) {
				restoreProvider(provider, existing);
				registeredProviders.delete(provider);
			}
			const previous = modelRegistry.getRegisteredProviderConfig(provider) as RegisteredProviderConfig | undefined;
			if (
				previous?.streamSimple &&
				previous.api !== "openai-responses" &&
				previous.api !== "openai-codex-responses"
			)
				continue;
			const fallbackProvider = modelRegistry.getProvider(provider);
			if (!fallbackProvider) throw new Error(`Cannot overlay missing provider: ${provider}`);
			const overlayStream: NonNullable<RegisteredProviderConfig["streamSimple"]> = (model, context, options) => {
				const plan = resolveCodexRuntimePlanForState(
					{ model },
					{
						config: getConfig(),
						executionMode: getExecutionMode() ?? getConfig().executionMode,
						availableToolNames: getAvailableToolNames(),
					},
				);
				if (
					model.api === "openai-responses" &&
					plan.transport === "responses-lite"
				)
					return streamCodeModeResponsesProxy(model, context, options);
				const stream = fallbackProvider.streamSimple(
					model as never,
					context,
					options,
				);
				return plan.contextManagement && hasContextNamespaceRouters(context)
					? routeContextNamespaceToolStream(stream)
					: stream;
			};
			pi.registerProvider(provider, {
				api,
				streamSimple: overlayStream,
			});
			registeredProviders.set(provider, {
				previous,
				overlayStream,
				modelRegistry,
				api,
			});
		}
		for (const [provider, registration] of registeredProviders) {
			if (desiredProviders.has(provider)) continue;
			restoreProvider(provider, registration);
			registeredProviders.delete(provider);
		}
	};

	return { applyConfig, shutdown };
}