/**
 * research-web.ts
 * Pi-Extension: Web-Recherche via Perplexity Sonar
 *
 * Platzieren in: ~/.pi/agent/extensions/research-web.ts
 * Nach Änderungen in Pi: /reload
 *
 * Kostenstruktur (Stand April 2026):
 *   sonar:     1 USD/M Input + 1 USD/M Output + 0.005 USD/web_search
 *   sonar-pro: 3 USD/M Input + 15 USD/M Output + 0.005 USD/web_search
 
   Übersicht aller drei Contexts: Zusatz im Prompt: "Nutze context=<context> .                  
                                                                                                                                           
  ┌─────────┬────────────────┬──────────────────────────────────────────────────┬───────────────────────────────────┐
  │ context │ Modell-Default │                      Format                      │            Optimal für            │                      
  ├─────────┼────────────────┼──────────────────────────────────────────────────┼───────────────────────────────────┤
  │ facts   │ sonar/fast     │ Verbatim + [N] inline                            │ Presse, Faktencheck (default)     │                      
  ├─────────┼────────────────┼──────────────────────────────────────────────────┼───────────────────────────────────┤
  │ code    │ sonar/fast     │ Snippet unter jeder Quelle                       │ APIs, Libraries, Doku             │                      
  ├─────────┼────────────────┼──────────────────────────────────────────────────┼───────────────────────────────────┤                      
  │ legal   │ sonar-pro/deep │ Verbatim + [N] + Paragraf + Snippet + Disclaimer │ Gesetze, DSGVO, EU-Recht, Urteile │                      
  └─────────┴────────────────┴──────────────────────────────────────────────────┴───────────────────────────────────┘                      
                                                            
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";

const PARAMS = Type.Object({
	query: Type.String({ description: "Die Recherchefrage auf Deutsch oder Englisch" }),
	recency: Type.Optional(
		Type.String({
			description: "Aktualitätsfilter: day, week, month oder year. Weglassen = kein Filter",
		})
	),
	allowWikipedia: Type.Optional(
		Type.Boolean({
			description:
				"Wikipedia als Quelle erlauben. Nur setzen wenn der User das explizit anfordert. Standard: false",
		})
	),
	mode: Type.Optional(
		Type.String({
			description:
				"fast (Standard): sonar, kostengünstig, für die meisten Anfragen. " +
				"deep: sonar-pro, für komplexe, mehrstufige oder heikle Recherchen.",
		})
	),
	context: Type.Optional(
		Type.String({
			description:
				"facts (Standard): Presse/Fakten-Recherche — Inline-Zitierungen [1][2][3] im Text, verbatim ausgeben. " +
				"code: Programmier-Dokumentation — Quellen mit relevantem Textauszug (Snippet). " +
				"legal: Gesetze/Urteile/Verordnungen — offizielle Quellen, Paragrafangaben, Gesetzestext als Snippet. Nutzt automatisch mode=deep.",
		})
	),
});

type ResearchSource = { url: string; title?: string; snippet?: string };

type PerplexityResponse = {
	model?: string;
	citations?: string[];
	search_results?: Array<{ url?: string; title?: string; snippet?: string }>;
	choices?: Array<{ message?: { content?: string } }>;
	usage?: {
		prompt_tokens?: number;
		completion_tokens?: number;
		total_tokens?: number;
		search_queries?: number;
	};
};

// Kostenmodell (USD pro 1 Mio Tokens)
const PRICING = {
	sonar: { inputPerM: 1, outputPerM: 1 },
	"sonar-pro": { inputPerM: 3, outputPerM: 15 },
} as const;
const SEARCH_COST_PER_CALL = 0.005;

class RetryableError extends Error {
	constructor(message: string) {
		super(message);
		this.name = "RetryableError";
	}
}

function stripQuotes(s: string): string {
	return s.replace(/^["']+|["']+$/g, "").trim();
}

function normalizeMode(raw: string | undefined, fallback: "fast" | "deep"): "fast" | "deep" {
	if (!raw) return fallback;
	const v = stripQuotes(raw).toLowerCase();
	return v === "deep" ? "deep" : "fast";
}

function normalizeContext(raw: string | undefined): "facts" | "code" | "legal" {
	if (!raw) return "facts";
	const v = stripQuotes(raw).toLowerCase();
	if (v === "code") return "code";
	if (v === "legal") return "legal";
	return "facts";
}

function normalizeRecency(raw: string | undefined): "day" | "week" | "month" | "year" | undefined {
	if (!raw) return undefined;
	const v = stripQuotes(raw).toLowerCase() as "day" | "week" | "month" | "year";
	return ["day", "week", "month", "year"].includes(v) ? v : undefined;
}

async function sleep(ms: number) {
	await new Promise((r) => setTimeout(r, ms));
}

function estimateCostUSD(
	model: string,
	promptTokens: number,
	completionTokens: number,
	searchQueries: number
): number {
	const p = PRICING[model as keyof typeof PRICING] ?? PRICING["sonar"];
	const tokenCost =
		(promptTokens / 1_000_000) * p.inputPerM + (completionTokens / 1_000_000) * p.outputPerM;
	const searchCost = searchQueries * SEARCH_COST_PER_CALL;
	return tokenCost + searchCost;
}

function buildSystemPrompt(context: "facts" | "code" | "legal"): string {
	if (context === "code") {
		return (
			"Du bist ein Recherche-Tool für Softwareentwickler. " +
			"Antworte knapp und präzise auf Deutsch oder Englisch. " +
			"Fokussiere auf die technisch relevante Information und benenne, welche Dokumentation die Frage beantwortet."
		);
	}
	if (context === "legal") {
		return (
			"Du bist ein juristisches Recherche-Tool. " +
			"Antworte sachlich und präzise auf Deutsch. " +
			"Fokussiere ausschließlich auf offizielle Rechtsquellen: Gesetze, Verordnungen, EU-Recht, amtliche Bekanntmachungen und Gerichtsurteile. " +
			"Nenne Gesetz und Paragraf direkt im Text (z.B. § 13 DSGVO, Art. 5 GG). " +
			"Setze Inline-Zitierungen [1][2][3] hinter jeden Satz mit Rechtsgrundlage. " +
			"Weise ausdrücklich auf relevante Rechtsänderungen oder Ausnahmen hin. " +
			"Keine Rechtsberatung — nur Darstellung der Rechtslage."
		);
	}
	return (
		"Du bist ein Recherche-Tool für einen KI-Agenten. " +
		"Antworte sachlich und prägnant auf Deutsch. " +
		"Setze Inline-Zitierungen [1][2][3] direkt hinter jeden Satz, den du auf Quellen stützt. " +
		"Begrenze die Antwort auf das Wesentliche."
	);
}

async function callPerplexity(
	query: string,
	recency: string | undefined,
	allowWikipedia: boolean | undefined,
	context: "facts" | "code" | "legal",
	model: "sonar" | "sonar-pro",
	contextSize: "low" | "high",
	maxTokens: number,
	apiKey: string,
	signal: AbortSignal
): Promise<PerplexityResponse> {
	const body: Record<string, unknown> = {
		model,
		messages: [
			{
				role: "system",
				content: buildSystemPrompt(context),
			},
			{
				role: "user",
				content: `Recherchefrage: ${query}\n\nKurze, neutrale Zusammenfassung der wichtigsten Fakten.`,
			},
		],
		max_tokens: maxTokens,
		temperature: 0.2,
		web_search_options: {
			search_context_size: contextSize,
		},
	};

	if (recency) {
		body.search_recency_filter = recency;
	}

	// Wikipedia nur auf explizite Anforderung — standardmäßig ausschließen
	if (!allowWikipedia) {
		body.search_domain_filter = [
			"-en.wikipedia.org",
			"-de.wikipedia.org",
			"-simple.wikipedia.org",
		];
	}

	const resp = await fetch("https://api.perplexity.ai/chat/completions", {
		method: "POST",
		headers: {
			Authorization: `Bearer ${apiKey}`,
			"Content-Type": "application/json",
		},
		body: JSON.stringify(body),
		signal,
	});

	if (!resp.ok) {
		const text = await resp.text().catch(() => "");
		// Nur bei transienten Fehlern retryen — 4xx-Konfigurationsfehler nicht wiederholen
		if (resp.status === 429 || resp.status >= 500) {
			throw new RetryableError(`Perplexity API Fehler ${resp.status}: ${text}`);
		}
		throw new Error(`Perplexity API Fehler ${resp.status}: ${text}`);
	}

	return (await resp.json()) as PerplexityResponse;
}

function dedupeSources(sources: ResearchSource[]): ResearchSource[] {
	const seen = new Set<string>();
	const out: ResearchSource[] = [];
	for (const s of sources) {
		if (!s.url || seen.has(s.url)) continue;
		seen.add(s.url);
		out.push(s);
	}
	return out;
}

function parseSources(data: PerplexityResponse, withSnippets: boolean, maxSources = 8): ResearchSource[] {
	const fromSearchResults =
		data.search_results
			?.filter((r) => !!r?.url)
			.map((r) => ({
				url: r.url!,
				title: r.title,
				snippet: withSnippets ? r.snippet : undefined,
			})) ?? [];

	if (fromSearchResults.length > 0) {
		return dedupeSources(fromSearchResults).slice(0, maxSources);
	}

	return dedupeSources(
		data.citations
			?.filter((u) => typeof u === "string" && /^https?:\/\//.test(u))
			.map((url) => ({ url })) ?? []
	).slice(0, maxSources);
}

function formatCostLine(
	model: string,
	mode: string,
	promptTokens: number,
	completionTokens: number,
	estimatedCostUSD: number
): string {
	const tokenInfo =
		promptTokens || completionTokens ? ` · ${promptTokens}+${completionTokens} Tokens` : "";
	return `_[Perplexity: ${model}/${mode}${tokenInfo} · ~$${estimatedCostUSD.toFixed(4)}]_`;
}

function formatSourcesFacts(sources: ResearchSource[]): string[] {
	return sources.map((s, i) => {
		const num = `[${i + 1}]`;
		return s.title ? `${num} [${s.title}](${s.url})` : `${num} ${s.url}`;
	});
}

function formatSourcesWithSnippets(sources: ResearchSource[], snippetLen = 200): string[] {
	const lines: string[] = [];
	sources.forEach((s, i) => {
		const num = `[${i + 1}]`;
		const titleLine = s.title ? `${num} [${s.title}](${s.url})` : `${num} ${s.url}`;
		lines.push(titleLine);
		if (s.snippet) {
			const trimmed = s.snippet.length > snippetLen
				? s.snippet.slice(0, snippetLen - 3) + "…"
				: s.snippet;
			lines.push(`    > "${trimmed}"`);
		}
		lines.push("");
	});
	return lines;
}

function formatResult(
	summary: string,
	sources: ResearchSource[],
	context: "facts" | "code" | "legal",
	costLine: string
): string {
	const lines: string[] = [];

	if (context === "facts" || context === "legal") {
		// Verbatim-Wrapper: Agent soll diesen Block unverändert ausgeben
		const label = context === "legal"
			? "⚠️ RECHTLICHE RECHERCHE — Bitte unverändert und vollständig ausgeben (keine Rechtsberatung):\n"
			: "⚠️ RECHERCHE-ERGEBNIS — Bitte unverändert und vollständig ausgeben:\n";
		lines.push(label);
		lines.push(summary);
	} else {
		lines.push(summary);
	}

	if (sources.length > 0) {
		lines.push("\n**Quellen:**");
		if (context === "code") {
			lines.push(...formatSourcesWithSnippets(sources, 200));
		} else if (context === "legal") {
			// Legal: längere Snippets (300 Z.) um Gesetzestext lesbar zu machen
			lines.push(...formatSourcesWithSnippets(sources, 300));
		} else {
			lines.push(...formatSourcesFacts(sources));
		}
	} else {
		lines.push("\n_(Keine Quellen im Response enthalten)_");
	}

	lines.push(`\n${costLine}`);

	return lines.join("\n");
}

export default function researchWebExtension(pi: ExtensionAPI) {
	pi.registerTool({
		name: "research_web",
		label: "Web-Recherche",
		description:
			"Recherchiert live im Internet via Perplexity Sonar. " +
			"Nutze dieses Tool wenn: aktuelle Fakten, Preise, Versionen, News, Gesetze, " +
			"Produktänderungen oder unsichere/veraltete Informationen gefragt sind. " +
			"Nicht nutzen für reine Logik, Mathematik oder stabile Grundlagen.",
		promptGuidelines: [
			"Use research_web when the question requires up-to-date or potentially outdated information.",
			"Do NOT use research_web for stable knowledge, math, or logic questions.",
			"Default to mode=fast. Use mode=deep only for complex multi-part questions, legal/medical/financial topics, or when fast results are clearly insufficient.",
			"Default to context=facts. Use context=code for programming/API/documentation questions. Use context=legal for questions about laws, regulations, EU law, court rulings, or compliance.",
			"For context=facts: Output the research result VERBATIM and COMPLETE — do not rewrite, paraphrase, or restructure. Preserve all [N] inline citation markers exactly where they appear in the text.",
			"For context=code: Show each source with its indented snippet (> \"...\") exactly as provided — this saves the user from opening each URL.",
			"For context=legal: Output the result VERBATIM. Always add a disclaimer that this is not legal advice. Preserve paragraph references (§ 13 DSGVO etc.) and [N] inline citations exactly.",
			"Always include the complete numbered **Quellen:** section verbatim from the tool result, with all [N] numbers and URLs as clickable links.",
			"Always include the cost line (the italicized [Perplexity: ...] line) verbatim at the end of your response.",
			"If research_web returns no sources, flag the answer as potentially uncertain.",
			"Set allowWikipedia=true ONLY if the user explicitly asks to use Wikipedia as a source.",
		],
		parameters: PARAMS,
		async execute(_toolCallId, params, signal) {
			const apiKey = process.env.PERPLEXITY_API_KEY;
			if (!apiKey) {
				return {
					content: [{ type: "text", text: "Fehler: PERPLEXITY_API_KEY ist nicht gesetzt." }],
				};
			}

			const context = normalizeContext(params.context);
			// legal braucht immer deep — Gesetze erfordern gründliche Recherche
			const mode = normalizeMode(params.mode, context === "legal" ? "deep" : "fast");
			const recency = normalizeRecency(params.recency);
			const model: "sonar" | "sonar-pro" = mode === "deep" ? "sonar-pro" : "sonar";
			const contextSize: "low" | "high" = mode === "deep" ? "high" : "low";
			const maxTokens = mode === "deep" ? 600 : 300;

			const maxAttempts = 3;
			let lastError: unknown;

			for (let attempt = 1; attempt <= maxAttempts; attempt++) {
				try {
					const data = await callPerplexity(
						params.query,
						recency,
						params.allowWikipedia,
						context,
						model,
						contextSize,
						maxTokens,
						apiKey,
						signal
					);
					const summary = data.choices?.[0]?.message?.content?.trim() ?? "";

					if (!summary) {
						throw new Error("Leere Antwort von Perplexity erhalten");
					}

					const sources = parseSources(data, context === "code" || context === "legal");

					const usage = data.usage ?? {};
					const promptTokens = usage.prompt_tokens ?? 0;
					const completionTokens = usage.completion_tokens ?? 0;
					const searchQueries = usage.search_queries ?? 0;
					const finalModel = data.model ?? model;
					const estimatedCostUSD = Number(
						estimateCostUSD(finalModel, promptTokens, completionTokens, searchQueries).toFixed(6)
					);

					const costLine = formatCostLine(finalModel, mode, promptTokens, completionTokens, estimatedCostUSD);
					const text = formatResult(summary, sources, context, costLine);

					return {
						content: [{ type: "text", text }],
						details: {
							model: finalModel,
							mode,
							context,
							sourceCount: sources.length,
							promptTokens: promptTokens || null,
							completionTokens: completionTokens || null,
							totalTokens: usage.total_tokens ?? null,
							searchQueries: searchQueries || null,
							estimatedCostUSD,
						},
					};
				} catch (err) {
					lastError = err;
					// Nur bei transienten Fehlern warten und nochmals versuchen
					if (err instanceof RetryableError && attempt < maxAttempts) {
						await sleep(400 * 2 ** (attempt - 1));
					} else {
						break;
					}
				}
			}

			const msg = lastError instanceof Error ? lastError.message : "Unbekannter Fehler";
			return { content: [{ type: "text", text: `Recherchefehler: ${msg}` }] };
		},
	});
}