import type { Api, AssistantMessage, Model, SimpleStreamOptions, TranscriptContext } from "@earendil-works/pi-ai";
import type { ResponseCreateParamsStreaming } from "openai/resources/responses/responses.js";
import type { CodexCompactionDiagnostic } from "../../adapter/compaction/diagnostics.ts";
import type { CodexCacheKeepaliveStrategy } from "../../adapter/activation/cache-keepalive.ts";
import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import type { WebSocketContinuationDecision, WebSocketContinuationMismatch } from "./continuation-diagnostics.ts";
export type { WebSocketContinuationDecision } from "./continuation-diagnostics.ts";

export interface WebSocketLike {
	readyState?: number | undefined;
	send(data: string): void;
	close(code?: number, reason?: string): void;
	addEventListener(type: string, listener: (event: unknown) => void): void;
	removeEventListener(type: string, listener: (event: unknown) => void): void;
}

export interface WebSocketConstructorLike {
	new (url: string, options?: { headers?: Record<string, string> | undefined } | string | string[]): WebSocketLike;
}

export interface SessionWebSocketCacheEntry {
	socket: WebSocketLike;
	busy: boolean;
	createdAtMs: number;
	continuation?: CachedWebSocketContinuationState | undefined;
}

export interface AcquiredWebSocket {
	socket: WebSocketLike;
	routeKey?: string | undefined;
	entry?: SessionWebSocketCacheEntry | undefined;
	reused: boolean;
	socketAgeMs: number;
	release: (options?: { keep?: boolean | undefined }) => void;
}

export interface CachedWebSocketContinuationState {
	lastRequestBody: ResponsesBody;
	lastResponseId: string;
	lastResponseItems: unknown[];
}

export type CanonicalHistoryDecision =
	| "compaction"
	| "identity_mismatch"
	| "input_shorter_than_baseline"
	| "request_prefix_mismatch"
	| "response_prefix_mismatch"
	| "validated";

export type CodexDiagnosticsLane = "response" | "compaction" | "prewarm";
export type CodexPrewarmKind = "ordinary" | "compaction" | "keepalive";
export interface CodexPrewarmDiagnostics {
	kind: CodexPrewarmKind;
	keepaliveStrategy?: CodexCacheKeepaliveStrategy | undefined;
	requestSource?: "captured" | "reconstructed" | undefined;
}
export type CodexDiagnosticsTransport = "websocket" | "sse";
export type CodexDiagnosticsFailureCategory =
	| "aborted"
	| "authentication"
	| "connection"
	| "connection_limit"
	| "message_too_big"
	| "overload"
	| "previous_response_missing"
	| "protocol"
	| "rate_limit"
	| "timeout"
	| "transport"
	| "unknown";
export interface CodexDiagnosticsFailure {
	category: CodexDiagnosticsFailureCategory;
	code?: string | undefined;
	status?: number | undefined;
}
export interface CodexRequestDiagnostics {
	instructionsFingerprint: string;
	toolsFingerprint: string;
	cacheKeyFingerprint: string;
	routingFingerprint: string;
	optionsFingerprint: string;
	headersSent: boolean;
	reasoningEffort: string;
	serviceTier: string;
	accessProgram: "standard" | "daybreak_blue" | "daybreak_red" | "unset";
	responsesLite: boolean;
	routingHint: "present" | "absent" | "guardian_suppressed";
	guardianReviewer: boolean;
	parentThreadLinked: boolean;
	subagentKind: string;
}

