import { mkdir, writeFile } from "node:fs/promises";
import path from "node:path";
import {
	createUIMessageStream,
	createUIMessageStreamResponse,
	type UIMessage,
	type UIMessageStreamWriter,
} from "ai";
import {
	accountSessionInputRoot,
	accountSessionOutputRoot,
	accountSessionRoot,
	accountSessionScratchpadRoot,
	getCurrentAccount,
} from "@/lib/claw-auth";

export const runtime = "nodejs";
export const maxDuration = 3600;

type ClawChatResponse = {
	final_output?: string;
	session_id?: string;
	tool_calls?: number;
	stop_reason?: string;
	elapsed_ms?: number;
	usage?: UsageSummary;
	transcript?: ClawTranscriptEntry[];
	error?: string;
	detail?: string;
};

type ClawRuntimeEvent = {
	type?: string;
	run_id?: string;
	tool_call_index?: number | null;
	tool_name?: string;
	tool_call_id?: string;
	arguments?: unknown;
	arguments_delta?: string;
	assistant_content?: string;
	metadata?: Record<string, unknown>;
	ok?: boolean;
	delta?: string;
	finish_reason?: string | null;
	elapsed_ms?: number;
};

type ClawStreamItem =
	| { kind: "event"; event?: ClawRuntimeEvent }
	| { kind: "result"; payload?: ClawChatResponse }
	| { kind: "error"; error?: string; detail?: string; status_code?: number };

type StreamState = {
	textStarted: boolean;
	textStreamed: boolean;
	phase: "waiting" | "queued" | "running" | "done";
	runId?: string;
	announcedToolCallKeys: Set<string>;
	streamedToolCallIds: Set<string>;
};

type ChatRequestBody = {
	id?: string;
	messages: UIMessage[];
	resumeSessionId?: string;
};

type ClawTranscriptEntry = {
	role?: string;
	content?: string;
	tool_call_id?: string;
	tool_calls?: ClawToolCall[];
};

type ClawToolCall = {
	id?: string;
	name?: string;
	arguments?: unknown;
	function?: {
		name?: string;
		arguments?: unknown;
	};
};

type UsageSummary = {
	input_tokens?: number;
	output_tokens?: number;
	total_tokens?: number;
	reasoning_tokens?: number;
};

const CLAW_API_URL = process.env.CLAW_API_URL ?? "http://127.0.0.1:8765";

export async function POST(req: Request) {
	const account = await getCurrentAccount();
	if (!account) return new Response("请先登录账号", { status: 401 });

	const { id, messages, resumeSessionId }: ChatRequestBody = await req.json();
	const resumeId = resumeSessionId ?? getLastSessionId(messages);
	const sessionId = safeSessionId(resumeId ?? id ?? crypto.randomUUID());
	await ensureSessionDirectories(account.id, sessionId);
	const userPrompt = await getLastUserText(messages, account.id, sessionId);

	if (!userPrompt) {
		return new Response("Prompt is empty", { status: 400 });
	}
	const runtimeContext = renderSessionRuntimeContext(account.id, sessionId);

	const stream = createUIMessageStream({
		execute: async ({ writer }) => {
			const streamState: StreamState = {
				textStarted: false,
				textStreamed: false,
				phase: "waiting",
				announcedToolCallKeys: new Set<string>(),
				streamedToolCallIds: new Set<string>(),
			};
			writer.write({ type: "start" });
			writer.write({ type: "start-step" });

			try {
				const startedAt = Date.now();
				writer.write({ type: "reasoning-start", id: "reasoning-1" });
				writer.write({
					type: "reasoning-delta",
					id: "reasoning-1",
					delta: "请求已发送，等待后端开始处理。",
				});
				const payload = await callClawBackendStream(
					userPrompt,
					runtimeContext,
					account.id,
					sessionId,
					resumeId,
					writer,
					streamState,
				);
				const elapsedMs = payload.elapsed_ms ?? Date.now() - startedAt;
				const text = formatClawResponse(payload);

				writer.write({
					type: "reasoning-delta",
					id: "reasoning-1",
					delta: `\n思考并执行完成，用时 ${formatDuration(elapsedMs)}。`,
				});
				writer.write({ type: "reasoning-end", id: "reasoning-1" });
				writeToolTrace(writer, payload.transcript, streamState);
				if (!streamState.textStreamed) {
					writer.write({ type: "text-start", id: "text-1" });
					writer.write({ type: "text-delta", id: "text-1", delta: text });
					writer.write({ type: "text-end", id: "text-1" });
				} else if (streamState.textStarted) {
					writer.write({ type: "text-end", id: "text-1" });
					streamState.textStarted = false;
				}
				writer.write({ type: "finish-step" });
				writer.write({
					type: "finish",
					finishReason: payload.error ? "error" : "stop",
					messageMetadata: {
						sessionId: payload.session_id,
						toolCalls: payload.tool_calls,
						stopReason: payload.stop_reason,
						elapsedMs,
						usage: payload.usage,
					},
				});
			} catch (err) {
				const text =
					err instanceof Error && err.name === "AbortError"
						? "请求已取消。"
						: err instanceof Error
							? err.message
							: "请求失败。";
				writer.write({ type: "reasoning-end", id: "reasoning-1" });
				writer.write({ type: "text-start", id: "text-1" });
				writer.write({ type: "text-delta", id: "text-1", delta: text });
				writer.write({ type: "text-end", id: "text-1" });
				writer.write({ type: "finish-step" });
				writer.write({
					type: "finish",
					finishReason: "error",
				});
			}
		},
	});

	return createUIMessageStreamResponse({
		stream,
		headers: {
			"Cache-Control": "no-cache, no-transform",
			"X-Accel-Buffering": "no",
		},
	});
}

