/**
 * Pure state logic of the chat: a reducer folding server messages and
 * connection changes into the UI state. Keeping it a plain function makes the
 * whole conversation flow testable without React or a WebSocket.
 */

import type {
	AcpAvailableCommand,
	AcpContentBlock,
	AcpPlanEntry,
	AcpPermissionOption,
	AcpToolCallContent,
	AcpToolCallLocation,
	ServerMessage,
} from "./protocol";

/** The UI's view of one tool call, folded from tool_call / tool_call_update. */
export interface ToolCall {
	toolCallId: string;
	title: string;
	toolKind: string;
	status: string;
	contents: AcpToolCallContent[];
	locations: AcpToolCallLocation[];
}

/** One entry of the conversation thread. */
export type ThreadItem =
	| { kind: "user"; text: string }
	| { kind: "agent"; text: string; closed: boolean }
	| { kind: "thought"; text: string; closed: boolean }
	| { kind: "tool"; call: ToolCall }
	| { kind: "error"; text: string };

/** A permission request awaiting the user's decision. */
export interface PermissionRequestState {
	requestId: string;
	title: string;
	options: AcpPermissionOption[];
}

export type ConnectionStatus = "connecting" | "online" | "offline";

/** The whole state the UI renders from. */
export interface ChatState {
	connection: ConnectionStatus;
	sessionId: string | null;
	turnActive: boolean;
	thread: ThreadItem[];
	plan: AcpPlanEntry[];
	permissions: PermissionRequestState[];
	/** Slash commands the agent announced (available_commands_update). */
	commands: AcpAvailableCommand[];
}

export const initialState: ChatState = {
	connection: "connecting",
	sessionId: null,
	turnActive: false,
	thread: [],
	plan: [],
	permissions: [],
	commands: [],
};

/** Everything that can change the state. */
export type ChatEvent =
	| { type: "server"; message: ServerMessage }
	| { type: "connection"; status: ConnectionStatus };

/**
 * reduce returns the next state for an event. It never mutates its input.
 *
 * @example
 * let state = initialState;
 * state = reduce(state, { type: "server", message: { type: "hello", sessionId: "s1" } });
 * state.sessionId; // "s1"
 */
export function reduce(state: ChatState, event: ChatEvent): ChatState {
	if (event.type === "connection") {
		return { ...state, connection: event.status };
	}
	// Unknown or not-yet-handled message types must never break the UI.
	const handler = serverHandlers[event.message.type];
	return handler ? handler(state, event.message) : state;
}

type MessageHandler = (state: ChatState, msg: ServerMessage) => ChatState;

/** One handler per server message type; unknown types fall through in reduce. */
const serverHandlers: Record<string, MessageHandler> = {
	// A hello opens every (re)connection, right before the full history
	// replay: resetting everything here makes reconnects idempotent.
	hello: (state, msg) => ({
		...state,
		connection: "online",
		sessionId: msg.sessionId ?? null,
		turnActive: msg.turnActive ?? false,
		thread: [],
		plan: [],
		permissions: [],
		commands: [],
	}),
	user_message: (state, msg) =>
		withItem(state, { kind: "user", text: msg.text ?? "" }),
	turn_started: (state) => ({ ...state, turnActive: true }),
	turn_ended: (state) => ({
		...state,
		turnActive: false,
		thread: closeStreamItems(state.thread),
	}),
	session_update: reduceUpdate,
	permission_request: addPermission,
	permission_resolved: (state, msg) => ({
		...state,
		permissions: state.permissions.filter((p) => p.requestId !== msg.requestId),
	}),
	error: (state, msg) =>
		withItem(state, { kind: "error", text: msg.message ?? "unknown error" }),
};

/**
 * One handler per ACP update kind; the rest (current_mode_update, ...) are
 * not rendered yet, and ignoring them keeps the UI forward-compatible.
 */
const updateHandlers: Record<string, MessageHandler> = {
	agent_message_chunk: (state, msg) =>
		appendStreamText(state, "agent", blockText(msg.update?.content)),
	agent_thought_chunk: (state, msg) =>
		appendStreamText(state, "thought", blockText(msg.update?.content)),
	tool_call: upsertToolCall,
	tool_call_update: upsertToolCall,
	plan: (state, msg) => ({ ...state, plan: msg.update?.entries ?? [] }),
	available_commands_update: (state, msg) => ({
		...state,
		commands: msg.update?.availableCommands ?? [],
	}),
};

function reduceUpdate(state: ChatState, msg: ServerMessage): ChatState {
	const handler = msg.update
		? updateHandlers[msg.update.sessionUpdate]
		: undefined;
	return handler ? handler(state, msg) : state;
}

/** blockText extracts the text of a single content block, if any. */
function blockText(
	content: AcpContentBlock | AcpToolCallContent[] | undefined,
): string {
	if (!content || Array.isArray(content)) {
		return "";
	}
	return content.text ?? "";
}

function withItem(state: ChatState, item: ThreadItem): ChatState {
	return { ...state, thread: [...state.thread, item] };
}

/** Streamed chunks accumulate into the last open item of the same kind. */
function appendStreamText(
	state: ChatState,
	kind: "agent" | "thought",
	text: string,
): ChatState {
	const thread = [...state.thread];
	const last = thread[thread.length - 1];
	if (last && last.kind === kind && !last.closed) {
		thread[thread.length - 1] = { ...last, text: last.text + text };
		return { ...state, thread };
	}
	thread.push({ kind, text, closed: false });
	return { ...state, thread };
}

/**
 * upsertToolCall folds a tool_call or tool_call_update into the thread:
 * updates merge only the fields they carry, and an update for an unknown id
 * creates the entry (the adapter may resume a call after a reconnect).
 */
function upsertToolCall(state: ChatState, msg: ServerMessage): ChatState {
	const update = msg.update;
	if (!update || !update.toolCallId) {
		return state;
	}

	const contents = Array.isArray(update.content) ? update.content : undefined;
	const thread = [...state.thread];
	const index = thread.findIndex(
		(item) =>
			item.kind === "tool" && item.call.toolCallId === update.toolCallId,
	);

	if (index >= 0) {
		const existing = thread[index] as Extract<ThreadItem, { kind: "tool" }>;
		thread[index] = {
			kind: "tool",
			call: {
				...existing.call,
				title: update.title ?? existing.call.title,
				toolKind: update.kind ?? existing.call.toolKind,
				status: update.status ?? existing.call.status,
				contents: contents ?? existing.call.contents,
				locations: update.locations ?? existing.call.locations,
			},
		};
		return { ...state, thread };
	}

	thread.push({
		kind: "tool",
		call: {
			toolCallId: update.toolCallId,
			title: update.title ?? "Tool call",
			toolKind: update.kind ?? "other",
			status: update.status ?? "pending",
			contents: contents ?? [],
			locations: update.locations ?? [],
		},
	});
	return { ...state, thread };
}

function addPermission(state: ChatState, msg: ServerMessage): ChatState {
	if (
		!msg.requestId ||
		state.permissions.some((p) => p.requestId === msg.requestId)
	) {
		return state;
	}
	return {
		...state,
		permissions: [
			...state.permissions,
			{
				requestId: msg.requestId,
				title: msg.request?.toolCall?.title ?? "The agent asks for permission",
				options: msg.request?.options ?? [],
			},
		],
	};
}

/** turn_ended closes every streaming item so the next turn starts fresh. */
function closeStreamItems(thread: ThreadItem[]): ThreadItem[] {
	return thread.map((item) =>
		item.kind === "agent" || item.kind === "thought"
			? { ...item, closed: true }
			: item,
	);
}