export type CodexDiagnosticsEvent = {
	scope?: "session" | "auxiliary" | undefined;
	requestId?: string | undefined;
	sessionRequestOrdinal?: number | undefined;
	requestCounterProcessId?: string | undefined;
} & (
	| {
			type: "connection";
			socketId: string;
			phase: "created" | "observed" | "client_close" | "closed";
			initiator: "client" | "peer_or_unknown";
			cause?: string | undefined;
			closeCode?: number | undefined;
			closeReason?: "normal" | "going_away" | "message_too_big" | "abnormal" | "other" | undefined;
	  }
	| {
			type: "timing";
			lane: CodexDiagnosticsLane;
			transport: CodexDiagnosticsTransport;
			requestId: string;
			attempt: number;
			outcome: "completed" | "failed" | "aborted";
			requestElapsedMs: number;
			attemptElapsedMs: number;
			sendMs: number | null;
			firstEventMs: number | null;
			firstVisibleTextMs: number | null;
			completionMs: number | null;
			outputTokens: number | null;
			reasoningTokens: number | null;
			outputTokensPerSecond: number | null;
			backend: Record<string, number | string | null>;
	  }
	| {
			type: "request";
			lane: CodexDiagnosticsLane;
			transport: CodexDiagnosticsTransport;
			attempt: number;
			fullInputItems: number;
			sentInputItems: number;
			model?: string | undefined;
			socketReused?: boolean | undefined;
			socketId?: string | undefined;
			socketAgeMs?: number | undefined;
			socketRequestOrdinal?: number | undefined;
			socketLane?: "main" | "keepalive" | undefined;
			continuation?: WebSocketContinuationDecision | undefined;
			continuationMismatch?: WebSocketContinuationMismatch | undefined;
			continuationBaselineInputItems?: number | undefined;
			continuationBaselineResponseItems?: number | undefined;
			canonicalHistory?: CanonicalHistoryDecision | undefined;
			prewarm?: CodexPrewarmDiagnostics | undefined;
			compaction?: CodexCompactionDiagnostic | undefined;
			previousResponseId?: boolean | undefined;
			requestMetadata?: CodexRequestDiagnostics | undefined;
	  }
	| {
			type: "usage";
			lane: Exclude<CodexDiagnosticsLane, "prewarm">;
			transport: CodexDiagnosticsTransport;
			inputTokens: number;
			cachedInputTokens: number;
			cacheWriteInputTokens: number;
			outputTokens: number;
	  }
	| {
			type: "retry";
			lane: Exclude<CodexDiagnosticsLane, "prewarm">;
			transport: CodexDiagnosticsTransport;
			attempt: number;
			delayMs?: number | undefined;
			failure: CodexDiagnosticsFailure;
	  }
	| {
			type: "fallback";
			lane: CodexDiagnosticsLane;
			from: CodexDiagnosticsTransport;
			to: CodexDiagnosticsTransport;
			reason: "upgrade_required" | "message_too_big" | "unauthorized" | "retry_budget_exhausted";
	  }
	| {
			type: "failure";
			lane: CodexDiagnosticsLane;
			transport: CodexDiagnosticsTransport;
			failure: CodexDiagnosticsFailure;
	  }
	| {
			type: "prewarm-ready";
			transport: "websocket";
			socketId?: string | undefined;
			socketReused: boolean;
			socketAgeMs: number;
			socketLane: "main" | "keepalive";
			prewarm: CodexPrewarmDiagnostics;
			usage?: CodexPrewarmUsage | undefined;
	  }
	| {
			type: "keepalive";
			phase: "armed" | "started" | "applied" | "skipped";
			strategy: CodexCacheKeepaliveStrategy;
			intervalMs?: number | undefined;
			requestSource?: "captured" | "reconstructed" | undefined;
			action?: "generated-refresh" | undefined;
	  });

export type CodexDiagnosticsSink = (event: CodexDiagnosticsEvent) => void;

export type CodexUsageRecorder = (accountId: string, model: Model<Api>, usage: AssistantMessage["usage"]) => Promise<void>;

export interface CachedWebSocketRequestBodyResult {
	body: ResponsesBody;
	decision: WebSocketContinuationDecision;
	mismatch?: WebSocketContinuationMismatch | undefined;
}

