import type {
	LanguageModel,
	LanguageModelCallOptions,
	ModelMessage,
	Output,
	SystemModelMessage,
	TelemetryOptions,
	ToolCallRepairFunction,
	ToolSet,
} from 'ai';

import type { FinishReason, SerializableAgentState, TokenUsage } from '../../types';
import type { ExecutionOptions, RunOptions } from '../../types/sdk/agent';
import type { AgentMessage } from '../../types/sdk/message';
import type { JSONObject } from '../../types/utils/json';
import type { AgentMessageList } from '../model/message-list';
import type { ToolCallBatchResult, ToolCallSuspension } from '../tools/tool-call-executor';

type RunCallOptions = (RunOptions & ExecutionOptions) | undefined;

/**
 * Recognized ways a model turn can fail while the provider still returns
 * "successfully" (no thrown SDK error). Extend this as new failure modes are
 * recognized so consumers switch on `type` instead of probing separate fields.
 * - `prompt_blocked`: the provider rejected the prompt (e.g. a safety block).
 * - `no_output`: the model returned nothing with a non-stop finish reason.
 */
export type ModelTurnErrorType = 'prompt_blocked' | 'no_output';

/** A recognized model-turn failure: a `type` to branch on and a user-facing `message`. */
export interface ModelTurnError {
	type: ModelTurnErrorType;
	message: string;
}

/** Normalized result of a single LLM turn, independent of generate vs stream. */
export interface ModelTurnResult {
	/** Raw AI SDK finish reason (used to detect the `tool-calls` continuation). */
	aiFinishReason: string;
	finishReason: FinishReason;
	/** Normalized usage, including any cache-token details available from the provider. */
	usage: TokenUsage | undefined;
	newMessages: AgentMessage[];
	toolCalls: Array<{
		toolCallId: string;
		toolName: string;
		input: unknown;
		providerExecuted?: boolean;
	}>;
	/** Resolved structured output, only when an output spec is set and the turn finished. */
	structuredOutput: unknown;
	/**
	 * Set when the turn failed in a recognized way despite the provider returning
	 * without a thrown error (e.g. a prompt safety block, or empty output). The
	 * loop surfaces this instead of ending the run silently. Check this single
	 * field rather than probing per-failure-mode properties.
	 */
	errorReason?: ModelTurnError;
}

/** Per-iteration inputs for the LLM call, assembled by the shared loop. */
export interface ModelCallContext {
	model: LanguageModel;
	system: SystemModelMessage | SystemModelMessage[];
	messages: ModelMessage[];
	abortSignal: AbortSignal;
	hasTools: boolean;
	aiTools: ToolSet;
	reasoning?: LanguageModelCallOptions['reasoning'];
	providerOptions?: Record<string, JSONObject>;
	outputSpec?: ReturnType<typeof Output.object>;
	maxOutputTokens?: number;
	aiSdkOptions: {
		telemetry?: TelemetryOptions;
		repairToolCall?: ToolCallRepairFunction<NoInfer<ToolSet>>;
	};
}

/** Data needed to emit a terminal suspension result. */
export interface SuspendEmission {
	suspendRunId: string;
	list: AgentMessageList;
	usage: TokenUsage | undefined;
	suspensions: ToolCallSuspension[];
}

/** Data needed to emit a terminal completion result. */
export interface CompleteEmission {
	list: AgentMessageList;
	options: RunCallOptions;
	finishReason: FinishReason;
	usage: TokenUsage | undefined;
	structuredOutput: unknown;
}

/**
 * Runtime services the sinks call into for shared terminal concerns (cost,
 * memory persistence, title generation, telemetry flush, state). Implemented by
 * `AgentRuntime` and passed to each sink.
 */
export interface RunServices {
	readonly runId: string;
	readonly modelId: string;
	applyCost(usage: TokenUsage | undefined): TokenUsage | undefined;
	saveToMemory(list: AgentMessageList, options: RunCallOptions): Promise<void>;
	maybeGenerateTitle(list: AgentMessageList, options: RunCallOptions): Promise<void>;
	flushTelemetry(options: RunCallOptions): Promise<void>;
	cleanupRun(): Promise<void>;
	updateState(patch: Partial<SerializableAgentState>): void;
	emitAgentEnd(messages: AgentMessage[]): void;
	getState(): SerializableAgentState;
}

/**
 * Strategy that adapts the shared agentic loop to a concrete output channel:
 * `GenerateSink` accumulates a `GenerateResult`; `StreamSink` writes
 * `StreamChunk`s through the stream writer guard.
 */
export interface RunOutputSink<TResult> {
	/** Run one LLM turn. Streaming implementations also emit text/tool chunks here. */
	callModel(ctx: ModelCallContext): Promise<ModelTurnResult>;
	/**
	 * Report the run's cumulative token usage after each turn. Streaming
	 * implementations keep it so an aborted run can still emit a terminal
	 * finish chunk carrying the tokens consumed before the stop.
	 */
	reportUsage(usage: TokenUsage | undefined): void;
	/**
	 * Signal that the just-returned turn's messages have been folded into the list.
	 * Streaming implementations drop the retained streamed text here so a later abort
	 * can't duplicate a turn that is already in the list — while a stop landing before
	 * this call still recovers the turn from the retained text.
	 */
	onTurnFolded?(): void;
	/** Emit the results/errors of a completed tool-call batch. */
	emitToolBatch(batch: ToolCallBatchResult): Promise<void>;
	/** Produce the terminal result when the run suspends for HITL / suspend-resume. */
	finishSuspended(emission: SuspendEmission): Promise<TResult>;
	/** Produce the terminal result when the run completes normally. */
	finishComplete(emission: CompleteEmission): Promise<TResult>;
}