/**
 * Snowflake Query — Snowflake + dbt access for Pi
 *
 * Registers:
 *   - tool `snowflake_query`     — single read-only SQL statement via snow CLI
 *   - tool `snowflake_describe`  — DESC TABLE shortcut for schema discovery
 *   - tool `snowflake_create_raw_external_table` — creates new raw external-table sources
 *   - tool `dbt`                 — wrapper for dbt commands (run/test/build/
 *                                  compile/seed/snapshot/debug/deps/parse/
 *                                  clean/docs-generate/run-operation/list/
 *                                  source-freshness) executed against the
 *                                  configured local dbt env
 *   - command `/snowflake`       — ad-hoc read-only query from the prompt
 *   - command `/snowflake-conn`  — print configured Snowflake + dbt context
 *   - tool_call gate             — blocks `bash` invocations of `snow sql`
 *     that bypass the `claude` connection or attempt destructive SQL
 *
 * Safety model (matches the dbt repo's CLAUDE.md rules):
 *   - All Snowflake queries run via `snow sql -c <SNOWFLAKE_CONNECTION>`
 *     (default `claude`). The connection is hard-defaulted; override only
 *     by setting SNOWFLAKE_CONNECTION explicitly.
 *   - Only one SQL statement per `snowflake_query` call. No `;`-chained
 *     statements. First keyword must be SELECT/WITH/SHOW/DESC/DESCRIBE/
 *     EXPLAIN/LIST.
 *   - dbt full-refresh is opt-in (fullRefresh: true). Selectors that
 *     mention raw tables are rejected even with the opt-in, matching
 *     the "never full-refresh raw tables" rule.
 *   - All output is truncated; the bounded captured output tail is saved to a
 *     temp file when the LLM-visible payload would exceed configured limits.
 *
 * Usage:
 *   pi -e extensions/snowflake-query/snowflake-query.ts
 *   Use `/skill:pi-snowflake-query` in the UI, or send messages invoking the tools directly.
 */

import type {
	AgentToolResult,
	ExtensionAPI,
} from "@earendil-works/pi-coding-agent";
import {
	DEFAULT_MAX_BYTES,
	DEFAULT_MAX_LINES,
	formatSize,
	isToolCallEventType,
	truncateTail,
} from "@earendil-works/pi-coding-agent";
import { StringEnum } from "@earendil-works/pi-ai";
import { Text } from "@earendil-works/pi-tui";
import { spawn } from "node:child_process";
import { mkdtemp, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Type, type Static } from "typebox";
import { hasContextModePiAdapter } from "./_shared/context-mode.ts";
import {
	buildCreateRawExternalTableSql,
	RAW_DATABASES,
	RAW_FILE_FORMATS,
	RAW_S3_LOCATIONS,
	type CreateRawExternalTableParams,
} from "./_shared/raw-external-table.ts";
import { validateReadOnlySql } from "./_shared/sql-gate.ts";
import {
	appendBounded,
	handleBashSnowflakeToolCall,
	resultOrThrow,
	selectSnowflakeQueryTool,
	validateFullRefreshRequest,
	validateResolvedFullRefreshCapture,
} from "./_shared/safety.ts";

// ── Config ──────────────────────────────────────────────────────────

const DEFAULT_CONNECTION = "claude";
const CONNECTION =
	process.env.SNOWFLAKE_CONNECTION?.trim() || DEFAULT_CONNECTION;
const SNOW_BIN = process.env.SNOW_CLI_BIN?.trim() || "snow";
const DEFAULT_TIMEOUT_MS = Math.max(
	5_000,
	Number.parseInt(process.env.SNOWFLAKE_QUERY_TIMEOUT_MS || "120000", 10) ||
		120_000,
);
const MAX_TIMEOUT_MS = 600_000;
// Keep process memory bounded even if the child produces gigabytes. The tail is
// what is useful for diagnostics, and normal tool output is smaller still.
const MAX_CAPTURE_BYTES = Math.max(DEFAULT_MAX_BYTES * 4, 256 * 1024);

// dbt: `DBT_BIN` may point directly at `dbt` (when the user has the venv
// already on PATH) or at a wrapper script like `./.claude/tools/dbt-run.sh`
// that activates the venv before invoking dbt. `DBT_PROJECT_DIR` is the
// working directory for dbt commands; defaults to the process cwd.
const DBT_BIN = process.env.DBT_BIN?.trim() || "dbt";
const DBT_PROJECT_DIR = process.env.DBT_PROJECT_DIR?.trim() || process.cwd();
const DEFAULT_DBT_TIMEOUT_MS = Math.max(
	5_000,
	Number.parseInt(process.env.DBT_TIMEOUT_MS || "1800000", 10) || 1_800_000,
);
const MAX_DBT_TIMEOUT_MS = 3_600_000;