export type ServiceTier = ResponseCreateParamsStreaming["service_tier"] | "ultrafast";
export interface AccessPrograms {
	cyber: "standard" | "daybreak_blue" | "daybreak_red";
}
export type ProviderEnv = Record<string, string>;
export type CodexProviderStreamOptions = Omit<SimpleStreamOptions, "toolChoice"> & {
	serviceTier?: ServiceTier | undefined;
	accessPrograms?: AccessPrograms | undefined;
	textVerbosity?: string | undefined;
	reasoningSummary?: string | null | undefined;
	toolChoice?: "auto" | "none" | "required" | undefined;
};
export type CodexReasoningEffort = "none" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max";
export type OpenAICodexStreamOptions = CodexProviderStreamOptions & {
	reasoningEffort?: CodexReasoningEffort | undefined;
	responsesLite?: boolean | undefined;
	grammarToolInputProperties?: ReadonlyMap<string, string> | undefined;
	onOutputItemDone?: ((item: unknown) => void) | undefined;
	websocketConnectTimeoutMs?: number | undefined;
	env?: ProviderEnv | undefined;
	canonicalCompaction?: boolean | undefined;
	/** Host-only operation context. Never serialized into a provider request. */
	remoteDeliveryContext?: ExtensionContext | undefined;
	compactionDiagnostics?: CodexCompactionDiagnostic | undefined;
};

export interface ResponsesBody {
	model: string;
	store: boolean;
	stream: boolean;
	instructions?: string | undefined;
	previous_response_id?: string | undefined;
	input: unknown[];
	text: { verbosity: string };
	include: string[];
	prompt_cache_key?: string | undefined;
	tool_choice: "auto" | "none" | "required";
	parallel_tool_calls: boolean;
	temperature?: number | undefined;
	service_tier?: string | undefined;
	access_programs?: AccessPrograms | undefined;
	tools?: unknown[] | undefined;
	reasoning?: {
		effort?: string | undefined;
		summary?: string | undefined;
		context?: "all_turns" | undefined;
	} | undefined;
	client_metadata?: Record<string, string> | undefined;
	[key: string]: unknown;
}

export type BeforeCodexRequestSend = (
	model: Model<Api>,
	context: TranscriptContext,
	payload: ResponsesBody,
	options: OpenAICodexStreamOptions | undefined,
	responsesLite: boolean,
) => void | Promise<void>;

export interface CodexPrewarmUsage {
	inputTokens: number;
	cachedInputTokens: number;
	cacheWriteInputTokens: number;
	outputTokens?: number | undefined;
}

export interface CodexPrewarmResult {
	socketReused: boolean;
	usage?: CodexPrewarmUsage | undefined;
}

export interface ResponseEnvelope {
	id?: string | undefined;
	status?: string | undefined;
	usage?: {
		input_tokens?: number | undefined;
		output_tokens?: number | undefined;
		total_tokens?: number | undefined;
		input_tokens_details?: { cached_tokens?: number | undefined; cache_write_tokens?: number | undefined } | undefined;
		output_tokens_details?: { reasoning_tokens?: number | undefined } | undefined;
	} | undefined;
	service_tier?: string | undefined;
	error?: {
		code?: string | undefined;
		type?: string | undefined;
		message?: string | undefined;
		status?: number | string | undefined;
		status_code?: number | string | undefined;
		[key: string]: unknown;
	} | undefined;
	[key: string]: unknown;
}

export interface StreamEventShape {
	type?: string | undefined;
	headers?: Record<string, unknown> | undefined;
	response?: ResponseEnvelope | undefined;
	item?: {
		id?: string | undefined;
		type?: string | undefined;
		result?: string | null | undefined;
		output_format?: string | undefined;
		revised_prompt?: string | undefined;
		status?: string | undefined;
		[key: string]: unknown;
	} | undefined;
	code?: string | undefined;
	message?: string | undefined;
	[key: string]: unknown;
}

export function createInitialAssistantMessage(model: { provider: string; id: string }): AssistantMessage {
	return {
		role: "assistant",
		content: [],
		api: "openai-codex-responses",
		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(),
	};
}