/**
 * File Tree Extension
 *
 * Provides a file_tree tool that returns a tree structure of a requested path.
 *
 * Features:
 * - Toggle between folders only or all files (includeFiles)
 * - Include or exclude hidden files starting with dot (includeHidden)
 * - ASCII tree visualization with ├── and └── connectors
 * - Smart auto-depth: auto-expands to fill available context
 *
 * Parameters:
 * - path: Path to generate tree for (relative or absolute)
 * - includeFiles: Include files in the tree (default: false, folders only)
 * - includeHidden: Include hidden files/dirs starting with dot (default: false)
 * - depth: Max directory depth to traverse (default: null = smart auto-depth)
 */

import * as fs from "node:fs/promises";
import * as path from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "@sinclair/typebox";

/** Hard-coded folder names that are always excluded from the tree. */
const EXCLUDED_FOLDERS = ["node_modules", "dist"];

const MAX_LINES = 400;

interface TreeEntry {
	name: string;
	isDirectory: boolean;
	size?: number;
}

const TEXT_EXTS = new Set([
	".txt", ".md", ".tsx", ".ts", ".jsx", ".js", ".json", ".css", ".scss",
	".html", ".yml", ".yaml", ".xml", ".svg", ".toml", ".ini", ".sh", ".py",
	".rs", ".go", ".java", ".c", ".cpp", ".h", ".hpp", ".rb", ".php",
	".sql", ".vue", ".svelte",
]);

function isTextFile(name: string): boolean {
	return TEXT_EXTS.has(path.extname(name).toLowerCase());
}

async function countLines(filePath: string): Promise<number> {
	try {
		const buf = await fs.readFile(filePath);
		let count = buf.length === 0 ? 0 : 1;
		for (let i = 0; i < buf.length; i++) if (buf[i] === 10) count++;
		return count;
	} catch {
		return 0;
	}
}

