import { mkdirSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs";
import { dirname, join } from "node:path";
import { getAgentDir } from "@earendil-works/pi-coding-agent";
import type { UsageRecord } from "./scan.ts";

/** 每百万 token 的美元单价，与 pi 自带定价表同单位 */
export interface Rate {
	input: number;
	output: number;
	cacheRead: number;
	cacheWrite: number;
	/** 阶梯：单次请求 prompt 量越过 inputTokensAbove 后改用这组单价 */
	tiers?: Tier[];
}

export interface Tier extends Omit<Rate, "tiers"> {
	inputTokensAbove: number;
}

/** 单价的出处，可信度依次递减 */
export type PriceSource = "manual" | "manual-wildcard" | "catalog" | "borrowed";

export interface ResolvedRate {
	rate: Rate;
	source: PriceSource;
	/** source 为 borrowed 时，单价实际借自哪个 provider */
	borrowedFrom?: string;
}

export const PRICE_SOURCE_LABELS: Record<PriceSource, string> = {
	manual: "手工配置",
	"manual-wildcard": "手工通配",
	catalog: "pi 定价表",
	borrowed: "借用同名模型",
};

export function channelKey(provider: string, model: string): string {
	return `${provider}/${model}`;
}

function isFiniteNonNegative(v: unknown): v is number {
	return typeof v === "number" && Number.isFinite(v) && v >= 0;
}

/** 把任意来源的对象收敛成 Rate；字段缺失按 0 计，非法值则拒绝整条 */
function toRate(raw: unknown, errors: string[], where: string): Rate | null {
	if (!raw || typeof raw !== "object") {
		errors.push(`${where}: 不是对象`);
		return null;
	}
	const o = raw as Record<string, unknown>;
	const fields = ["input", "output", "cacheRead", "cacheWrite"] as const;
	const rate: Rate = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };

	for (const f of fields) {
		if (o[f] === undefined) continue;
		if (!isFiniteNonNegative(o[f])) {
			errors.push(`${where}.${f}: 期望非负数字，实际是 ${JSON.stringify(o[f])}`);
			return null;
		}
		rate[f] = o[f];
	}

	if (Array.isArray(o.tiers)) {
		const tiers: Tier[] = [];
		for (const [i, t] of o.tiers.entries()) {
			const inner = toRate(t, errors, `${where}.tiers[${i}]`);
			if (!inner) return null;
			const above = (t as Record<string, unknown>).inputTokensAbove;
			if (!isFiniteNonNegative(above)) {
				errors.push(`${where}.tiers[${i}].inputTokensAbove: 期望非负数字`);
				return null;
			}
			tiers.push({ ...inner, inputTokensAbove: above });
		}
		rate.tiers = tiers;
	}

	return rate;
}

export interface PriceBookOptions {
	/** 用户配置里的 pricing 段，key 为 "provider/model" 或 "provider/*" */
	manual?: Record<string, Rate>;
	/** pi 目录：provider -> (model id -> Rate) */
	catalog?: Map<string, Map<string, Rate>>;
	/** 加载过程中发现的问题，原样展示给用户，避免改了配置却不知为何不生效 */
	errors?: string[];
	configPath?: string;
}

export class PriceBook {
	readonly errors: string[];
	readonly configPath: string;
	private readonly manual: Map<string, Rate>;
	private readonly manualWildcard: Map<string, Rate>;
	private readonly catalog: Map<string, Map<string, Rate>>;
	/** model id -> 提供该 id 的 provider 列表，按名字典序，保证借价结果可重现 */
	private readonly byModelId: Map<string, string[]>;
	private readonly cache = new Map<string, ResolvedRate | null>();

	constructor(opts: PriceBookOptions = {}) {
		this.errors = opts.errors ?? [];
		this.configPath = opts.configPath ?? "";
		this.catalog = opts.catalog ?? new Map();
		this.manual = new Map();
		this.manualWildcard = new Map();

		for (const [key, rate] of Object.entries(opts.manual ?? {})) {
			const slash = key.lastIndexOf("/");
			if (slash <= 0 || slash === key.length - 1) {
				this.errors.push(`pricing["${key}"]: key 应形如 "provider/model" 或 "provider/*"`);
				continue;
			}
			const provider = key.slice(0, slash);
			const model = key.slice(slash + 1);
			if (model === "*") this.manualWildcard.set(provider, rate);
			else this.manual.set(key, rate);
		}

		this.byModelId = new Map();
		for (const [provider, models] of [...this.catalog].sort((a, b) => a[0].localeCompare(b[0]))) {
			for (const id of models.keys()) {
				const list = this.byModelId.get(id);
				if (list) list.push(provider);
				else this.byModelId.set(id, [provider]);
			}
		}
	}