// ── Process invocation helpers ──────────────────────────────────────

interface ProcRunResult {
	stdout: string;
	stderr: string;
	code: number;
	stdoutTruncated: boolean;
	stderrTruncated: boolean;
}

interface ProcRunOptions {
	bin: string;
	args: string[];
	signal: AbortSignal | undefined;
	timeoutMs: number;
	cwd?: string;
	missingHint?: string;
	onChunk?: (chunk: string, stream: "stdout" | "stderr") => void;
}

function runProcess(opts: ProcRunOptions): Promise<ProcRunResult> {
	return new Promise((resolve, reject) => {
		if (opts.signal?.aborted) {
			reject(new Error(`Aborted before spawning \`${opts.bin}\`.`));
			return;
		}

		let stdout = "";
		let stderr = "";
		let stdoutRawBytes = 0;
		let stderrRawBytes = 0;
		let settled = false;

		let child: ReturnType<typeof spawn>;
		try {
			child = spawn(opts.bin, opts.args, {
				signal: opts.signal,
				env: process.env,
				cwd: opts.cwd,
				stdio: ["ignore", "pipe", "pipe"],
			});
		} catch (err) {
			reject(err instanceof Error ? err : new Error(String(err)));
			return;
		}

		const finish = (result: ProcRunResult | Error) => {
			if (settled) return;
			settled = true;
			clearTimeout(timer);
			if (result instanceof Error) reject(result);
			else resolve(result);
		};

		const timer = setTimeout(() => {
			child.kill("SIGTERM");
			finish(
				new Error(`${opts.bin} command timed out after ${opts.timeoutMs}ms`),
			);
		}, opts.timeoutMs);

		child.stdout?.on("data", (chunk) => {
			const text = chunk.toString();
			stdoutRawBytes += Buffer.byteLength(text);
			stdout = appendBounded(stdout, text, MAX_CAPTURE_BYTES);
			opts.onChunk?.(text, "stdout");
		});
		child.stderr?.on("data", (chunk) => {
			const text = chunk.toString();
			stderrRawBytes += Buffer.byteLength(text);
			stderr = appendBounded(stderr, text, MAX_CAPTURE_BYTES);
			opts.onChunk?.(text, "stderr");
		});
		child.on("error", (err) => {
			if ((err as NodeJS.ErrnoException).code === "ENOENT") {
				finish(
					new Error(
						opts.missingHint ??
							`Could not find the \`${opts.bin}\` executable on PATH.`,
					),
				);
				return;
			}
			finish(err);
		});
		child.on("close", (code) => {
			finish({
				stdout,
				stderr,
				code: code ?? 0,
				stdoutTruncated: stdoutRawBytes > MAX_CAPTURE_BYTES,
				stderrTruncated: stderrRawBytes > MAX_CAPTURE_BYTES,
			});
		});
	});
}

function runSnow(
	args: string[],
	signal: AbortSignal | undefined,
	timeoutMs: number,
): Promise<ProcRunResult> {
	return runProcess({
		bin: SNOW_BIN,
		args,
		signal,
		timeoutMs,
		missingHint: `Could not find the \`${SNOW_BIN}\` executable. Install the Snowflake CLI and verify \`${SNOW_BIN} --version\` works, or set SNOW_CLI_BIN.`,
	});
}

function runDbt(
	args: string[],
	signal: AbortSignal | undefined,
	timeoutMs: number,
	onChunk?: (chunk: string, stream: "stdout" | "stderr") => void,
): Promise<ProcRunResult> {
	return runProcess({
		bin: DBT_BIN,
		args,
		signal,
		timeoutMs,
		cwd: DBT_PROJECT_DIR,
		onChunk,
		missingHint: `Could not find the \`${DBT_BIN}\` executable. Either point DBT_BIN at your dbt binary or at a wrapper script (e.g. ./.claude/tools/dbt-run.sh) that activates your dbt virtualenv.`,
	});
}

// ── Output helpers ──────────────────────────────────────────────────

interface SnowflakeQueryDetails {
	connection: string;
	statement: string;
	exitCode: number;
	truncated?: boolean;
	capturedOutputTailPath?: string;
	format: "json" | "csv" | "table";
}

