import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { setCharlesThrottling, THROTTLE_PRESETS } from "./client.js";
import { calcEntropy, getHitLocations, simplifyEntry } from "./process.js";
import {
	authorizeKeyword,
	dataHint,
	findEntry,
	getCache,
	harvestData,
	isKeywordAuthed,
	listCheckpoints,
	loadCheckpoint,
	loadRecording,
} from "./store.js";

function ok(data: unknown) {
	return {
		content: [{ type: "text" as const, text: JSON.stringify(data, null, 2) }],
		details: data as Record<string, unknown>,
	};
}

function top(
	d: Map<string, number>,
	n = 10,
): Array<{ key: string; count: number }> {
	return [...d.entries()]
		.sort((a, b) => b[1] - a[1])
		.slice(0, n)
		.map(([key, count]) => ({ key, count }));
}

/** Register all Charles analysis tools on the extension API. */
export function registerCharlesTools(pi: ExtensionAPI): void {
	// ── harvest ──────────────────────────────────────────────
	pi.registerTool({
		name: "harvest_data",
		label: "Harvest Data",
		description:
			"Sync incremental traffic from Charles and create a new checkpoint on the timeline. " +
			"fresh_start=false (default) keeps only entries after the last harvest; " +
			"fresh_start=true inserts a reset point and only future traffic is returned. " +
			"Each call clears the Charles session and restarts recording.",
		promptSnippet:
			"Incrementally harvest Charles traffic and create a checkpoint",
		promptGuidelines: [
			"Call harvest_data() to sync traffic before analyzing Charles captures.",
			"Use harvest_data(fresh_start=true) to reset the window when switching analysis targets.",
		],
		parameters: Type.Object({
			fresh_start: Type.Optional(
				Type.Boolean({
					description:
						"true = insert a reset point without loading entries; default false",
				}),
			),
		}),
		async execute(_id, params) {
			return ok(await harvestData(params.fresh_start ?? false));
		},
	});

	// ── timeline ─────────────────────────────────────────────
	pi.registerTool({
		name: "list_checkpoints",
		label: "List Checkpoints",
		description:
			"List all checkpoints on the harvest timeline. " +
			"The active checkpoint is marked by current_checkpoint_id.",
		promptSnippet: "List the Charles harvest timeline",
		parameters: Type.Object({}),
		async execute() {
			return ok(listCheckpoints());
		},
	});

	pi.registerTool({
		name: "load_checkpoint",
		label: "Load Checkpoint",
		description:
			"Switch the visible window to a checkpoint so all filter tools operate in that time range. " +
			"Read-only — ARCHIVE is not modified. Call harvest_data() to return to the latest increment.",
		promptSnippet: "Switch to a historical checkpoint time window",
		parameters: Type.Object({
			checkpoint_id: Type.Number({
				description: "id field from list_checkpoints",
			}),
		}),
		async execute(_id, params) {
			return ok(loadCheckpoint(params.checkpoint_id));
		},
	});

	pi.registerTool({
		name: "load_recording",
		label: "Load Recording",
		description:
			"Load a local .chlsj recording into ARCHIVE, create a checkpoint, and switch to it. " +
			"Recording entries use the rec:filename:id namespace so they never clash with live ids.",
		promptSnippet: "Load a local .chlsj recording for analysis",
		parameters: Type.Object({
			file_path: Type.String({
				description: "Absolute or relative path to a .chlsj file",
			}),
		}),
		async execute(_id, params) {
			return ok(await loadRecording(params.file_path));
		},
	});

	// ── overview ─────────────────────────────────────────────
	pi.registerTool({
		name: "summarize_traffic",
		label: "Summarize Traffic",
		description:
			"Global stats for the current checkpoint (top hosts/paths, status and method distributions). " +
			"Does not load bodies — useful for quick orientation.",
		promptSnippet: "Summarize current traffic (host/path/status/method)",
		parameters: Type.Object({}),
		async execute() {
			const hint = dataHint();
			if (hint) return ok({ warn: hint, total: 0 });

			const data = getCache().data;
			const hostCount = new Map<string, number>();
			const pathCount = new Map<string, number>();
			const statusCount = new Map<string, number>();
			const methodCount = new Map<string, number>();

			for (const e of data) {
				const host = e.host || "unknown";
				const path = e.path || "/";
				const method = e.method || "unknown";
				const status = String(e.response?.status ?? "unknown");
				const parts = path.split("/").filter(Boolean);
				const pathKey = parts.length ? "/" + parts.slice(0, 2).join("/") : "/";

				hostCount.set(host, (hostCount.get(host) ?? 0) + 1);
				pathCount.set(pathKey, (pathCount.get(pathKey) ?? 0) + 1);
				statusCount.set(status, (statusCount.get(status) ?? 0) + 1);
				methodCount.set(method, (methodCount.get(method) ?? 0) + 1);
			}

			return ok({
				checkpoint_id: getCache().checkpoint_id,
				total: data.length,
				top_hosts: top(hostCount),
				top_paths: top(pathCount),
				status_dist: top(statusCount),
				method_dist: top(methodCount),
			});
		},
	});

	// ── filters ──────────────────────────────────────────────
	pi.registerTool({
		name: "check_keyword_exists",
		label: "Check Keyword Exists",
		description:
			"[Interlock-1] Probe which entries in the current checkpoint contain a keyword. " +
			"Returns a lightweight index of id + hit locations only (no bodies). " +
			"Unlocks filter_by_keyword for this keyword for 5 minutes.",
		promptSnippet: "Probe keyword locations (unlocks filter_by_keyword)",
		parameters: Type.Object({
			keyword: Type.String({ description: "Keyword to probe" }),
		}),
		async execute(_id, params) {
			const hint = dataHint();
			if (hint) return ok({ warn: hint, total: 0, matches: [] });

			const matches = [];
			for (const e of getCache().data) {
				const hits = getHitLocations(e, params.keyword);
				if (hits.length) {
					matches.push({ id: e.id, path: e.path, hit_at: hits });
				}
			}
			authorizeKeyword(params.keyword);
			return ok({ total: matches.length, matches: matches.slice(0, 50) });
		},
	});

	pi.registerTool({
		name: "filter_by_keyword",
		label: "Filter By Keyword",
		description:
			"[Interlock-2] Return simplified entries that contain the keyword (with body previews). " +
			"If more than 30 matches, call check_keyword_exists first to unlock. " +
			"limit defaults to 10, max 50; returns the most recent N matches.",
		promptSnippet: "Filter traffic entries by keyword",
		parameters: Type.Object({
			keyword: Type.String({ description: "Search keyword" }),
			limit: Type.Optional(
				Type.Number({
					description: "Number of results to return; default 10, max 50",
				}),
			),
		}),
		async execute(_id, params) {
			const hint = dataHint();
			if (hint) return ok({ warn: hint, results: [], total_matched: 0 });

			const k = params.keyword.toLowerCase();
			const authed = isKeywordAuthed(k);
			const filtered = getCache()
				.data.filter((e) => JSON.stringify(e).toLowerCase().includes(k))
				.map((e) => simplifyEntry(e));

			if (filtered.length > 30 && !authed) {
				return ok({
					error: "PRE_CHECK_REQUIRED",
					message: `Matched ${filtered.length} entries. Call check_keyword_exists('${params.keyword}') first, then retry.`,
				});
			}

			const limit = Math.min(params.limit ?? 10, 50);
			return ok({
				results: filtered.slice(-limit),
				total_matched: filtered.length,
				returned: Math.min(filtered.length, limit),
			});
		},
	});

	pi.registerTool({
		name: "filter_by_path",
		label: "Filter By Path",
		description:
			"Filter by URL path keyword and return simplified entries. " +
			"Useful for locating API routes such as '/api/sign' or '/login'.",
		promptSnippet: "Filter traffic by URL path",
		parameters: Type.Object({
			path_keyword: Type.String({ description: "Path keyword" }),
			limit: Type.Optional(
				Type.Number({
					description: "Number of results to return; default 15, max 50",
				}),
			),
		}),
		async execute(_id, params) {
			const hint = dataHint();
			if (hint) return ok({ warn: hint, results: [], total_matched: 0 });

			const k = params.path_keyword.toLowerCase();
			const results = getCache()
				.data.filter((e) => (e.path ?? "").toLowerCase().includes(k))
				.map((e) => simplifyEntry(e));
			const limit = Math.min(params.limit ?? 15, 50);
			return ok({
				results: results.slice(-limit),
				total_matched: results.length,
				returned: Math.min(results.length, limit),
			});
		},
	});

	pi.registerTool({
		name: "filter_by_host",
		label: "Filter By Host",
		description:
			"Filter by the top-level host field only (never body/headers). " +
			"Partial match is supported, e.g. 'example.com' matches all subdomains.",
		promptSnippet: "Filter traffic by host",
		parameters: Type.Object({
			host_keyword: Type.String({ description: "Host keyword" }),
			limit: Type.Optional(
				Type.Number({
					description: "Number of results to return; default 15, max 50",
				}),
			),
		}),
		async execute(_id, params) {
			const hint = dataHint();
			if (hint) return ok({ warn: hint, results: [], total_matched: 0 });

			const k = params.host_keyword.toLowerCase();
			const results = getCache()
				.data.filter((e) => (e.host ?? "").toLowerCase().includes(k))
				.map((e) => simplifyEntry(e));
			const limit = Math.min(params.limit ?? 15, 50);
			return ok({
				results: results.slice(-limit),
				total_matched: results.length,
				returned: Math.min(results.length, limit),
			});
		},
	});

	pi.registerTool({
		name: "filter_by_status",
		label: "Filter By Status",
		description:
			"Filter by HTTP status code and return simplified entries. " +
			"Useful for investigating 403/500 and other error responses.",
		promptSnippet: "Filter traffic by HTTP status code",
		parameters: Type.Object({
			status_code: Type.Number({
				description: "HTTP status code, e.g. 403",
			}),
			limit: Type.Optional(
				Type.Number({
					description: "Number of results to return; default 15, max 50",
				}),
			),
		}),
		async execute(_id, params) {
			const hint = dataHint();
			if (hint) return ok({ warn: hint, results: [], total_matched: 0 });

			const results = getCache()
				.data.filter((e) => e.response?.status === params.status_code)
				.map((e) => simplifyEntry(e));
			const limit = Math.min(params.limit ?? 15, 50);
			return ok({
				results: results.slice(-limit),
				total_matched: results.length,
				returned: Math.min(results.length, limit),
			});
		},
	});

	pi.registerTool({
		name: "filter_by_method",
		label: "Filter By Method",
		description:
			"Filter by HTTP method (GET / POST / PUT / DELETE, etc.). Case-insensitive.",
		promptSnippet: "Filter traffic by HTTP method",
		parameters: Type.Object({
			method: Type.String({ description: "HTTP method, e.g. POST" }),
			limit: Type.Optional(
				Type.Number({
					description: "Number of results to return; default 15, max 50",
				}),
			),
		}),
		async execute(_id, params) {
			const hint = dataHint();
			if (hint) return ok({ warn: hint, results: [], total_matched: 0 });

			const m = params.method.toUpperCase();
			const results = getCache()
				.data.filter((e) => (e.method ?? "").toUpperCase() === m)
				.map((e) => simplifyEntry(e));
			const limit = Math.min(params.limit ?? 15, 50);
			return ok({
				results: results.slice(-limit),
				total_matched: results.length,
				returned: Math.min(results.length, limit),
			});
		},
	});

	// ── deep analysis ────────────────────────────────────────
	pi.registerTool({
		name: "filter_by_encryption",
		label: "Filter By Encryption",
		description:
			"Scan the current checkpoint for bodies that look encrypted/compressed/encoded. " +
			"Sorted by max(req_entropy, res_entropy) descending. Default threshold 3.9: " +
			"plain JSON is usually 2.5–3.5, Base64 about 4.0–5.0, encrypted > 5.0.",
		promptSnippet: "Find likely encrypted bodies via Shannon entropy",
		parameters: Type.Object({
			threshold: Type.Optional(
				Type.Number({ description: "Entropy threshold; default 3.9" }),
			),
			limit: Type.Optional(
				Type.Number({
					description: "Number of results to return; default 20, max 100",
				}),
			),
		}),
		async execute(_id, params) {
			const hint = dataHint();
			if (hint) return ok({ warn: hint, results: [], total_matched: 0 });

			const threshold = params.threshold ?? 3.9;
			const hits: Array<{
				ent: number;
				row: ReturnType<typeof simplifyEntry>;
			}> = [];

			for (const e of getCache().data) {
				const reqTxt = e.request?.body?.text ?? "";
				const resTxt = e.response?.body?.text ?? "";
				const reqEnt = calcEntropy(reqTxt);
				const resEnt = calcEntropy(resTxt);
				const maxEnt = Math.max(reqEnt, resEnt);
				if (maxEnt > threshold) {
					const row = simplifyEntry(e);
					row.req_entropy = Math.round(reqEnt * 100) / 100;
					row.res_entropy = Math.round(resEnt * 100) / 100;
					row.req_body_len = reqTxt.length;
					row.res_body_len = resTxt.length;
					hits.push({ ent: maxEnt, row });
				}
			}

			hits.sort((a, b) => b.ent - a.ent);
			const results = hits.map((h) => h.row);
			const limit = Math.min(params.limit ?? 20, 100);
			return ok({
				results: results.slice(0, limit),
				total_matched: results.length,
				returned: Math.min(results.length, limit),
				threshold,
			});
		},
	});

	pi.registerTool({
		name: "get_raw_data",
		label: "Get Raw Data",
		description:
			"Fetch the full raw entry (headers, body, timing, and all other fields). " +
			"Looks in the current checkpoint first, then the global ARCHIVE.",
		promptSnippet: "Get the full raw data for a traffic entry",
		parameters: Type.Object({
			entry_id: Type.String({
				description: "Entry id (live numeric id or rec:file:id)",
			}),
		}),
		async execute(_id, params) {
			const entry = findEntry(params.entry_id);
			if (!entry) {
				return ok({ error: "NOT_FOUND", entry_id: params.entry_id });
			}
			return ok({ entry });
		},
	});

	// ── environment ──────────────────────────────────────────
	pi.registerTool({
		name: "set_throttling",
		label: "Set Throttling",
		description:
			"Toggle Charles network bandwidth limits. " +
			"preset = preset name enables throttling; null/omitted disables it. " +
			`Available presets: ${THROTTLE_PRESETS.join(" | ")}. ` +
			"Throttling is automatically disabled when the session ends.",
		promptSnippet: "Enable or disable Charles network throttling",
		parameters: Type.Object({
			preset: Type.Optional(
				Type.Union([Type.String(), Type.Null()], {
					description: "Throttle preset name; null/omitted = disable",
				}),
			),
		}),
		async execute(_id, params, _signal, _onUpdate, ctx) {
			const preset = params.preset ?? null;
			const result = await setCharlesThrottling(preset);
			if (result.message) {
				ctx.ui.notify(result.message, result.success ? "info" : "error");
			}
			return ok({
				action: preset ? "activate" : "deactivate",
				preset,
				success: result.success,
				message: result.message,
			});
		},
	});
}