	/** 四级优先：手工精确 > 手工通配 > 目录精确 > 跨 provider 借价 */
	resolve(provider: string, model: string): ResolvedRate | null {
		const key = channelKey(provider, model);
		const cached = this.cache.get(key);
		if (cached !== undefined) return cached;

		const exactManual = this.manual.get(key);
		const wildcard = this.manualWildcard.get(provider);

		let result: ResolvedRate | null;
		if (exactManual) result = { rate: exactManual, source: "manual" };
		else if (wildcard) result = { rate: wildcard, source: "manual-wildcard" };
		else result = this.resolveFromCatalog(provider, model);

		this.cache.set(key, result);
		return result;
	}

	/**
	 * 只查 pi 的模型目录，忽略手工配置。
	 * 生成官方价快照时用它，避免把用户改过的价当成官方价写回文件。
	 */
	resolveFromCatalog(provider: string, model: string): ResolvedRate | null {
		const exact = this.catalog.get(provider)?.get(model);
		if (exact) return { rate: exact, source: "catalog" };

		const from = this.byModelId.get(model)?.[0];
		const borrowed = from ? this.catalog.get(from)?.get(model) : undefined;
		return from && borrowed ? { rate: borrowed, source: "borrowed", borrowedFrom: from } : null;
	}

	/**
	 * 单条记录的金额。返回 null 表示该渠道无单价——注意这不等于 0。
	 * 阶梯按本次请求自身的 prompt 量判定，因此必须逐条计算后再累加。
	 */
	cost(record: UsageRecord): number | null {
		const resolved = this.resolve(record.provider, record.model);
		if (!resolved) return null;

		const promptTokens = record.input + record.cacheRead + record.cacheWrite;
		let rates: Rate | Tier = resolved.rate;
		let matched = -1;
		for (const tier of resolved.rate.tiers ?? []) {
			if (promptTokens > tier.inputTokensAbove && tier.inputTokensAbove > matched) {
				rates = tier;
				matched = tier.inputTokensAbove;
			}
		}

		return (
			(rates.input * record.input +
				rates.output * record.output +
				rates.cacheRead * record.cacheRead +
				rates.cacheWrite * record.cacheWrite) /
			1_000_000
		);
	}
}

function readJson(path: string, errors: string[]): unknown {
	try {
		return JSON.parse(readFileSync(path, "utf8"));
	} catch (err: any) {
		// 文件不存在是正常状态（没配过价），其余错误必须暴露
		if (err?.code !== "ENOENT") errors.push(`${path}: ${err?.message ?? String(err)}`);
		return null;
	}
}

export function parseManualPricing(raw: unknown, errors: string[]): Record<string, Rate> {
	const pricing = (raw as any)?.pricing;
	if (pricing === undefined) return {};
	if (!pricing || typeof pricing !== "object" || Array.isArray(pricing)) {
		errors.push("config.pricing: 期望是对象");
		return {};
	}

	const out: Record<string, Rate> = {};
	for (const [key, value] of Object.entries(pricing)) {
		const rate = toRate(value, errors, `pricing["${key}"]`);
		if (rate) out[key] = rate;
	}
	return out;
}

export function parseCatalog(raw: unknown): Map<string, Map<string, Rate>> {
	const catalog = new Map<string, Map<string, Rate>>();
	if (!raw || typeof raw !== "object") return catalog;

	for (const [provider, entry] of Object.entries(raw as Record<string, any>)) {
		const models = new Map<string, Rate>();
		for (const m of entry?.models ?? []) {
			// 目录由 pi 维护，格式问题不该打扰用户，静默跳过即可
			const rate = toRate(m?.cost, [], "");
			if (m?.id && rate) models.set(m.id, rate);
		}
		if (models.size > 0) catalog.set(provider, models);
	}
	return catalog;
}

export function configPathFor(agentDir = getAgentDir()): string {
	return join(agentDir, "extensions", "pi-token-use", "config.json");
}

/**
 * 读出整份配置以便原地改写。
 * 文件存在但无法解析时**抛错而不是从空对象重来**——那会把 AI 批量写入的整份价格表冲掉。
 */
