pi-map/src/init.ts

import { discoverProject, type DirectoryEntry } from "./discover.js";

export { discoverProject };
import type { FileEntry } from "./directory-model.js";
import { renderDirectoryIndex, renderDirectoryMap } from "./format.js";
import { extractFileLLM, extractPackageLLM } from "./llm/llm-extract.js";
import { extractFileAST } from "./ast/ast-extract.js";
import { mergeFileData } from "./merge.js";
import { processFiles } from "./llm/llm-batch.js";
import { writeFileSync } from "fs";
import { join } from "path";
import type { LLMClient } from "./llm/llm-client.js";
import {
	createDirectoryModel,
	type DirectoryArtifactModel,
	type RoutingMetadataOptions,
} from "./directory-model.js";
import { populateRoutingMetadata } from "./routing-metadata.js";
import { loadConfig } from "./config.js";

export interface ProgressInfo {
	message: string;
	completed: number;
	total: number;
	currentFile?: string;
	dir?: string;
}

export interface InitOptions {
	verbose?: boolean;
	llmClient?: LLMClient;
	cacheDir?: string;
	onProgress?: (info: ProgressInfo) => void;
	tagCap?: number;
	workflowHintCap?: number;
}

export async function initProject(
	rootPath: string,
	options: InitOptions = {},
): Promise<void> {
	const entries = discoverProject(rootPath);
	const totalFiles = entries.reduce((sum, e) => sum + e.files.length, 0);
	let globalCompleted = 0;

	// Load project config and apply defaults when not overridden in options
	const config = loadConfig(rootPath);
	const routingOpts: RoutingMetadataOptions = {
		tagCap: options.tagCap ?? config.tagCap,
		workflowHintCap: options.workflowHintCap ?? config.workflowHintCap,
	};

	options.onProgress?.({
		message: `Scanning ${entries.length} directories (${totalFiles} files)...`,
		completed: 0,
		total: totalFiles,
	});

	// Build directory relationships
	const dirSet = new Set(entries.map((e) => e.relativePath));
	const parentMap = buildParentMap(entries);
	const childrenMap = buildChildrenMap(entries);

	for (let i = 0; i < entries.length; i++) {
		const entry = entries[i];
		await generateDirectoryArtifacts(
			entry,
			{
				dirSet,
				parentMap,
				childrenMap,
				isRoot: entry.relativePath === ".",
			},
			options.llmClient,
			options.cacheDir,
			(info) => {
				globalCompleted =
					info.completed +
					entries.slice(0, i).reduce((sum, e) => sum + e.files.length, 0);
				options.onProgress?.({
					...info,
					completed: globalCompleted,
					total: totalFiles,
					dir: entry.relativePath,
				});
			},
			routingOpts,
		);
	}

	options.onProgress?.({
		message: `Generated ${entries.length} directory map/index pairs`,
		completed: totalFiles,
		total: totalFiles,
	});
	if (options.verbose !== false) {
		console.log(`Generated ${entries.length} directory map/index pairs`);
	}
}

export interface DirectoryContext {
	dirSet: Set<string>;
	parentMap: Map<string, string | undefined>;
	childrenMap: Map<string, string[]>;
	isRoot: boolean;
}

export type ArtifactWriteMode = "both" | "map" | "index";

export function buildDirectoryContext(
	entries: DirectoryEntry[],
	targetEntry: DirectoryEntry,
): DirectoryContext {
	return {
		dirSet: new Set(entries.map((e) => e.relativePath)),
		parentMap: buildParentMap(entries),
		childrenMap: buildChildrenMap(entries),
		isRoot: targetEntry.relativePath === ".",
	};
}

function buildParentMap(
	entries: DirectoryEntry[],
): Map<string, string | undefined> {
	const map = new Map<string, string | undefined>();
	for (const entry of entries) {
		const rel = entry.relativePath;
		if (rel === ".") {
			map.set(rel, undefined);
		} else {
			const lastSep = rel.lastIndexOf("/");
			map.set(rel, lastSep >= 0 ? rel.slice(0, lastSep) : ".");
		}
	}
	return map;
}

