// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
// SPDX-FileCopyrightText: Armin Ronacher <armin.ronacher@active-4.com>
//
// SPDX-License-Identifier: Apache-2.0

import type { complete } from "@mariozechner/pi-ai";
import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
import type { ExtractedQuestion, ExtractionResult } from "./prompt.js";

// Preferred model for extraction — lightweight is fine for structured JSON output.
// Falls back to the session model if this one isn't available.
const PREFERRED_EXTRACTION_MODEL_ID = "nemotron-3-super-120b-a12b";

/**
 * Resolve the extraction model: prefer a lightweight model from the registry,
 * fall back to the current session model.
 */
export function resolveExtractionModel(ctx: ExtensionContext, fallback: NonNullable<ExtensionContext["model"]>) {
	const available = ctx.modelRegistry.getAvailable();
	const preferred = available.find((m) => m.id === PREFERRED_EXTRACTION_MODEL_ID);
	return preferred ?? fallback;
}

/**
 * Parse a tool call response into an ExtractionResult.
 */
export function parseToolCallResult(response: Awaited<ReturnType<typeof complete>>): ExtractionResult | null {
	const toolCall = response.content.find((c) => c.type === "toolCall" && c.name === "extract_questions");

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

	const args = toolCall.arguments as Record<string, unknown>;
	if (!Array.isArray(args.questions)) {
		console.error("[answer] expected questions array, got:", typeof args.questions, args);
		return null;
	}

	// Validate each question item — model may return plain strings instead
	// of objects if it ignores the schema, so we handle both gracefully.
	const questions: ExtractedQuestion[] = [];
	for (const item of args.questions) {
		if (typeof item === "string") {
			// Model returned bare string instead of object — use as question text
			questions.push({ question: item });
			continue;
		}
		if (typeof item !== "object" || item === null) continue;
		const obj = item as Record<string, unknown>;
		if (typeof obj.question !== "string") {
			console.error("[answer] skipping item with no 'question' string:", Object.keys(obj));
			continue;
		}
		questions.push({
			question: obj.question,
			context: typeof obj.context === "string" ? obj.context : undefined,
		});
	}

	return { questions };
}

/** Escape characters that would break pseudo-XML tag boundaries. */
export function escapeXml(s: string): string {
	return s.replace(/&/g, "&amp;").replace(/</g, "&lt;").replace(/>/g, "&gt;");
}