function formatSize(bytes: number): string {
	if (bytes < 1024) return `${bytes}b`;
	if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed()}Kb`;
	if (bytes < 1024 * 1024 * 1024)
		return `${(bytes / (1024 * 1024)).toFixed(1)}Mb`;
	return `${(bytes / (1024 * 1024 * 1024)).toFixed(1)}Gb`;
}

async function getDirectoryEntries(
	dirPath: string,
	includeHidden: boolean,
): Promise<TreeEntry[]> {
	try {
		const entries = await fs.readdir(dirPath);
		const result: TreeEntry[] = [];

		for (const entry of entries) {
			if (!includeHidden && entry.startsWith(".")) continue;
			const stats = await fs.stat(path.join(dirPath, entry));
			if (stats.isDirectory() && EXCLUDED_FOLDERS.includes(entry)) continue;
			result.push({
				name: entry,
				isDirectory: stats.isDirectory(),
				size: stats.isDirectory() ? undefined : stats.size,
			});
		}

		return result.sort((a, b) => {
			if (a.isDirectory !== b.isDirectory) return a.isDirectory ? -1 : 1;
			return a.name.localeCompare(b.name);
		});
	} catch {
		return [];
	}
}

async function countEntries(
	dir: string,
	incFiles: boolean,
	incHidden: boolean,
): Promise<{ dirs: number; files: number }> {
	const entries = await getDirectoryEntries(dir, incHidden);
	let dirs = 0,
		files = 0;
	for (const e of entries) {
		if (e.isDirectory) {
			dirs++;
			const sub = await countEntries(
				path.join(dir, e.name),
				incFiles,
				incHidden,
			);
			dirs += sub.dirs;
			files += sub.files;
		} else if (incFiles) files++;
	}
	return { dirs, files };
}

async function buildTree(
	dirPath: string,
	prefix: string,
	includeFiles: boolean,
	includeHidden: boolean,
	maxDepth: number | null,
	currentDepth: number = 0,
): Promise<string> {
	const entries = await getDirectoryEntries(dirPath, includeHidden);
	const visible = includeFiles ? entries : entries.filter((e) => e.isDirectory);
	const lines: string[] = [];

	for (let i = 0; i < visible.length; i++) {
		const entry = visible[i];
		const isLast = i === visible.length - 1;
		const connector = isLast ? "└─ " : "├─ ";

		let label: string;
		if (entry.isDirectory) {
			label = entry.name;
		} else {
			label = `${entry.name} · ${formatSize(entry.size ?? 0)}`;
			if (isTextFile(entry.name)) {
				label += ` · ${await countLines(path.join(dirPath, entry.name))} lines`;
			}
		}
		lines.push(prefix + connector + label);

		if (entry.isDirectory && (maxDepth === null || currentDepth < maxDepth)) {
			const child = await buildTree(
				path.join(dirPath, entry.name),
				prefix + (isLast ? "   " : "│  "),
				includeFiles,
				includeHidden,
				maxDepth,
				currentDepth + 1,
			);
			if (child) lines.push(child);
		}
	}

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

async function generateTree(
	rootPath: string,
	includeFiles: boolean,
	includeHidden: boolean,
	maxDepth: number | null,
): Promise<string> {
	const tree = await buildTree(
		rootPath,
		"",
		includeFiles,
		includeHidden,
		maxDepth,
	);
	const { dirs, files } = await countEntries(
		rootPath,
		includeFiles,
		includeHidden,
	);
	return `${dirs} folder(s), ${files} file(s)\n${path.basename(rootPath)}\n${tree}`;
}

/**
 * Returns the deepest tree that fits within MAX_LINES.
 * Caches each passing result so there's no wasted regeneration.
 */
async function findSmartTree(
	rootPath: string,
	includeFiles: boolean,
	includeHidden: boolean,
): Promise<string> {
	let best = await generateTree(rootPath, includeFiles, includeHidden, 0);

	for (let depth = 1; depth <= 50; depth++) {
		const tree = await generateTree(
			rootPath,
			includeFiles,
			includeHidden,
			depth,
		);
		if (tree.split("\n").length > MAX_LINES) break;
		best = tree;
	}

	return best;
}

export default function fileTreeExtension(pi: ExtensionAPI) {
	pi.registerTool({
		name: "file_tree",
		label: "File Tree",
		description:
			"Returns a tree structure of the requested path with ASCII visualization. Useful for understanding project structure.",
		promptSnippet:
			"Display a visual tree structure of folders (and optionally files) at a given path",
		promptGuidelines: [
			"includeFiles: true for files + folders.",
			"includeHidden: only when user asks for dotfiles (.gitignore, .env).",
			"Omit depth unless you need a specific limit — auto-sizes to fit context.",
		],
		parameters: Type.Object({
			path: Type.String({
				description: "Path to generate tree for (relative or absolute)",
			}),
			includeFiles: Type.Optional(
				Type.Boolean({
					description:
						"Include files in the tree (default: false, folders only)",
				}),
			),
			includeHidden: Type.Optional(
				Type.Boolean({
					description:
						"Include hidden files/dirs starting with dot (default: false)",
				}),
			),
			depth: Type.Optional(
				Type.Number({
					description:
						"Max directory depth to traverse. When not set (default), auto-expands to fill context. Output is always capped to fit.",
				}),
			),
		}),

		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const includeFiles = params.includeFiles ?? false;
			const includeHidden = params.includeHidden ?? false;

			let targetPath = params.path;
			if (!path.isAbsolute(targetPath)) {
				targetPath = path.resolve(ctx.cwd, targetPath);
			}

			try {
				const stats = await fs.stat(targetPath);
				if (!stats.isDirectory()) {
					throw new Error(`Path is not a directory: ${params.path}`);
				}

				const tree =
					params.depth != null
						? await generateTree(
								targetPath,
								includeFiles,
								includeHidden,
								params.depth,
							)
						: await findSmartTree(targetPath, includeFiles, includeHidden);

				const allLines = tree.split("\n");
				const output =
					allLines.length > MAX_LINES
						? allLines.slice(0, MAX_LINES).join("\n") +
							"\n\n! Output truncated. Use a smaller depth to see more detail."
						: tree;

				return {
					content: [{ type: "text", text: output }],
					details: {
						path: targetPath,
						includeFiles,
						includeHidden,
						depth: params.depth ?? null,
					},
				};
			} catch (error) {
				throw new Error(
					`Failed to generate file tree: ${error instanceof Error ? error.message : String(error)}`,
				);
			}
		},
	});
}