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

import { complete, type Message, type Tool } from "@mariozechner/pi-ai";
import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";

const SYSTEM_PROMPT = `You're helping transfer context between coding sessions. The next session starts fresh with no memory of this conversation, so extract what matters.

Consider these questions:
- What was just done or implemented?
- What decisions were made and why?
- What technical details were discovered (APIs, methods, patterns)?
- What constraints, limitations, or caveats were found?
- What patterns or approaches are being followed?
- What is still unfinished or unresolved?
- What open questions or risks are worth flagging?

Rules:
- Be concrete: prefer file paths, specific decisions, and actual commands over vague summaries.
- Only include files from the provided candidate list.
- Only include skills from the provided loaded skills list.
- If something wasn't explicitly discussed, don't include it.`;

/**
 * Tool definition passed to the extraction LLM call. Models produce
 * structured tool-call arguments far more reliably than free-form JSON
 * in a text response, so we use a single tool instead of responseFormat.
 */
const HANDOFF_EXTRACTION_TOOL: Tool = {
	name: "extract_handoff_context",
	description: "Extract handoff context from the conversation for a new session.",
	parameters: Type.Object({
		relevantFiles: Type.Array(Type.String(), {
			description: "File paths relevant to continuing this work, chosen from the provided candidate list",
		}),
		skillsInUse: Type.Array(Type.String(), {
			description: "Skills relevant to the goal, chosen from the provided list of loaded skills",
		}),
		context: Type.String({
			description:
				"Key context: what was done, decisions made, technical details discovered, constraints, and patterns being followed",
		}),
		openItems: Type.Array(Type.String(), {
			description: "Unfinished work, open questions, known risks, and things to watch out for",
		}),
	}),
};

export type HandoffExtraction = {
	files: string[];
	skills: string[];
	context: string;
	openItems: string[];
};

/**
 * Run the extraction LLM call and return structured context, or null if
 * aborted or the model didn't produce a valid tool call.
 */
export async function extractHandoffContext(
	model: NonNullable<ExtensionContext["model"]>,
	apiKey: string | undefined,
	headers: Record<string, string> | undefined,
	conversationText: string,
	goal: string,
	candidateFiles: string[],
	loadedSkills: string[],
	signal?: AbortSignal,
): Promise<HandoffExtraction | null> {
	const filesContext =
		candidateFiles.length > 0
			? `\n\n## Candidate Files\n\nThese files were touched or mentioned during the session. Return only the ones relevant to the goal.\n\n${candidateFiles.map((f) => `- ${f}`).join("\n")}`
			: "";

	const skillsContext =
		loadedSkills.length > 0
			? `\n\n## Skills Loaded During This Session\n\n${loadedSkills.map((s) => `- ${s}`).join("\n")}\n\nReturn only the skills from this list that are relevant to the goal.`
			: "";

	const userMessage: Message = {
		role: "user",
		content: [
			{
				type: "text",
				text: `## Conversation\n\n${conversationText}\n\n## Goal for Next Session\n\n${goal}${filesContext}${skillsContext}`,
			},
		],
		timestamp: Date.now(),
	};

	const response = await complete(
		model,
		{
			systemPrompt: SYSTEM_PROMPT,
			messages: [userMessage],
			tools: [HANDOFF_EXTRACTION_TOOL],
		},
		{ apiKey, headers, signal },
	);

	if (response.stopReason === "aborted") return null;

	const toolCall = response.content.find((c) => c.type === "toolCall" && c.name === "extract_handoff_context");

	if (!toolCall || toolCall.type !== "toolCall") {
		console.error("Model did not call extract_handoff_context:", response.content);
		return null;
	}

	const args = toolCall.arguments as Record<string, unknown>;

	if (
		!Array.isArray(args.relevantFiles) ||
		!Array.isArray(args.skillsInUse) ||
		typeof args.context !== "string" ||
		!Array.isArray(args.openItems)
	) {
		console.error("Unexpected tool call arguments shape:", args);
		return null;
	}

	const candidateSet = new Set(candidateFiles);
	const loadedSkillSet = new Set(loadedSkills);

	return {
		files: (args.relevantFiles as string[]).filter((f) => typeof f === "string" && candidateSet.has(f)),
		skills: (args.skillsInUse as string[]).filter((s) => typeof s === "string" && loadedSkillSet.has(s)),
		context: args.context as string,
		openItems: (args.openItems as string[]).filter((item) => typeof item === "string"),
	};
}

/**
 * Assemble the handoff draft. The user's goal goes last so it has the
 * most weight in the new session (recency bias).
 */
export function assembleHandoffDraft(
	result: HandoffExtraction,
	goal: string,
	parentSessionFile: string | undefined,
): string {
	const parentBlock = parentSessionFile
		? `**Parent session:** \`${parentSessionFile}\`\n\nUse the \`session_query\` tool with this path if you need details from the prior thread.\n\n`
		: "";

	const filesSection = result.files.length > 0 ? `## Files\n\n${result.files.map((f) => `- ${f}`).join("\n")}\n\n` : "";

	const skillsSection =
		result.skills.length > 0 ? `## Skills in Use\n\n${result.skills.map((s) => `- ${s}`).join("\n")}\n\n` : "";

	const contextSection = `## Context\n\n${result.context}\n\n`;

	const openItemsSection =
		result.openItems.length > 0 ? `## Open Items\n\n${result.openItems.map((item) => `- ${item}`).join("\n")}\n\n` : "";

	return `${parentBlock}${filesSection}${skillsSection}${contextSection}${openItemsSection}${goal}`.trim();
}