interface DbtRunDetails {
	command: string;
	args: string[];
	cwd: string;
	exitCode: number;
	truncated?: boolean;
	capturedOutputTailPath?: string;
}

async function buildToolPayload<TDetails extends object>(opts: {
	tmpPrefix: string;
	label: string;
	exitName: string;
	result: ProcRunResult;
	details: TDetails;
}): Promise<AgentToolResult<TDetails>> {
	const { result, details, label, tmpPrefix, exitName } = opts;
	const combined = [result.stdout, result.stderr].filter(Boolean).join("\n");
	const truncation = truncateTail(combined, {
		maxLines: DEFAULT_MAX_LINES,
		maxBytes: DEFAULT_MAX_BYTES,
	});
	const finalDetails: TDetails & {
		exitCode: number;
		truncated?: boolean;
		capturedOutputTailPath?: string;
	} = { ...details, exitCode: result.code };

	let text = truncation.content || "(no output)";

	if (truncation.truncated) {
		const dir = await mkdtemp(join(tmpdir(), tmpPrefix));
		const capturedOutputTailPath = join(dir, "output-tail.txt");
		await writeFile(capturedOutputTailPath, combined, "utf8");
		finalDetails.truncated = true;
		finalDetails.capturedOutputTailPath = capturedOutputTailPath;
		text += `\n\n[${label} output truncated: showing ${truncation.outputLines} of ${truncation.totalLines} captured lines (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}). Captured output tail: ${capturedOutputTailPath}]`;
	}

	if (result.code !== 0) {
		text = `${exitName} exited ${result.code}\n${text}`;
	}

	return resultOrThrow(result.code, text, finalDetails) as AgentToolResult<TDetails>;
}

// ── Tool schemas ────────────────────────────────────────────────────

const QueryParams = Type.Object({
	query: Type.String({
		description:
			"Single read-only SQL statement. Allowed leading keywords: SELECT, WITH, SHOW, DESC, DESCRIBE, EXPLAIN, LIST. Use LIMIT 10 for discovery queries. Avoid SELECT *. No `;`-chained statements.",
	}),
	database: Type.Optional(
		Type.String({
			description:
				"Optional database to USE before the query (e.g. 'analytics', 'raw').",
		}),
	),
	schema: Type.Optional(
		Type.String({
			description: "Optional schema to USE before the query (e.g. 'dbt').",
		}),
	),
	warehouse: Type.Optional(
		Type.String({
			description:
				"Optional warehouse override. Defaults to whatever the connection profile selects.",
		}),
	),
	format: Type.Optional(
		StringEnum(["json", "csv", "table"] as const, {
			description: "Output format from snow CLI. Defaults to 'json'.",
		}),
	),
	timeoutMs: Type.Optional(
		Type.Integer({
			minimum: 1000,
			maximum: MAX_TIMEOUT_MS,
			description: `Override the per-query timeout (ms). Default ${DEFAULT_TIMEOUT_MS}, max ${MAX_TIMEOUT_MS}.`,
		}),
	),
});

type QueryParams = Static<typeof QueryParams>;

const CONTEXT_MODE_QUERY_TOOL = "ctx_snowflake_query";

const DescribeParams = Type.Object({
	table: Type.String({
		description:
			"Fully or partially qualified table name (e.g. 'analytics.dbt.model_aspyn_customers'). Whitespace is rejected.",
	}),
});

type DescribeParams = Static<typeof DescribeParams>;

const CreateRawExternalTableParams = Type.Object({
	database: StringEnum(RAW_DATABASES, {
		description: "Target raw database: `raw` for production or `raw_staging` for staging.",
	}),
	schema: Type.String({
		description: "Raw schema name, such as `aspyn` or `pestroutes`. Must be an unquoted identifier.",
	}),
	table: Type.String({
		description: "Raw table name. Must be an unquoted identifier.",
	}),
	s3Location: Type.Optional(StringEnum(RAW_S3_LOCATIONS, {
		description: "Named external stage to use. Defaults to `s3_etl_tmp`.",
	})),
	fileFormat: Type.Optional(StringEnum(RAW_FILE_FORMATS, {
		description: "Source file format and matching virtual-column layout. Defaults to `json`.",
	})),
	customS3Path: Type.Optional(Type.String({
		description: "Optional path below the selected stage. Defaults to the repository's standard schema/table layout.",
	})),
});