function readConfigForWrite(path: string): Record<string, any> {
	let text: string;
	try {
		text = readFileSync(path, "utf8");
	} catch (err: any) {
		if (err?.code === "ENOENT") return {};
		throw err;
	}

	let parsed: unknown;
	try {
		parsed = JSON.parse(text);
	} catch (err: any) {
		throw new Error(`配置文件无法解析，已拒绝写入以免丢失原有内容：${err?.message ?? String(err)}`);
	}

	if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
		throw new Error("配置文件顶层不是对象，已拒绝写入以免丢失原有内容");
	}
	return parsed as Record<string, any>;
}

/**
 * 先写临时文件再 rename。rename 在同一文件系统上是原子的，
 * 因此写入被中断（崩溃、磁盘满）时要么是旧内容要么是新内容，
 * 不会留下截断的 JSON 把整份价格表毁掉。
 */
function writeConfig(path: string, config: Record<string, any>): void {
	mkdirSync(dirname(path), { recursive: true });
	const tmp = `${path}.tmp`;
	writeFileSync(tmp, `${JSON.stringify(config, null, 2)}\n`, "utf8");
	try {
		renameSync(tmp, path);
	} catch (err) {
		try {
			rmSync(tmp, { force: true });
		} catch {
			// 清理失败无关紧要，真正的错误是下面这个
		}
		throw err;
	}
}

/** 写入单个渠道的手工单价，文件里的其余内容原样保留 */
export function saveManualRate(channel: string, rate: Rate, agentDir = getAgentDir()): void {
	const path = configPathFor(agentDir);
	const config = readConfigForWrite(path);
	const pricing = config.pricing && typeof config.pricing === "object" ? config.pricing : {};
	config.pricing = { ...pricing, [channel]: rate };
	writeConfig(path, config);
}

/** 移除单个渠道的手工单价，使其回退到自动匹配 */
export function removeManualRate(channel: string, agentDir = getAgentDir()): void {
	const path = configPathFor(agentDir);
	const config = readConfigForWrite(path);
	if (!config.pricing || typeof config.pricing !== "object") return;
	if (!(channel in config.pricing)) return;

	const { [channel]: _removed, ...rest } = config.pricing as Record<string, unknown>;
	config.pricing = rest;
	writeConfig(path, config);
}

const CATALOG_NOTE =
	"catalog 段由插件在每次打开面板时自动同步为 pi 的官方定价，直接修改它不会生效；要覆盖某个渠道的价格，请写到 pricing 段。";

/**
 * 官方价快照：**忽略手工配置**，纯粹反映 pi 模型目录的解析结果（含跨 provider 借价）。
 * 目的是让配置文件里能直接看到官方价，方便对照修改或让 AI 参考。
 */
export function buildCatalogSnapshot(
	channels: Array<{ provider: string; model: string }>,
	book: PriceBook,
): Record<string, Rate> {
	const out: Record<string, Rate> = {};
	for (const { provider, model } of channels) {
		const resolved = book.resolveFromCatalog(provider, model);
		// 解析不出单价的渠道不写入：写成 0 会把「未知」伪装成「免费」
		if (!resolved) continue;
		// 快照里只留四个基础单价。阶梯规则仍由 pi 的模型目录提供并参与计价，
		// 放进这里只会让文件难读——手工价本来就是整条替换、不带阶梯的
		const { tiers: _tiers, ...flat } = resolved.rate;
		out[channelKey(provider, model)] = flat;
	}
	return out;
}

/**
 * 把官方价快照同步进配置文件的 catalog 段，pricing 段原样保留。
 * 配置文件损坏时抛错由调用方决定如何提示——绝不覆盖写，以免毁掉已有内容。
 */
export function syncCatalogSnapshot(snapshot: Record<string, Rate>, agentDir = getAgentDir()): void {
	const path = configPathFor(agentDir);
	const config = readConfigForWrite(path);

	if (JSON.stringify(config.catalog) === JSON.stringify(snapshot) && config._note === CATALOG_NOTE) {
		return; // 内容没变就不碰文件，避免每次打开都刷新 mtime
	}

	config._note = CATALOG_NOTE;
	config.catalog = snapshot;
	config.pricing ??= {};
	writeConfig(path, config);
}

/** 每次打开面板时重新加载，改完配置文件重开面板即生效 */
export function loadPriceBook(agentDir = getAgentDir()): PriceBook {
	const errors: string[] = [];
	const configPath = configPathFor(agentDir);
	const manual = parseManualPricing(readJson(configPath, errors), errors);
	const catalog = parseCatalog(readJson(join(agentDir, "models-store.json"), errors));
	return new PriceBook({ manual, catalog, errors, configPath });
}