function formatDuration(ms: number) {
	const safeMs = Math.max(0, Math.round(ms));
	if (safeMs < 1000) return `${safeMs}ms`;
	const totalSeconds = Math.round(safeMs / 1000);
	if (totalSeconds < 60) return `${totalSeconds}s`;
	const minutes = Math.floor(totalSeconds / 60);
	const seconds = totalSeconds % 60;
	return seconds > 0 ? `${minutes}min${seconds}s` : `${minutes}min`;
}

function writeToolTrace(
	writer: UIMessageStreamWriter<UIMessage>,
	transcript?: ClawTranscriptEntry[],
	streamState?: StreamState,
) {
	if (!transcript?.length) return;

	const toolResults = new Map<string, string>();
	for (const entry of transcript) {
		if (entry.role === "tool" && entry.tool_call_id) {
			toolResults.set(entry.tool_call_id, entry.content ?? "");
		}
	}

	for (const entry of transcript) {
		if (!entry.tool_calls?.length) continue;
		const stageNote = entry.role === "assistant" ? entry.content?.trim() : "";
		for (const call of entry.tool_calls) {
			const toolCallId = call.id;
			const toolName = call.function?.name ?? call.name;
			if (!toolCallId || !toolName) continue;
			if (streamState?.streamedToolCallIds.has(toolCallId)) continue;

			writer.write({
				type: "tool-input-available",
				toolCallId,
				toolName,
				input: attachStageNote(
					parseToolInput(call.function?.arguments ?? call.arguments),
					stageNote,
				),
			});

			if (toolResults.has(toolCallId)) {
				writer.write({
					type: "tool-output-available",
					toolCallId,
					output: toolResults.get(toolCallId),
				});
			}
		}
	}
}

function attachStageNote(input: unknown, stageNote?: string) {
	const note = stageNote?.trim();
	if (!note) return input;
	if (input && typeof input === "object" && !Array.isArray(input)) {
		return { ...input, __claw_stage_note: note };
	}
	return { value: input, __claw_stage_note: note };
}

function parseToolInput(value: unknown) {
	if (typeof value !== "string") return value ?? {};
	try {
		return JSON.parse(value);
	} catch {
		return value;
	}
}

async function getLastUserText(
	messages: UIMessage[],
	accountId: string,
	threadId?: string,
) {
	const lastUserMessage = [...messages]
		.reverse()
		.find((msg) => msg.role === "user");
	if (!lastUserMessage) return "";

	const chunks: string[] = [];
	for (const part of lastUserMessage.parts) {
		if (part.type === "text") chunks.push(part.text);
		if (part.type === "file") {
			chunks.push(
				await saveFilePart(part, accountId, threadId ?? lastUserMessage.id),
			);
		}
	}

	return chunks.filter(Boolean).join("\n\n").trim();
}

async function ensureSessionDirectories(accountId: string, sessionId: string) {
	await Promise.all([
		mkdir(accountSessionRoot(accountId, sessionId), { recursive: true }),
		mkdir(accountSessionInputRoot(accountId, sessionId), { recursive: true }),
		mkdir(accountSessionOutputRoot(accountId, sessionId), { recursive: true }),
		mkdir(accountSessionScratchpadRoot(accountId, sessionId), {
			recursive: true,
		}),
	]);
}