function buildChildrenMap(entries: DirectoryEntry[]): Map<string, string[]> {
	const map = new Map<string, string[]>();
	for (const entry of entries) {
		map.set(entry.relativePath, []);
	}
	for (const entry of entries) {
		const rel = entry.relativePath;
		if (rel === ".") continue;
		const lastSep = rel.lastIndexOf("/");
		const parent = lastSep >= 0 ? rel.slice(0, lastSep) : ".";
		const siblings = map.get(parent);
		if (siblings) {
			siblings.push(rel);
		}
	}
	return map;
}

export async function buildDirectoryArtifactModel(
	entry: DirectoryEntry,
	ctx: DirectoryContext,
	llmClient?: LLMClient,
	cacheDir?: string,
	onProgress?: (info: ProgressInfo) => void,
	routingOpts?: RoutingMetadataOptions,
): Promise<DirectoryArtifactModel> {
	const fileData = await processFiles(
		entry.files,
		async (file) => {
			const filePath = join(entry.dirPath, file);
			const llmData = await extractFileLLM(filePath, llmClient, cacheDir);
			const astData = await extractFileAST(filePath);
			return mergeFileData(file, llmData, astData);
		},
		{ concurrency: 4, batchDelayMs: 100, maxRetries: 2 },
		(completed, total, currentFile) => {
			onProgress?.({
				message: `  → ${currentFile}`,
				completed,
				total,
				currentFile,
				dir: entry.relativePath,
			});
		},
	);

	const packageData = await extractPackageLLM(
		entry.relativePath,
		fileData,
		llmClient,
		cacheDir,
	);

	const model: DirectoryArtifactModel = createDirectoryModel({
		dir: entry.relativePath,
		role: packageData.role,
		files: fileData,
		arch: packageData.arch,
		parent: ctx.parentMap.get(entry.relativePath),
		children: ctx.childrenMap.get(entry.relativePath) ?? [],
		isRoot: ctx.isRoot,
		dirty: "-",
	});

	populateRoutingMetadata(model, routingOpts);
	return model;
}

export function writeDirectoryArtifacts(
	entry: DirectoryEntry,
	model: DirectoryArtifactModel,
	writeMode: ArtifactWriteMode = "both",
): void {
	const mapPath = join(entry.dirPath, ".pi-map.md");
	const indexPath = join(entry.dirPath, ".pi-map.index.md");

	if (writeMode === "both" || writeMode === "map") {
		writeFileSync(mapPath, renderDirectoryMap(model));
	}
	if (writeMode === "both" || writeMode === "index") {
		writeFileSync(indexPath, renderDirectoryIndex(model));
	}
}

export async function generateDirectoryArtifacts(
	entry: DirectoryEntry,
	ctx: DirectoryContext,
	llmClient?: LLMClient,
	cacheDir?: string,
	onProgress?: (info: ProgressInfo) => void,
	routingOpts?: RoutingMetadataOptions,
	writeMode: ArtifactWriteMode = "both",
): Promise<FileEntry[]> {
	const model = await buildDirectoryArtifactModel(
		entry,
		ctx,
		llmClient,
		cacheDir,
		onProgress,
		routingOpts,
	);
	writeDirectoryArtifacts(entry, model, writeMode);
	return model.files;
}

export async function reinitPath(
	path: string,
	options: InitOptions = {},
): Promise<void> {
	// Full regeneration clears all dirty markers by overwriting every map/index pair
	await initProject(path, options);
}

// Backward-compatible wrapper for patch/validate compatibility
export async function generateDirectoryMap(
	entry: DirectoryEntry,
	llmClient?: LLMClient,
	cacheDir?: string,
): Promise<FileEntry[]> {
	const entries = [entry];
	const ctx = buildDirectoryContext(entries, entry);
	const config = loadConfig(entry.dirPath);
	const routingOpts: RoutingMetadataOptions = {
		tagCap: config.tagCap,
		workflowHintCap: config.workflowHintCap,
	};
	return generateDirectoryArtifacts(
		entry,
		ctx,
		llmClient,
		cacheDir,
		undefined,
		routingOpts,
	);
}