type CreateRawExternalTableToolParams = Static<typeof CreateRawExternalTableParams>;

// ── dbt command schema and safety helpers ───────────────────────────

const DBT_COMMANDS = [
	"run",
	"test",
	"build",
	"compile",
	"seed",
	"snapshot",
	"debug",
	"deps",
	"parse",
	"clean",
	"docs-generate",
	"run-operation",
	"list",
	"source-freshness",
] as const;

const DbtParams = Type.Object({
	command: StringEnum(DBT_COMMANDS, {
		description:
			"dbt subcommand to invoke. `run-operation` requires `operation`. `source-freshness` is shorthand for `source freshness`.",
	}),
	select: Type.Optional(
		Type.String({
			description:
				"dbt --select expression (e.g. 'tag:Aspyn', '+model_name', 'path:models/clean').",
		}),
	),
	exclude: Type.Optional(
		Type.String({ description: "dbt --exclude expression." }),
	),
	target: Type.Optional(
		StringEnum(["dev", "prod"] as const, {
			description:
				"dbt profile target. Defaults to whatever profiles.yml has set; usually 'dev'.",
		}),
	),
	fullRefresh: Type.Optional(
		Type.Boolean({
			description:
				"Pass --full-refresh. Rejected if the selector looks like a raw table (the post-hook `external_table_delete_files()` permanently deletes source S3 files).",
		}),
	),
	vars: Type.Optional(
		Type.String({
			description:
				"YAML/JSON string passed verbatim to --vars (e.g. '{allow_full_refresh: True}').",
		}),
	),
	operation: Type.Optional(
		Type.String({
			description:
				"Macro name for run-operation (required when command is 'run-operation').",
		}),
	),
	operationArgs: Type.Optional(
		Type.String({
			description: "JSON args passed via --args for run-operation.",
		}),
	),
	timeoutMs: Type.Optional(
		Type.Integer({
			minimum: 1000,
			maximum: MAX_DBT_TIMEOUT_MS,
			description: `Override per-invocation timeout (ms). Default ${DEFAULT_DBT_TIMEOUT_MS}, max ${MAX_DBT_TIMEOUT_MS}.`,
		}),
	),
});

type DbtParams = Static<typeof DbtParams>;


async function validateResolvedFullRefresh(
	params: Pick<DbtParams, "select" | "exclude" | "target" | "vars">,
	signal: AbortSignal | undefined,
	timeoutMs: number,
): Promise<string | undefined> {
	// dbt resolves graph selectors (including indirect selectors); this closes the
	// regex-only bypass without executing a model or contacting Snowflake.
	const args = ["ls", "--select", params.select!, "--output", "json"];
	if (params.exclude) args.push("--exclude", params.exclude);
	if (params.target) args.push("--target", params.target);
	if (params.vars) args.push("--vars", params.vars);
	const result = await runDbt(args, signal, timeoutMs);
	if (result.code !== 0) {
		return `Unable to resolve the --full-refresh selector safely (dbt ls exited ${result.code}). Refusing to run it.`;
	}
	return validateResolvedFullRefreshCapture(result.stdout, result.stdoutTruncated);
}

function buildDbtArgs(
	params: DbtParams,
): { args: string[] } | { error: string } {
	const args: string[] = [];

	if (params.command === "docs-generate") {
		args.push("docs", "generate");
	} else if (params.command === "source-freshness") {
		args.push("source", "freshness");
	} else {
		args.push(params.command);
	}

	if (params.command === "run-operation") {
		if (!params.operation) {
			return {
				error: "`operation` is required when command is 'run-operation'.",
			};
		}
		args.push(params.operation);
		if (params.operationArgs) {
			args.push("--args", params.operationArgs);
		}
	} else if (params.operation || params.operationArgs) {
		return {
			error:
				"`operation` and `operationArgs` are only valid with command 'run-operation'.",
		};
	}

	if (params.select) args.push("--select", params.select);
	if (params.exclude) args.push("--exclude", params.exclude);
	if (params.target) args.push("--target", params.target);
	if (params.vars) args.push("--vars", params.vars);

	if (params.fullRefresh) {
		const fullRefreshError = validateFullRefreshRequest(params.select);
		if (fullRefreshError) return { error: fullRefreshError };
		// Some commands don't accept --full-refresh, but dbt just warns; don't reject here.
		args.push("--full-refresh");
	}

	return { args };
}

// ── Extension factory ───────────────────────────────────────────────