function renderSessionRuntimeContext(accountId: string, sessionId: string) {
	return [
		"[当前 session 工作区]",
		`- session_root: ${accountSessionRoot(accountId, sessionId)}`,
		`- 输入目录: ${accountSessionInputRoot(accountId, sessionId)}`,
		`- 输出目录: ${accountSessionOutputRoot(accountId, sessionId)}`,
		`- 临时目录: ${accountSessionScratchpadRoot(accountId, sessionId)}`,
		"当前 session 目录是默认可写工作区。",
		"交付产物必须优先写入输出目录。",
		"临时脚本、缓存和中间结果必须写入临时目录。",
		"可以读取完成任务所需的外部资料，但不要修改平台服务代码目录。",
	].join("\n\n");
}

function getLastSessionId(messages: UIMessage[]) {
	for (const message of [...messages].reverse()) {
		if (message.role !== "assistant") continue;
		const metadata = message.metadata;
		if (metadata && typeof metadata === "object" && "sessionId" in metadata) {
			const sessionId = (metadata as { sessionId?: unknown }).sessionId;
			if (typeof sessionId === "string" && sessionId) return sessionId;
		}
	}
	return undefined;
}

async function saveFilePart(
	part: Extract<UIMessage["parts"][number], { type: "file" }>,
	accountId: string,
	threadId: string,
) {
	const filename = safeFilename(part.filename ?? "attachment");
	if (!part.url.startsWith("data:")) {
		return `[文件] ${filename}\n- 类型: ${part.mediaType}\n- URL: ${part.url}`;
	}

	const match = part.url.match(/^data:([^;,]+)?(;base64)?,(.*)$/);
	if (!match) {
		return `[文件] ${filename}\n- 类型: ${part.mediaType}\n- 状态: 无法解析 data URL`;
	}

	const isBase64 = Boolean(match[2]);
	const body = match[3] ?? "";
	const bytes = isBase64
		? Buffer.from(body, "base64")
		: Buffer.from(decodeURIComponent(body), "utf8");
	const uploadDir = accountSessionInputRoot(accountId, safeFilename(threadId));
	await mkdir(uploadDir, { recursive: true });
	const filePath = path.join(uploadDir, `${Date.now()}-${filename}`);
	await writeFile(filePath, bytes);

	return [
		`[文件已上传] ${filename}`,
		`- 类型: ${part.mediaType}`,
		`- 本地路径: ${filePath}`,
		"请在需要读取文件内容时使用 read_file 工具读取该本地路径。",
	].join("\n");
}

function safeFilename(value: string) {
	return value.replace(/[^a-zA-Z0-9._-]+/g, "_").slice(0, 120) || "attachment";
}

async function callClawBackendStream(
	prompt: string,
	runtimeContext: string,
	accountId: string,
	sessionId: string,
	resumeSessionId?: string,
	writer?: UIMessageStreamWriter<UIMessage>,
	streamState?: StreamState,
): Promise<ClawChatResponse> {
	try {
		const response = await fetch(`${CLAW_API_URL}/api/chat/stream`, {
			method: "POST",
			headers: { "content-type": "application/json" },
			body: JSON.stringify({
				prompt,
				runtime_context: runtimeContext,
				account_id: accountId,
				session_id: sessionId,
				...(resumeSessionId ? { resume_session_id: resumeSessionId } : {}),
			}),
		});
		if (!response.ok) {
			const payload = (await response
				.json()
				.catch(() => ({}))) as ClawChatResponse;
			return {
				error:
					payload.detail ??
					payload.error ??
					`Claw backend returned ${response.status}`,
			};
		}
		if (!response.body)
			return { error: "Claw backend returned an empty stream" };
		return await consumeClawStream(response.body, writer, streamState);
	} catch (err) {
		return {
			error:
				err instanceof Error
					? `Unable to reach Claw backend: ${err.message}`
					: "Unable to reach Claw backend",
		};
	}
}

async function consumeClawStream(
	body: ReadableStream<Uint8Array>,
	writer?: UIMessageStreamWriter<UIMessage>,
	streamState?: StreamState,
) {
	const reader = body.getReader();
	const decoder = new TextDecoder();
	let buffer = "";
	let finalPayload: ClawChatResponse | undefined;
	while (true) {
		const { value, done } = await reader.read();
		if (value) {
			buffer += decoder.decode(value, { stream: !done });
			const lines = buffer.split("\n");
			buffer = lines.pop() ?? "";
			for (const line of lines) {
				const payload = parseStreamLine(line);
				if (!payload) continue;
				if (payload.kind === "event" && payload.event && writer) {
					writeRuntimeEvent(writer, payload.event, streamState);
				}
				if (payload.kind === "result") {
					finalPayload = payload.payload;
				}
				if (payload.kind === "error") {
					return {
						error:
							payload.detail ??
							payload.error ??
							`Claw backend returned ${payload.status_code ?? 500}`,
					};
				}
			}
		}
		if (done) break;
	}
	if (buffer.trim()) {
		const payload = parseStreamLine(buffer);
		if (payload?.kind === "result") finalPayload = payload.payload;
		if (payload?.kind === "error") {
			return {
				error:
					payload.detail ??
					payload.error ??
					`Claw backend returned ${payload.status_code ?? 500}`,
			};
		}
	}
	return (
		finalPayload ?? { error: "Claw backend stream ended without a result" }
	);
}

