// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
// SPDX-FileCopyrightText: Petr Baudis <pasky@ucw.cz>
//
// SPDX-License-Identifier: MIT

import { complete, type Message } from "@mariozechner/pi-ai";
import type { ExtensionAPI, SessionEntry } from "@mariozechner/pi-coding-agent";
import { SessionManager, convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";
import * as fs from "node:fs";
import { getFallbackSessionsRoot, getSessionsRoot, normalizeSessionPath, sessionPathAllowed } from "./session-paths.js";

const QUERY_SYSTEM_PROMPT = `You answer questions about a prior pi session.

Rules:
- Use only facts from the provided conversation.
- Prefer concrete outputs: file paths, decisions, TODOs, errors.
- If not present, say explicitly: "Not found in provided session.".
- Keep answer concise.`;

export function registerSessionQueryTool(pi: ExtensionAPI) {
	pi.registerTool({
		name: "session_query",
		label: "Session Query",
		description:
			"Query a prior pi session file. Use when handoff prompt references a parent session and you need details.",
		parameters: Type.Object({
			sessionPath: Type.String({
				description:
					"Session .jsonl path. Absolute path, or relative to sessions root (e.g. 2026-02-16/foo/session.jsonl)",
			}),
			question: Type.String({ description: "Question about that session" }),
		}),
		async execute(_toolCallId, params, signal, onUpdate, ctx) {
			const currentSessionFile = ctx.sessionManager.getSessionFile();
			const sessionsRoot = getSessionsRoot(currentSessionFile) ?? getFallbackSessionsRoot();
			const resolvedPath = normalizeSessionPath(params.sessionPath, sessionsRoot);

			const error = (text: string) => ({
				content: [{ type: "text" as const, text }],
				details: { error: true } as const,
			});

			const cancelled = () => ({
				content: [{ type: "text" as const, text: "Session query cancelled." }],
				details: { cancelled: true } as const,
			});

			if (signal?.aborted) {
				return cancelled();
			}

			if (!resolvedPath.endsWith(".jsonl")) {
				return error(`Invalid session path (expected .jsonl): ${params.sessionPath}`);
			}

			if (!sessionPathAllowed(resolvedPath, sessionsRoot)) {
				return error(`Session path outside allowed sessions directory: ${params.sessionPath}`);
			}

			if (!fs.existsSync(resolvedPath)) {
				return error(`Session file not found: ${resolvedPath}`);
			}

			let fileStats: fs.Stats;
			try {
				fileStats = fs.statSync(resolvedPath);
			} catch (err) {
				return error(`Failed to stat session file: ${String(err)}`);
			}

			if (!fileStats.isFile()) {
				return error(`Session path is not a file: ${resolvedPath}`);
			}

			onUpdate?.({
				content: [{ type: "text", text: `Querying: ${resolvedPath}` }],
				details: { status: "loading", sessionPath: resolvedPath },
			});

			let sessionManager: SessionManager;
			try {
				sessionManager = SessionManager.open(resolvedPath);
			} catch (err) {
				return error(`Failed to open session: ${String(err)}`);
			}

			const branch = sessionManager.getBranch();
			const messages = branch
				.filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
				.map((entry) => entry.message);

			if (messages.length === 0) {
				return {
					content: [{ type: "text" as const, text: "Session has no messages." }],
					details: { empty: true, sessionPath: resolvedPath },
				};
			}

			if (!ctx.model) {
				return error("No model selected for session query.");
			}

			const conversationText = serializeConversation(convertToLlm(messages));
			try {
				const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model);
				if (!auth.ok) {
					return error(`Failed to get API key: ${auth.error}`);
				}
				const userMessage: Message = {
					role: "user",
					content: [
						{
							type: "text",
							text: `## Session\n\n${conversationText}\n\n## Question\n\n${params.question}`,
						},
					],
					timestamp: Date.now(),
				};

				const response = await complete(
					ctx.model,
					{ systemPrompt: QUERY_SYSTEM_PROMPT, messages: [userMessage] },
					{ apiKey: auth.apiKey, headers: auth.headers, signal: signal as AbortSignal },
				);

				if (response.stopReason === "aborted") {
					return cancelled();
				}

				const answer = response.content
					.filter((c): c is { type: "text"; text: string } => c.type === "text")
					.map((c) => c.text)
					.join("\n")
					.trim();

				return {
					content: [{ type: "text" as const, text: answer || "No answer generated." }],
					details: {
						sessionPath: resolvedPath,
						question: params.question,
						messageCount: messages.length,
					},
				};
			} catch (err) {
				if (signal?.aborted) {
					return cancelled();
				}
				if (err instanceof Error && err.name === "AbortError") {
					return cancelled();
				}
				return error(`Session query failed: ${String(err)}`);
			}
		},
	});
}