export default function snowflakeQuery(pi: ExtensionAPI) {
	const connectionLabel = `snowflake (${CONNECTION})`;
	const isContextModeActive = () => hasContextModePiAdapter(pi.getAllTools());

	function syncQueryToolRouting() {
		const contextModeActive = isContextModeActive();
		const active = pi.getActiveTools();
		const next = active.filter(
			(name) =>
				name !== "snowflake_query" && name !== CONTEXT_MODE_QUERY_TOOL,
		);
		next.push(contextModeActive ? CONTEXT_MODE_QUERY_TOOL : "snowflake_query");
		pi.setActiveTools(next);
	}

	pi.on("session_start", (_event, ctx) => {
		if (ctx.hasUI) ctx.ui.setStatus("snowflake-query", connectionLabel);
		syncQueryToolRouting();
	});
	pi.on("before_agent_start", () => syncQueryToolRouting());
	pi.on("session_shutdown", (_event, ctx) => {
		if (ctx.hasUI) ctx.ui.setStatus("snowflake-query", undefined);
	});

	// Block direct `bash` calls that reach Snowflake outside this tool. We
	// catch:
	//   - the new `snow sql` CLI (require -c <CONNECTION> + read-only SQL),
	//   - the legacy `snowsql` CLI (block entirely; it has no equivalent
	//     connection flag we can validate cheaply),
	//   - python imports of `snowflake.connector` / `snowflake.snowpark`
	//     issued via `python -c` / `python -m` one-liners.
	// Python scripts invoked by path are not intercepted; that's an LLM-
	// discipline issue documented in the agent persona, not a hard gate.
	pi.on("tool_call", (event) => {
		if (event.toolName === "snowflake_query" && isContextModeActive()) {
			return {
				block: true,
				reason: `context-mode is active. Use the \`${CONTEXT_MODE_QUERY_TOOL}\` tool so Snowflake output is wrapped by the context-mode session.`,
			};
		}
		if (event.toolName === CONTEXT_MODE_QUERY_TOOL && !isContextModeActive()) {
			return {
				block: true,
				reason: `\`${CONTEXT_MODE_QUERY_TOOL}\` is only available when the context-mode Pi adapter is active. Use \`snowflake_query\` instead.`,
			};
		}

		if (!isToolCallEventType("bash", event)) return;
		const command = String(event.input.command ?? "");
		return handleBashSnowflakeToolCall(command, CONNECTION);
	});

	async function executeQuery(
		params: QueryParams,
		signal: AbortSignal | undefined,
		onUpdate: ((result: AgentToolResult<SnowflakeQueryDetails>) => void) | undefined,
		label: string,
	): Promise<AgentToolResult<SnowflakeQueryDetails>> {
		const validation = validateReadOnlySql(params.query);
		if (!validation.ok) throw new Error(validation.reason);

		const format = params.format ?? "json";
		const timeoutMs = Math.min(
			params.timeoutMs ?? DEFAULT_TIMEOUT_MS,
			MAX_TIMEOUT_MS,
		);
		const args = [
			"sql",
			"-c",
			CONNECTION,
			"--format",
			format,
			"-q",
			validation.statement,
		];
		if (params.database) args.push("--database", params.database);
		if (params.schema) args.push("--schema", params.schema);
		if (params.warehouse) args.push("--warehouse", params.warehouse);

		onUpdate?.({
			content: [
				{
					type: "text",
					text: `Running on connection \`${CONNECTION}\`...`,
				},
			],
			details: {
				connection: CONNECTION,
				statement: validation.statement,
				exitCode: -1,
				format,
			},
		});

		const result = await runSnow(args, signal, timeoutMs);
		return buildToolPayload<SnowflakeQueryDetails>({
			tmpPrefix: "pi-snowflake-",
			label,
			exitName: "snow",
			result,
			details: {
				connection: CONNECTION,
				statement: validation.statement,
				exitCode: -1,
				format,
			},
		});
	}

	pi.registerTool({
		name: "snowflake_query",
		label: "Snowflake Query",
		description:
			"Run a single read-only SQL statement against Snowflake using the team `claude` connection. Reject DDL/DML up front. Use for SELECT/SHOW/DESC/EXPLAIN/LIST/WITH. Output is truncated; only the bounded captured tail (not the full result) is saved to a temp file when oversized.",
		promptSnippet:
			"Use snowflake_query to inspect Snowflake data with read-only SQL (SELECT/SHOW/DESC). Always include LIMIT on discovery SELECTs and avoid SELECT *.",
		promptGuidelines: [
			"snowflake_query is the only sanctioned way to run Snowflake SQL — do not invoke `snow sql` via bash.",
			"snowflake_query: prefer DESC TABLE, SHOW TABLES, SHOW COLUMNS for schema discovery before issuing a SELECT.",
			"snowflake_query: cap discovery SELECTs with LIMIT 10 (max LIMIT 100 unless the user asks for more) and truncate long text/JSON columns with LEFT(col, 100).",
		],
		parameters: QueryParams,

		async execute(_toolCallId, params: QueryParams, signal, onUpdate) {
			return executeQuery(params, signal, onUpdate, "snowflake_query");
		},

		renderCall(args, theme) {
			const preview = String(args.query ?? "")
				.replace(/\s+/g, " ")
				.slice(0, 120);
			return new Text(
				`${theme.fg("toolTitle", theme.bold("snowflake_query "))}${theme.fg("accent", preview)}`,
				0,
				0,
			);
		},

		renderResult(result, { expanded, isPartial }, theme) {
			if (isPartial)
				return new Text(theme.fg("warning", "Querying Snowflake..."), 0, 0);
			const details = result.details as SnowflakeQueryDetails | undefined;
			let text =
				details?.exitCode === 0
					? theme.fg("success", `ok (${details?.format ?? "json"})`)
					: theme.fg("error", `exit ${details?.exitCode ?? "unknown"}`);
			if (details?.truncated) text += theme.fg("warning", " (truncated)");
			if (expanded && details?.capturedOutputTailPath) {
				text += `\n${theme.fg("dim", `Captured output tail: ${details.capturedOutputTailPath}`)}`;
			}
			return new Text(text, 0, 0);
		},
	});

	pi.registerTool({
		name: CONTEXT_MODE_QUERY_TOOL,
		label: "Context-mode Snowflake Query",
		description:
			"Run a single read-only Snowflake SQL statement through the context-mode session. Available only when the context-mode Pi adapter is active; uses the team `claude` connection and the same SQL safety gate as snowflake_query.",
		promptSnippet:
			"When context-mode is active, use ctx_snowflake_query for all read-only Snowflake SQL.",
		promptGuidelines: [
			"When context-mode is active, ctx_snowflake_query is the only sanctioned way to run Snowflake SQL. Do not call snowflake_query or `snow sql` directly.",
		],
		parameters: QueryParams,

		async execute(_toolCallId, params: QueryParams, signal, onUpdate) {
			return executeQuery(params, signal, onUpdate, CONTEXT_MODE_QUERY_TOOL);
		},

		renderCall(args, theme) {
			const preview = String(args.query ?? "")
				.replace(/\s+/g, " ")
				.slice(0, 120);
			return new Text(
				`${theme.fg("toolTitle", theme.bold(`${CONTEXT_MODE_QUERY_TOOL} `))}${theme.fg("accent", preview)}`,
				0,
				0,
			);
		},

		renderResult(result, { expanded, isPartial }, theme) {
			if (isPartial)
				return new Text(theme.fg("warning", "Querying Snowflake..."), 0, 0);
			const details = result.details as SnowflakeQueryDetails | undefined;
			let text =
				details?.exitCode === 0
					? theme.fg("success", `ok (${details?.format ?? "json"})`)
					: theme.fg("error", `exit ${details?.exitCode ?? "unknown"}`);
			if (details?.truncated) text += theme.fg("warning", " (truncated)");
			if (expanded && details?.capturedOutputTailPath) {
				text += `\n${theme.fg("dim", `Captured output tail: ${details.capturedOutputTailPath}`)}`;
			}
			return new Text(text, 0, 0);
		},
	});

	pi.registerTool({
		name: "snowflake_describe",
		label: "Snowflake Describe",
		description:
			"Shortcut for `DESC TABLE <name>` on Snowflake. Use this for schema discovery — it is a metadata operation with near-zero compute and is cheaper than querying INFORMATION_SCHEMA.",
		promptSnippet:
			"Use snowflake_describe to inspect a Snowflake table's columns and types without spinning up a warehouse.",
		parameters: DescribeParams,

		async execute(_toolCallId, params: DescribeParams, signal) {
			const cleaned = params.table.trim();
			if (!cleaned || /\s/.test(cleaned)) {
				throw new Error(
					"Table name must be a single identifier (e.g. 'analytics.dbt.my_model'). No whitespace.",
				);
			}
			if (!/^[A-Za-z0-9_.$"]+$/.test(cleaned)) {
				throw new Error(
					"Table name contains unsupported characters. Allowed: letters, digits, underscore, period, dollar sign, double-quote.",
				);
			}
			const statement = `DESC TABLE ${cleaned}`;
			const result = await runSnow(
				["sql", "-c", CONNECTION, "--format", "json", "-q", statement],
				signal,
				DEFAULT_TIMEOUT_MS,
			);
			return buildToolPayload<SnowflakeQueryDetails>({
				tmpPrefix: "pi-snowflake-",
				label: "snowflake_describe",
				exitName: "snow",
				result,
				details: {
					connection: CONNECTION,
					statement,
					exitCode: -1,
					format: "json",
				},
			});
		},

		renderCall(args, theme) {
			return new Text(
				`${theme.fg("toolTitle", theme.bold("snowflake_describe "))}${theme.fg("accent", String(args.table ?? ""))}`,
				0,
				0,
			);
		},
	});

	pi.registerTool({
		name: "snowflake_create_raw_external_table",
		label: "Create Raw External Table",
		description:
			"Create a non-destructive, IF NOT EXISTS Snowflake external table for the dbt raw ingestion pattern. This creates only the external source table; create the dbt raw model and schema YAML separately, then run that selected dbt model. Uses the team `claude` connection.",
		promptSnippet:
			"Create a raw external-table source using snowflake_create_raw_external_table. Confirm the S3 path and file format first; then create the matching dbt raw model separately.",
		promptGuidelines: [
			"Use this only to create a new external source table in raw or raw_staging. It never replaces or drops an existing table.",
			"Before calling it, confirm the staged files exist and choose the exact database, S3 location, file format, and optional custom path.",
			"After creation, create the dbt raw model with full-refresh protection. Never full-refresh a raw model unless the user explicitly authorizes it and source files are confirmed present.",
		],
		parameters: CreateRawExternalTableParams,

		async execute(_toolCallId, params: CreateRawExternalTableToolParams, signal, onUpdate) {
			const statement = buildCreateRawExternalTableSql(params as CreateRawExternalTableParams);
			onUpdate?.({
				content: [{ type: "text", text: `Creating external table on connection \`${CONNECTION}\`...` }],
				details: { connection: CONNECTION, statement, exitCode: -1, format: "table" },
			});
			const result = await runSnow(
				["sql", "-c", CONNECTION, "--format", "table", "-q", statement],
				signal,
				DEFAULT_TIMEOUT_MS,
			);
			return buildToolPayload<SnowflakeQueryDetails>({
				tmpPrefix: "pi-snowflake-",
				label: "snowflake_create_raw_external_table",
				exitName: "snow",
				result,
				details: { connection: CONNECTION, statement, exitCode: -1, format: "table" },
			});
		},

		renderCall(args, theme) {
			return new Text(
				`${theme.fg("toolTitle", theme.bold("snowflake_create_raw_external_table "))}${theme.fg("accent", `${String(args.database ?? "")}.${String(args.schema ?? "")}.${String(args.table ?? "")}`)}`,
				0,
				0,
			);
		},
	});

	pi.registerTool({
		name: "dbt",
		label: "dbt",
		description:
			"Run a dbt subcommand (run, test, build, compile, seed, snapshot, debug, deps, parse, clean, docs-generate, run-operation, list, source-freshness) against the configured local dbt env. Output is truncated; the bounded captured tail is saved to a temp file when oversized.",
		promptSnippet:
			"Use the dbt tool to run dbt commands (run/test/build/compile/seed/snapshot/debug/deps/etc.) without shelling out to bash.",
		promptGuidelines: [
			"dbt: prefer `compile` or `parse` for cheap correctness checks before running `run` or `build`.",
			"dbt: never pass fullRefresh:true on a selector that touches raw tables — the tool will refuse, and bypassing it via bash risks permanent S3 data loss.",
			"dbt: scope `run`/`test`/`build` with `select` so you don't accidentally rebuild the whole project.",
		],
		parameters: DbtParams,

		async execute(_toolCallId, params: DbtParams, signal, onUpdate) {
			const build = buildDbtArgs(params);
			if ("error" in build) throw new Error(build.error);
			const timeoutMs = Math.min(
				params.timeoutMs ?? DEFAULT_DBT_TIMEOUT_MS,
				MAX_DBT_TIMEOUT_MS,
			);
			if (params.fullRefresh) {
				const resolutionError = await validateResolvedFullRefresh(
					params,
					signal,
					timeoutMs,
				);
				if (resolutionError) throw new Error(resolutionError);
			}

			const partialDetails: DbtRunDetails = {
				command: params.command,
				args: build.args,
				cwd: DBT_PROJECT_DIR,
				exitCode: -1,
			};
			onUpdate?.({
				content: [
					{
						type: "text",
						text: `Running \`${DBT_BIN} ${build.args.join(" ")}\` in ${DBT_PROJECT_DIR}...`,
					},
				],
				details: partialDetails,
			});

			let lastChunk = "";
			const result = await runDbt(build.args, signal, timeoutMs, (chunk) => {
				lastChunk = `${lastChunk}${chunk}`.slice(-2_000);
				onUpdate?.({
					content: [{ type: "text", text: lastChunk }],
					details: partialDetails,
				});
			});

			return buildToolPayload<DbtRunDetails>({
				tmpPrefix: "pi-dbt-",
				label: "dbt",
				exitName: DBT_BIN,
				result,
				details: {
					command: params.command,
					args: build.args,
					cwd: DBT_PROJECT_DIR,
					exitCode: -1,
				},
			});
		},

		renderCall(args, theme) {
			const command = String(args.command ?? "");
			const selector = args.select
				? ` ${theme.fg("accent", String(args.select))}`
				: "";
			return new Text(
				`${theme.fg("toolTitle", theme.bold("dbt "))}${theme.fg("toolTitle", command)}${selector}`,
				0,
				0,
			);
		},

		renderResult(result, { expanded, isPartial }, theme) {
			if (isPartial)
				return new Text(theme.fg("warning", "Running dbt..."), 0, 0);
			const details = result.details as DbtRunDetails | undefined;
			let text =
				details?.exitCode === 0
					? theme.fg("success", `dbt ${details?.command ?? ""} ok`)
					: theme.fg(
							"error",
							`dbt ${details?.command ?? ""} exit ${details?.exitCode ?? "unknown"}`,
						);
			if (details?.truncated) text += theme.fg("warning", " (truncated)");
			if (expanded && details?.capturedOutputTailPath) {
				text += `\n${theme.fg("dim", `Captured output tail: ${details.capturedOutputTailPath}`)}`;
			}
			return new Text(text, 0, 0);
		},
	});

	pi.registerCommand("snowflake", {
		description:
			"Run an ad-hoc read-only Snowflake query. Usage: /snowflake SELECT current_warehouse()",
		handler: async (args, ctx) => {
			await ctx.waitForIdle();
			if (!args || !args.trim()) {
				if (ctx.hasUI)
					ctx.ui.notify("Usage: /snowflake <read-only SQL>", "warning");
				return;
			}
			pi.sendUserMessage(
				`Run this Snowflake query for me using the ${selectSnowflakeQueryTool(isContextModeActive())} tool:\n\n${args.trim()}`,
			);
		},
	});

	pi.registerCommand("snowflake-conn", {
		description: "Show the configured Snowflake + dbt context.",
		handler: async (_args, ctx) => {
			await ctx.waitForIdle();
			const message = [
				`Snowflake connection: \`${CONNECTION}\` (override with SNOWFLAKE_CONNECTION).`,
				`snow CLI: \`${SNOW_BIN}\` (override with SNOW_CLI_BIN).`,
				`Snowflake query timeout: ${DEFAULT_TIMEOUT_MS} ms.`,
				`dbt binary: \`${DBT_BIN}\` (override with DBT_BIN).`,
				`dbt project dir: \`${DBT_PROJECT_DIR}\` (override with DBT_PROJECT_DIR).`,
				`dbt timeout: ${DEFAULT_DBT_TIMEOUT_MS} ms.`,
			].join(" ");
			if (ctx.hasUI) ctx.ui.notify(message, "info");
			else process.stdout.write(`${message}\n`);
		},
	});

	pi.registerCommand("dbt", {
		description:
			"Ask Pi to run a dbt command via the dbt tool. Usage: /dbt run --select tag:Aspyn",
		handler: async (args, ctx) => {
			await ctx.waitForIdle();
			if (!args || !args.trim()) {
				if (ctx.hasUI)
					ctx.ui.notify(
						"Usage: /dbt <command> [args]  e.g. /dbt test --select model_name",
						"warning",
					);
				return;
			}
			pi.sendUserMessage(
				`Run this dbt invocation using the dbt tool:\n\n${args.trim()}`,
			);
		},
	});
}