function parseStreamLine(line: string): ClawStreamItem | undefined {
	const trimmed = line.trim();
	if (!trimmed) return undefined;
	try {
		return JSON.parse(trimmed) as ClawStreamItem;
	} catch {
		return undefined;
	}
}

function writeRuntimeEvent(
	writer: UIMessageStreamWriter<UIMessage>,
	event: ClawRuntimeEvent,
	streamState?: StreamState,
) {
	if (event.type === "run_queued") {
		if (streamState) streamState.phase = "queued";
		if (streamState && event.run_id) streamState.runId = event.run_id;
		writer.write({
			type: "reasoning-delta",
			id: "reasoning-1",
			delta: "\n当前会话已有任务在执行，本轮正在排队。",
		});
	}
	if (event.type === "run_started") {
		if (streamState) streamState.phase = "running";
		if (streamState && event.run_id) streamState.runId = event.run_id;
		writer.write({
			type: "reasoning-delta",
			id: "reasoning-1",
			delta: "\n后端已开始执行本轮任务。",
		});
	}
	if (event.type === "server_heartbeat") {
		const elapsed =
			typeof event.elapsed_ms === "number"
				? `，已等待 ${formatDuration(event.elapsed_ms)}`
				: "";
		const phase = streamState?.phase === "queued" ? "仍在排队" : "后端仍在处理";
		writer.write({
			type: "reasoning-delta",
			id: "reasoning-1",
			delta: `\n${phase}${elapsed}。`,
		});
	}
	if (event.type === "content_delta" && event.delta) {
		writer.write({
			type: "reasoning-delta",
			id: "reasoning-1",
			delta: event.delta,
		});
	}
	if (event.type === "tool_call_delta") {
		const key =
			event.tool_call_id ??
			(event.tool_call_index === null || event.tool_call_index === undefined
				? undefined
				: `index:${event.tool_call_index}`);
		if (key && !streamState?.announcedToolCallKeys.has(key)) {
			streamState?.announcedToolCallKeys.add(key);
			writer.write({
				type: "reasoning-delta",
				id: "reasoning-1",
				delta: `\n准备调用 ${event.tool_name ?? "工具"}。`,
			});
		}
	}
	if (event.type === "final_text_start") {
		if (!streamState?.textStarted) {
			writer.write({ type: "text-start", id: "text-1" });
			if (streamState) streamState.textStarted = true;
		}
	}
	if (event.type === "final_text_delta" && event.delta) {
		if (!streamState?.textStarted) {
			writer.write({ type: "text-start", id: "text-1" });
			if (streamState) streamState.textStarted = true;
		}
		writer.write({ type: "text-delta", id: "text-1", delta: event.delta });
		if (streamState) streamState.textStreamed = true;
	}
	if (event.type === "final_text_end") {
		if (streamState?.textStarted) {
			writer.write({ type: "text-end", id: "text-1" });
			streamState.textStarted = false;
		}
	}
	if (event.type === "tool_start" && event.tool_call_id && event.tool_name) {
		streamState?.streamedToolCallIds.add(event.tool_call_id);
		writer.write({
			type: "tool-input-available",
			toolCallId: event.tool_call_id,
			toolName: event.tool_name,
			input: attachStageNote(event.arguments ?? {}, event.assistant_content),
		});
	}
	if (event.type === "tool_delta" && event.tool_call_id && event.delta) {
		writer.write({
			type: "reasoning-delta",
			id: "reasoning-1",
			delta: `\n${event.tool_name ?? "tool"} 输出中...`,
		});
	}
	if (event.type === "tool_result" && event.tool_call_id) {
		const preview = event.metadata?.output_preview;
		writer.write({
			type: "tool-output-available",
			toolCallId: event.tool_call_id,
			output:
				typeof preview === "string"
					? preview
					: event.ok === false
						? "工具执行失败，等待最终结果汇总。"
						: "工具调用完成，等待最终结果汇总。",
		});
	}
}

function formatClawResponse(payload: ClawChatResponse) {
	if (payload.error) return payload.error;
	return payload.final_output ?? "";
}

function safeSessionId(value: string) {
	return (
		value.replace(/[^a-zA-Z0-9._-]+/g, "_").slice(0, 120) || crypto.randomUUID()
	);
}
