/**
 * Context7 Extension for pi
 *
 * Adds two tools that let the LLM fetch up-to-date library documentation
 * from context7.com at query time, avoiding stale training data.
 *
 * Tools:
 *   - resolve_library_id  — search context7.com for a library by name
 *   - get_library_docs    — fetch documentation snippets for a resolved library ID
 *
 * Installation:
 *   Copy or symlink this directory into ~/.pi/agent/extensions/ or
 *   your project's .pi/extensions/, then restart pi (or /reload).
 *
 * Usage:
 *   Just ask pi about any library. If the LLM's training data might be stale
 *   (new versions, niche packages, etc.) it will automatically call these tools.
 *   You can also explicitly ask: "use context7 to look up the zustand docs".
 */

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

// ─── Constants ────────────────────────────────────────────────────────────────

const BASE_URL = "https://context7.com/api/v1";
const DEFAULT_TOKENS = 10_000;
const MAX_RESULTS = 10;

// ─── Types ────────────────────────────────────────────────────────────────────

interface SearchResult {
	id: string;
	title: string;
	description: string;
	totalTokens: number;
	totalSnippets: number;
	stars: number;
	trustScore: number;
	benchmarkScore: number;
	verified: boolean;
	versions: string[];
	score: number;
	vip: boolean;
}

interface SearchResponse {
	results: SearchResult[];
}

// ─── Helpers ──────────────────────────────────────────────────────────────────

async function searchLibraries(query: string, signal?: AbortSignal): Promise<SearchResponse> {
	const url = new URL(`${BASE_URL}/search`);
	url.searchParams.set("query", query);

	const response = await fetch(url.toString(), {
		headers: { "X-User-Agent": "pi-context7/1.0" },
		signal,
	});

	if (!response.ok) {
		const body = await response.text().catch(() => "");
		throw new Error(
			`Context7 search failed: ${response.status} ${response.statusText}${body ? ` — ${body}` : ""}`,
		);
	}

	return response.json() as Promise<SearchResponse>;
}

async function fetchLibraryDocs(
	libraryId: string,
	tokens: number,
	topic?: string,
	signal?: AbortSignal,
): Promise<string> {
	// Normalise the id: strip a leading "/" so callers can pass either form
	const normalised = libraryId.startsWith("/") ? libraryId.slice(1) : libraryId;
	const url = new URL(`${BASE_URL}/${normalised}`);
	url.searchParams.set("tokens", String(tokens));
	if (topic) url.searchParams.set("topic", topic);

	const response = await fetch(url.toString(), {
		headers: { "X-User-Agent": "pi-context7/1.0" },
		signal,
	});

	if (!response.ok) {
		const body = await response.text().catch(() => "");
		throw new Error(
			`Context7 docs fetch failed: ${response.status} ${response.statusText}${body ? ` — ${body}` : ""}`,
		);
	}

	return response.text();
}

/** Format a number with K/M suffix */
function fmt(n: number): string {
	if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
	if (n >= 1_000) return `${(n / 1_000).toFixed(0)}K`;
	return String(n);
}

// ─── System Prompt ────────────────────────────────────────────────────────────

/**
 * Appended to the system prompt every turn so the LLM knows exactly when and
 * how to use the context7 tools, mirroring the behaviour of the Context7 MCP server.
 */
const SYSTEM_PROMPT_SECTION = `
## Context7 — live library documentation

You have access to two tools that fetch real, up-to-date documentation from
context7.com at query time. Use them whenever the user's prompt contains the
phrase **"use context7"** or whenever you need current documentation for a
library, package, or framework (especially when your training data may be stale).

**Workflow (always follow this order):**
1. Call \`resolve_library_id\` with the library name to get a list of matching
   library IDs from context7.com.
2. Choose the best match (highest trustScore, prefer verified/featured entries).
3. Call \`get_library_docs\` with that ID (and an optional \`topic\`) to fetch
   real documentation snippets.
4. Use the returned documentation to answer the user's question accurately.

**When to use:**
- User prompt contains "use context7" → always fetch docs before answering.
- User asks about a specific library version, new API, or recent release.
- You are unsure whether your training data reflects the current API.

**Do not** answer library questions from memory when context7 tools are available
and the user has asked you to use context7.
`.trim();

// ─── Extension ────────────────────────────────────────────────────────────────

