import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import { DynamicBorder, getMarkdownTheme } from "@earendil-works/pi-coding-agent";
import type { AutocompleteItem } from "@earendil-works/pi-tui";
import { Container, Markdown, matchesKey, Text } from "@earendil-works/pi-tui";
import { existsSync, readdirSync, readFileSync, statSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";

type Usage = {
	input?: number;
	output?: number;
	cacheRead?: number;
	cacheWrite?: number;
	totalTokens?: number;
	cost?: {
		input?: number;
		output?: number;
		cacheRead?: number;
		cacheWrite?: number;
		total?: number;
	};
};

type AssistantMessage = {
	role?: string;
	provider?: string;
	model?: string;
	usage?: Usage;
	timestamp?: number;
};

type SessionLine = {
	type?: string;
	timestamp?: string;
	message?: AssistantMessage;
};

type ModelStats = {
	provider: string;
	model: string;
	requests: number;
	input: number;
	output: number;
	cacheRead: number;
	cacheWrite: number;
	totalTokens: number;
	costInput: number;
	costOutput: number;
	costCacheRead: number;
	costCacheWrite: number;
	costTotal: number;
};

type AggregateStats = {
	days: number;
	fromTime?: number;
	toTime: number;
	sessionDir: string;
	filesScanned: number;
	filesMatched: number;
	requests: number;
	input: number;
	output: number;
	cacheRead: number;
	cacheWrite: number;
	totalTokens: number;
	costInput: number;
	costOutput: number;
	costCacheRead: number;
	costCacheWrite: number;
	costTotal: number;
	models: Map<string, ModelStats>;
};

const EXTENSION_NAME = "token-stats";

const getSessionDir = (): string => {
	if (process.env.PI_CODING_AGENT_SESSION_DIR) {
		return process.env.PI_CODING_AGENT_SESSION_DIR;
	}

	const configDir = process.env.PI_CODING_AGENT_DIR ?? join(homedir(), ".pi", "agent");
	return join(configDir, "sessions");
};

const walkJsonlFiles = (dir: string): string[] => {
	if (!existsSync(dir)) {
		return [];
	}

	const result: string[] = [];
	const stack = [dir];

	while (stack.length > 0) {
		const current = stack.pop()!;
		let entries: string[] = [];
		try {
			entries = readdirSync(current);
		} catch {
			continue;
		}

		for (const entry of entries) {
			const fullPath = join(current, entry);
			let stats;
			try {
				stats = statSync(fullPath);
			} catch {
				continue;
			}

			if (stats.isDirectory()) {
				stack.push(fullPath);
			} else if (stats.isFile() && entry.endsWith(".jsonl")) {
				result.push(fullPath);
			}
		}
	}

	return result;
};

const parseDays = (args: string | undefined): number => {
	const text = (args ?? "").trim();
	if (!text) {
		return 30;
	}

	const flagMatch = text.match(/(?:^|\s)(?:--days|-d)\s+(-?\d+)(?:\s|$)/i);
	const compactMatch = text.match(/^(-?\d+)\s*d?$/i);
	const valueText = flagMatch?.[1] ?? compactMatch?.[1];
	const value = Number.parseInt(valueText ?? "", 10);

	if (!Number.isFinite(value) || value < 0) {
		return 30;
	}

	return value;
};

const addUsage = (stats: AggregateStats, provider: string, model: string, usage: Usage) => {
	const input = usage.input ?? 0;
	const output = usage.output ?? 0;
	const cacheRead = usage.cacheRead ?? 0;
	const cacheWrite = usage.cacheWrite ?? 0;
	const totalTokens = usage.totalTokens ?? input + output + cacheRead + cacheWrite;
	const costInput = usage.cost?.input ?? 0;
	const costOutput = usage.cost?.output ?? 0;
	const costCacheRead = usage.cost?.cacheRead ?? 0;
	const costCacheWrite = usage.cost?.cacheWrite ?? 0;
	const costTotal = usage.cost?.total ?? costInput + costOutput + costCacheRead + costCacheWrite;

	stats.requests += 1;
	stats.input += input;
	stats.output += output;
	stats.cacheRead += cacheRead;
	stats.cacheWrite += cacheWrite;
	stats.totalTokens += totalTokens;
	stats.costInput += costInput;
	stats.costOutput += costOutput;
	stats.costCacheRead += costCacheRead;
	stats.costCacheWrite += costCacheWrite;
	stats.costTotal += costTotal;

	const key = `${provider}/${model}`;
	let modelStats = stats.models.get(key);
	if (!modelStats) {
		modelStats = {
			provider,
			model,
			requests: 0,
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: 0,
			costInput: 0,
			costOutput: 0,
			costCacheRead: 0,
			costCacheWrite: 0,
			costTotal: 0,
		};
		stats.models.set(key, modelStats);
	}

	modelStats.requests += 1;
	modelStats.input += input;
	modelStats.output += output;
	modelStats.cacheRead += cacheRead;
	modelStats.cacheWrite += cacheWrite;
	modelStats.totalTokens += totalTokens;
	modelStats.costInput += costInput;
	modelStats.costOutput += costOutput;
	modelStats.costCacheRead += costCacheRead;
	modelStats.costCacheWrite += costCacheWrite;
	modelStats.costTotal += costTotal;
};

const collectStats = (days: number): AggregateStats => {
	const now = Date.now();
	const fromTime = days === 0 ? undefined : now - days * 24 * 60 * 60 * 1000;
	const sessionDir = getSessionDir();
	const files = walkJsonlFiles(sessionDir);
	const stats: AggregateStats = {
		days,
		fromTime,
		toTime: now,
		sessionDir,
		filesScanned: files.length,
		filesMatched: 0,
		requests: 0,
		input: 0,
		output: 0,
		cacheRead: 0,
		cacheWrite: 0,
		totalTokens: 0,
		costInput: 0,
		costOutput: 0,
		costCacheRead: 0,
		costCacheWrite: 0,
		costTotal: 0,
		models: new Map(),
	};

	for (const file of files) {
		let matchedInFile = false;
		let content = "";
		try {
			content = readFileSync(file, "utf8");
		} catch {
			continue;
		}

		for (const rawLine of content.split("\n")) {
			const line = rawLine.trim();
			if (!line) {
				continue;
			}

			let entry: SessionLine;
			try {
				entry = JSON.parse(line) as SessionLine;
			} catch {
				continue;
			}

			if (entry.type !== "message" || entry.message?.role !== "assistant" || !entry.message.usage) {
				continue;
			}

			const timestamp = entry.message.timestamp ?? (entry.timestamp ? Date.parse(entry.timestamp) : undefined);
			if (fromTime !== undefined && timestamp !== undefined && timestamp < fromTime) {
				continue;
			}

			const provider = entry.message.provider ?? "unknown";
			const model = entry.message.model ?? "unknown";
			addUsage(stats, provider, model, entry.message.usage);
			matchedInFile = true;
		}

		if (matchedInFile) {
			stats.filesMatched += 1;
		}
	}

	return stats;
};

const formatNumber = (value: number): string => new Intl.NumberFormat("en-US").format(Math.round(value));
const formatCompactTokens = (value: number): string => {
	const rounded = Math.round(value);
	if (rounded >= 100_000_000) {
		return `${(rounded / 100_000_000).toFixed(1)}亿`;
	}
	if (rounded >= 10_000) {
		return `${(rounded / 10_000).toFixed(1)}万`;
	}
	return formatNumber(rounded);
};
const formatCost = (value: number): string => `$${value.toFixed(6)}`;
const formatDateTime = (value: number): string => {
	const date = new Date(value);
	const pad = (part: number): string => part.toString().padStart(2, "0");
	return `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())} ${pad(date.getHours())}:${pad(date.getMinutes())}`;
};

const tableRow = (cells: Array<string | number>): string => `| ${cells.join(" | ")} |`;

const buildMarkdown = (stats: AggregateStats): string => {
	const range = stats.days === 0 ? "全部" : `最近 ${stats.days} 天（${formatDateTime(stats.fromTime!)} ~ ${formatDateTime(stats.toTime)}）`;
	const rows = [...stats.models.values()].sort((a, b) => b.totalTokens - a.totalTokens);

	const lines: string[] = [
		"# Pi Token 使用统计",
		"",
		`- 查询范围：${range}`,
		`- Session 目录：\`${stats.sessionDir}\``,
		`- 扫描文件：${formatNumber(stats.filesScanned)} 个，命中会话：${formatNumber(stats.filesMatched)} 个`,
		`- 模型请求：${formatNumber(stats.requests)} 次`,
		`- Token 总计：${formatNumber(stats.totalTokens)}`,
		`- 费用总计：${formatCost(stats.costTotal)}`,
		"",
		"## 总览",
		"",
		tableRow(["Input", "Output", "Cache Read", "Cache Write", "Total Tokens", "消耗 Tokens", "Cost"]),
		tableRow(["---:", "---:", "---:", "---:", "---:", "---:", "---:"]),
		tableRow([
			formatNumber(stats.input),
			formatNumber(stats.output),
			formatNumber(stats.cacheRead),
			formatNumber(stats.cacheWrite),
			formatNumber(stats.totalTokens),
			formatCompactTokens(stats.totalTokens),
			formatCost(stats.costTotal),
		]),
		"",
		"## 按模型统计",
		"",
	];

	if (rows.length === 0) {
		lines.push("未找到匹配范围内的 token 使用记录。");
		return lines.join("\n");
	}

	lines.push(tableRow(["模型", "请求", "Input", "Output", "Cache Read", "Cache Write", "Total", "消耗 Tokens", "Cost"]));
	lines.push(tableRow(["---", "---:", "---:", "---:", "---:", "---:", "---:", "---:", "---:"]));
	for (const row of rows) {
		lines.push(
			tableRow([
				`\`${row.provider}/${row.model}\``,
				formatNumber(row.requests),
				formatNumber(row.input),
				formatNumber(row.output),
				formatNumber(row.cacheRead),
				formatNumber(row.cacheWrite),
				formatNumber(row.totalTokens),
				formatCompactTokens(row.totalTokens),
				formatCost(row.costTotal),
			]),
		);
	}

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

const showMarkdown = async (markdown: string, ctx: ExtensionCommandContext) => {
	if (ctx.mode !== "tui") {
		console.log(markdown);
		return;
	}

	await ctx.ui.custom((_tui, theme, _kb, done) => {
		const container = new Container();
		const border = new DynamicBorder((s: string) => theme.fg("accent", s));
		const mdTheme = getMarkdownTheme();

		container.addChild(border);
		container.addChild(new Text(theme.fg("accent", theme.bold("Pi Token Stats")), 1, 0));
		container.addChild(new Markdown(markdown, 1, 1, mdTheme));
		container.addChild(new Text(theme.fg("dim", "Press Enter or Esc to close"), 1, 0));
		container.addChild(border);

		return {
			render: (width: number) => container.render(width),
			invalidate: () => container.invalidate(),
			handleInput: (data: string) => {
				if (matchesKey(data, "enter") || matchesKey(data, "escape")) {
					done(undefined);
				}
			},
		};
	});
};

const completions: AutocompleteItem[] = [
	{ value: "7", label: "7", description: "统计最近 7 天" },
	{ value: "30", label: "30", description: "统计最近 30 天（默认）" },
	{ value: "0", label: "0", description: "统计全部历史" },
];

export default function (pi: ExtensionAPI) {
	const command = {
		description: "统计 pi session token/cost 用量，参数为天数：默认 30，0 表示全部",
		getArgumentCompletions: (prefix: string): AutocompleteItem[] | null => {
			const filtered = completions.filter((item) => item.value.startsWith(prefix.trim()));
			return filtered.length > 0 ? filtered : null;
		},
		handler: async (args: string | undefined, ctx: ExtensionCommandContext) => {
			const days = parseDays(args);
			if (ctx.hasUI) {
				ctx.ui.notify(`正在统计 token 用量（${days === 0 ? "全部" : `最近 ${days} 天`}）...`, "info");
			}

			const stats = collectStats(days);
			const markdown = buildMarkdown(stats);
			await showMarkdown(markdown, ctx);
		},
	};

	pi.registerCommand("tokens", command);
	pi.registerCommand("token-stats", command);

}