export default function context7Extension(pi: ExtensionAPI) {
	// ── System prompt injection ──────────────────────────────────────────────
	pi.on("before_agent_start", async (event) => {
		return {
			systemPrompt: event.systemPrompt + "\n\n" + SYSTEM_PROMPT_SECTION,
		};
	});

	// ── Input transform — make "use context7" explicit ───────────────────────
	// When the user writes "use context7" anywhere in their prompt we append a
	// one-line reminder so the LLM sees a clear, unambiguous instruction even if
	// the system prompt is long and the phrase is buried.
	pi.on("input", async (event) => {
		if (event.source === "extension") return { action: "continue" };
		if (/\buse context7\b/i.test(event.text)) {
			return {
				action: "transform",
				text:
					event.text +
					"\n\n[context7] Call resolve_library_id then get_library_docs before answering.",
			};
		}
		return { action: "continue" };
	});

	// ── resolve_library_id ───────────────────────────────────────────────────
	pi.registerTool({
		name: "resolve_library_id",
		label: "Context7: Search Library",
		description: `Search context7.com for a library by name and return matching library IDs.

Use this tool FIRST whenever you need current documentation for any library, package, or framework — especially when:
- The user asks about a specific library, SDK, or API
- Your training data might be outdated (new major versions, recent releases)
- You are unsure which exact version of a library is in use

Returns a ranked list of matches with IDs, descriptions, token counts, trust scores, and available versions. Pick the best match (highest trustScore + benchmarkScore, prefer verified) and pass its ID to \`get_library_docs\`.`,

		parameters: Type.Object({
			query: Type.String({
				description:
					"Library name or keyword to search for (e.g. 'react', 'nextjs', 'drizzle-orm')",
			}),
		}),

		async execute(_toolCallId, params, signal) {
			let data: SearchResponse;
			try {
				data = await searchLibraries(params.query, signal);
			} catch (err) {
				return {
					content: [{ type: "text", text: `Error: ${(err as Error).message}` }],
					isError: true,
					details: { error: (err as Error).message },
				};
			}

			if (!data.results || data.results.length === 0) {
				return {
					content: [
						{
							type: "text",
							text: `No libraries found for query: "${params.query}"\n\nTry a shorter or different search term.`,
						},
					],
					details: { results: [] },
				};
			}

			const top = data.results.slice(0, MAX_RESULTS);

			const lines: string[] = [
				`Found ${data.results.length} result(s) for "${params.query}" (showing top ${top.length}):`,
				"",
			];

			for (let i = 0; i < top.length; i++) {
				const r = top[i];
				const badges: string[] = [];
				if (r.verified) badges.push("✓ verified");
				if (r.vip) badges.push("★ featured");
				if (r.stars > 0) badges.push(`⭐ ${fmt(r.stars)}`);

				const versionNote =
					r.versions.length > 0
						? `  versions: ${r.versions.slice(0, 4).join(", ")}${r.versions.length > 4 ? " …" : ""}`
						: "";

				lines.push(
					`${i + 1}. ${r.title}`,
					`   id: ${r.id}`,
					`   ${r.description.slice(0, 120)}${r.description.length > 120 ? "…" : ""}`,
					`   trust: ${r.trustScore}/10  benchmark: ${r.benchmarkScore}  tokens: ${fmt(r.totalTokens)}  snippets: ${fmt(r.totalSnippets)}${badges.length ? `  ${badges.join("  ")}` : ""}`,
					...(versionNote ? [versionNote] : []),
					"",
				);
			}

			lines.push(
				'Next step: call `get_library_docs` with the chosen library ID (e.g. "/websites/react_dev").',
			);

			return {
				content: [{ type: "text", text: lines.join("\n") }],
				details: { results: top },
			};
		},
	});

	// ── get_library_docs ─────────────────────────────────────────────────────
	pi.registerTool({
		name: "get_library_docs",
		label: "Context7: Fetch Docs",
		description: `Fetch up-to-date documentation snippets for a library from context7.com.

Use this tool after calling \`resolve_library_id\` to obtain the library ID.
The response contains real documentation excerpts (code examples, API references, guides) pulled directly from the library's official docs — not reconstructed from training data.

Guidance:
- Set \`tokens\` higher (e.g. 15000–20000) for complex topics or when you need broad coverage.
- Set \`topic\` to a focused keyword (e.g. "authentication", "streaming", "hooks") to retrieve the most relevant snippets.
- The library ID looks like "/websites/react_dev" or "/vercel/next.js".`,

		parameters: Type.Object({
			libraryId: Type.String({
				description:
					'Context7 library ID returned by resolve_library_id (e.g. "/websites/react_dev")',
			}),
			tokens: Type.Optional(
				Type.Number({
					description: `Maximum number of tokens to return (default: ${DEFAULT_TOKENS}). Increase for broader coverage.`,
					minimum: 1000,
					maximum: 100_000,
				}),
			),
			topic: Type.Optional(
				Type.String({
					description:
						'Specific topic or keyword to focus the docs on (e.g. "routing", "server actions", "useEffect"). Omit for a general overview.',
				}),
			),
		}),

		async execute(_toolCallId, params, signal, onUpdate) {
			const tokens = params.tokens ?? DEFAULT_TOKENS;

			onUpdate?.({
				content: [
					{
						type: "text",
						text: `Fetching context7 docs for \`${params.libraryId}\`${params.topic ? ` · topic: "${params.topic}"` : ""} · up to ${fmt(tokens)} tokens…`,
					},
				],
			});

			let raw: string;
			try {
				raw = await fetchLibraryDocs(params.libraryId, tokens, params.topic, signal);
			} catch (err) {
				return {
					content: [{ type: "text", text: `Error: ${(err as Error).message}` }],
					isError: true,
					details: { error: (err as Error).message },
				};
			}

			if (!raw || raw.trim().length === 0) {
				return {
					content: [
						{
							type: "text",
							text: `No documentation found for library ID "${params.libraryId}"${params.topic ? ` with topic "${params.topic}"` : ""}.\n\nTry a different topic or verify the library ID with resolve_library_id.`,
						},
					],
					details: { empty: true },
				};
			}

			// Respect pi's standard truncation limits to avoid context overflow
			const truncation = truncateHead(raw, {
				maxLines: DEFAULT_MAX_LINES,
				maxBytes: DEFAULT_MAX_BYTES,
			});

			let result = truncation.content;
			if (truncation.truncated) {
				result +=
					`\n\n[Documentation truncated: ${truncation.outputLines} of ${truncation.totalLines} lines returned. ` +
					`Request fewer tokens or a more specific topic to get more targeted content.]`;
			}

			return {
				content: [{ type: "text", text: result }],
				details: {
					libraryId: params.libraryId,
					topic: params.topic,
					tokens,
					truncated: truncation.truncated,
					lines: truncation.outputLines,
				},
			};
		},
	});

	// Notify on load (dev convenience)
	pi.on("session_start", async (_event, ctx) => {
		ctx.ui.setStatus("context7", "⚡ context7");
	});

	pi.on("session_shutdown", async () => {
		// nothing to clean up — all requests use AbortSignal passed by pi
	});
}