import * as fs from "node:fs/promises";
import * as os from "node:os";
import * as path from "node:path";
import { type } from "@oh-my-pi/omptype";
import type { AgentToolResult } from "@oh-my-pi/pi-agent-core";
import {
	type GeneratedImage,
	generateImage,
	type ImageGenerationRequest,
	type ImageGenerationResult,
	isImageGenerationApi,
	type Model,
	parseAntigravityCredentials,
} from "@oh-my-pi/pi-ai";
import { ProviderHttpError } from "@oh-my-pi/pi-ai/error";
import { getEffectiveTokenCost } from "@oh-my-pi/pi-catalog/models";
import {
	isEnoent,
	logger,
	parseImageMetadata,
	prompt,
	ptree,
	Snowflake,
	truncate,
	untilAborted,
} from "@oh-my-pi/pi-utils";
import { type RoleChainCandidate, resolveModelRoleValue, resolveRoleChain } from "../config/model-resolver";
import { roleCandidatePool } from "../config/model-roles";
import { isAuthenticated, type ModelRegistry } from "../config/model-registry";
import { settings } from "../config/settings";
import type { CustomTool } from "../extensibility/custom-tools/types";
import imageGenDescription from "../prompts/tools/image-gen.md" with { type: "text" };
import { resolveConfiguredModelTarget } from "../session/role-models";
import { resolveReadPath, specialFileKind } from "./path-utils";

const IMAGE_TIMEOUT = 3 * 60 * 1000;
const MAX_IMAGE_SIZE = 35 * 1024 * 1024;
/** 单条失败原因的字符上限：网关 HTML 错误页等原始响应体不应整段进入模型与 TUI。 */
const MAX_FAILURE_REASON_CHARS = 300;

const aspectRatioSchema = type('"1:1" | "3:4" | "4:3" | "9:16" | "16:9" | "3:2" | "2:3"').describe("aspect ratio");
const imageSizeSchema = type('"1024x1024" | "1536x1024" | "1024x1536"').describe("image size");
const inputImageSchema = type({
	"path?": type("string").describe("input image path"),
	"data?": type("string").describe("base64 image data"),
	"mime_type?": type("string").describe("mime type"),
});

export const imageGenSchema = type({
	subject: type("string").describe("main subject"),
	"action?": type("string").describe("what subject is doing"),
	"scene?": type("string").describe("location or environment"),
	"composition?": type("string").describe("camera angle and framing"),
	"lighting?": type("string").describe("lighting setup"),
	"style?": type("string").describe("artistic style"),
	"text?": type("string").describe("text to render"),
	"changes?": type("string[]").describe("edits to make"),
	"aspect_ratio?": aspectRatioSchema,
	"image_size?": imageSizeSchema,
	"input?": inputImageSchema.array().describe("input images"),
	"model?": type("string").describe("image model selector for this request"),
});
export type ImageGenParams = typeof imageGenSchema.infer;

interface ImageGenToolDetails {
	provider: string;
	model: string;
	imageCount: number;
	imagePaths: string[];
	images: GeneratedImage[];
	responseText?: string;
	usage?: ImageGenerationResult["usage"];
}

interface ImageInput {
	path?: string;
	data?: string;
	mime_type?: string;
}

function assemblePrompt(params: ImageGenParams): string {
	const parts: string[] = [];
	const subjectParts = [params.subject];
	if (params.action) subjectParts.push(params.action);
	if (params.scene) subjectParts.push(params.scene);
	parts.push(subjectParts.join(", "));
	if (params.composition) parts.push(params.composition);
	if (params.lighting) parts.push(params.lighting);
	if (params.style) parts.push(params.style);
	let result = `${parts.map(part => part.replace(/[.!,;:]+$/, "")).join(". ")}.`;
	if (params.text) result += `\n\nText: ${params.text}`;
	if (params.changes?.length) result += `\n\nChanges:\n${params.changes.map(change => `- ${change}`).join("\n")}`;
	return result;
}

function normalizeDataUrl(data: string): { data: string; mimeType?: string } {
	const match = data.match(/^data:([^;]+);base64,(.+)$/);
	if (!match) return { data };
	return { data: match[2] ?? "", mimeType: match[1] };
}

async function loadImageFromPath(imagePath: string, cwd: string): Promise<{ data: string; mimeType: string }> {
	const resolved = resolveReadPath(imagePath, cwd);
	try {
		const stat = await fs.stat(resolved);
		const kind = specialFileKind(stat);
		if (kind) throw new Error(`Cannot load '${imagePath}': it is a ${kind}, not a regular file or directory.`);
		if (stat.size > MAX_IMAGE_SIZE) throw new Error(`Image file too large: ${imagePath}`);
		const buffer = await Bun.file(resolved).bytes();
		if (buffer.length > MAX_IMAGE_SIZE) throw new Error(`Image file too large: ${imagePath}`);
		const mimeType = parseImageMetadata(buffer)?.mimeType;
		if (!mimeType) throw new Error(`Unsupported image type: ${imagePath}`);
		return { data: buffer.toBase64(), mimeType };
	} catch (error) {
		if (isEnoent(error)) throw new Error(`Image file not found: ${imagePath}`);
		throw error;
	}
}

async function resolveInputImage(input: ImageInput, cwd: string): Promise<{ data: string; mimeType: string }> {
	if (input.path) return loadImageFromPath(input.path, cwd);
	if (input.data) {
		const normalized = normalizeDataUrl(input.data.trim());
		const mimeType = normalized.mimeType ?? input.mime_type;
		if (!mimeType) throw new Error("mime_type is required when providing raw base64 data.");
		if (!normalized.data) throw new Error("Image data is empty.");
		return { data: normalized.data, mimeType };
	}
	throw new Error("input entries must include either path or data.");
}

function imageExtension(mimeType: string): string {
	const extensions: Record<string, string> = {
		"image/png": "png",
		"image/jpeg": "jpg",
		"image/gif": "gif",
		"image/webp": "webp",
	};
	return extensions[mimeType] ?? "png";
}

async function saveImagesToTemp(images: ImageGenerationResult["images"]): Promise<string[]> {
	return Promise.all(
		images.map(async image => {
			const filepath = path.join(os.tmpdir(), `omp-image-${Snowflake.next()}.${imageExtension(image.mimeType)}`);
			await Bun.write(filepath, Buffer.from(image.data, "base64"));
			return filepath;
		}),
	);
}

/**
 * Chat model that relays a hosted Responses image model's `image_generation`
 * call: the session model when it can carry the tool on the same provider,
 * otherwise the cheapest available carrier there.
 */
function resolveHostedImageCarrier(
	modelRegistry: ModelRegistry,
	selectedModel: Model,
	activeModel: Model | undefined,
): Model | undefined {
	if (activeModel?.provider === selectedModel.provider && activeModel.hostedImage) return activeModel;
	let carrier: Model | undefined;
	let carrierInput = Infinity;
	for (const model of modelRegistry.getAvailable()) {
		if (model.provider !== selectedModel.provider || !model.hostedImage) continue;
		const input = getEffectiveTokenCost(model.cost).input;
		if (!carrier || input < carrierInput) {
			carrier = model;
			carrierInput = input;
		}
	}
	return carrier;
}

/**
 * Default image candidates: explicit role entries first, then the session's
 * own provider (its `imageModel` swap, then the session model itself when it
 * carries the hosted image tool), then the remaining built-in chain.
 */
function defaultImageCandidates(chain: RoleChainCandidate[], sessionModel: Model | undefined, pool: Model[]): Model[] {
	const sessionCandidates = sessionModel
		? [
				resolveConfiguredModelTarget(sessionModel.imageModel, sessionModel, pool),
				sessionModel.hostedImage ? sessionModel : undefined,
			]
		: [];
	const ordered = [
		...chain.filter(candidate => candidate.explicit).map(candidate => candidate.model),
		...sessionCandidates,
		...chain.filter(candidate => !candidate.explicit).map(candidate => candidate.model),
	];
	const seen = new Set<string>();
	return ordered.filter((model): model is Model => {
		if (!model) return false;
		const key = `${model.provider}/${model.id}`;
		if (seen.has(key)) return false;
		seen.add(key);
		return true;
	});
}

async function buildToolResult(
	model: Model,
	result: ImageGenerationResult,
): Promise<AgentToolResult<ImageGenToolDetails, ImageGenParams>> {
	const imagePaths = await saveImagesToTemp(result.images);
	// Hosted backends may run a different model than the selected catalog entry; report what actually ran.
	const ranModel = result.model ?? model.id;
	if (imagePaths.length === 0) {
		return {
			content: [{ type: "text", text: `No image data returned.${result.text ? `\n\n${result.text}` : ""}` }],
			details: {
				provider: model.provider,
				model: ranModel,
				imageCount: 0,
				imagePaths: [],
				images: [],
				responseText: result.text,
				usage: result.usage,
			},
		};
	}
	const modelLine =
		ranModel === model.id ? `Model: ${ranModel}` : `Model: ${ranModel} (catalog entry ${model.provider}/${model.id})`;
	const lines = [`Provider: ${model.provider}`, modelLine, `Generated ${imagePaths.length} image(s):`];
	for (const [index, imagePath] of imagePaths.entries()) {
		const { size, quality } = result.images[index] ?? {};
		const meta = [size, quality && `quality ${quality}`].filter(Boolean).join(", ");
		lines.push(meta ? `  ${imagePath} (${meta})` : `  ${imagePath}`);
	}
	if (result.text) lines.push("", result.text.trim());
	return {
		content: [{ type: "text", text: lines.join("\n") }],
		details: {
			provider: model.provider,
			model: ranModel,
			imageCount: result.images.length,
			imagePaths,
			images: result.images,
			responseText: result.text,
			usage: result.usage,
		},
	};
}

export const imageGenTool: CustomTool<typeof imageGenSchema, ImageGenToolDetails> = {
	name: "generate_image",
	label: "GenerateImage",
	strict: false,
	approval: "write",
	description: prompt.render(imageGenDescription),
	parameters: imageGenSchema,
	async execute(_toolCallId, params, _onUpdate, ctx, signal) {
		return untilAborted(signal, async () => {
			const sessionId = ctx.sessionManager.getSessionId();
			const requestSignal = ptree.combineSignals(signal, IMAGE_TIMEOUT);
			const effectiveSettings = ctx.settings ?? settings;
			const pool = roleCandidatePool("image", effectiveSettings, ctx.modelRegistry);
			let candidates: Model[];
			if (params.model) {
				const selected = resolveModelRoleValue(params.model, pool, { settings: effectiveSettings }).model;
				if (!selected)
					throw new Error(`Image model selector did not match an available image model: ${params.model}`);
				candidates = [selected];
			} else {
				candidates = defaultImageCandidates(resolveRoleChain("image", effectiveSettings, pool), ctx.model, pool);
			}

			const failures: ProviderHttpError[] = [];
			const skipped: string[] = [];
			let inputImages: ImageGenerationRequest["inputImages"];
			for (const model of candidates) {
				if (!isImageGenerationApi(model.api)) {
					logger.warn("Skipping unsupported image model API", {
						provider: model.provider,
						model: model.id,
						api: model.api,
					});
					skipped.push(`${model.provider}/${model.id} (unsupported ${model.api})`);
					continue;
				}

				const initialKey = await ctx.modelRegistry.getApiKey(model, sessionId, { signal: requestSignal });
				if (!isAuthenticated(initialKey)) {
					skipped.push(`${model.provider}/${model.id} (credentials unavailable)`);
					continue;
				}
				if (model.api === "google-gemini-cli" && !parseAntigravityCredentials(initialKey)) {
					skipped.push(`${model.provider}/${model.id} (invalid credentials)`);
					continue;
				}

				let carrier: Model | undefined;
				let apiKey = ctx.modelRegistry.resolver(model, sessionId);
				if (model.api === "openai-responses" || model.api === "openai-codex-responses") {
					carrier = model.hostedImage ? model : resolveHostedImageCarrier(ctx.modelRegistry, model, ctx.model);
					if (!carrier) {
						skipped.push(`${model.provider}/${model.id} (hosted chat carrier unavailable)`);
						continue;
					}
					const carrierKey = await ctx.modelRegistry.getApiKey(carrier, sessionId, { signal: requestSignal });
					if (!isAuthenticated(carrierKey)) {
						skipped.push(`${model.provider}/${model.id} (carrier credentials unavailable)`);
						continue;
					}
					apiKey = ctx.modelRegistry.resolver(carrier, sessionId);
				}

				if (!inputImages) {
					inputImages = [];
					for (const input of params.input ?? []) {
						inputImages.push(await resolveInputImage(input, ctx.sessionManager.getCwd()));
					}
				}
				const request: ImageGenerationRequest = {
					prompt: assemblePrompt(params),
					inputImages,
					aspectRatio: params.aspect_ratio,
					imageSize: params.image_size,
					count: 1,
				};
				const resolvedModel = {
					...model,
					resolveHeaders: (headerSignal?: AbortSignal) =>
						ctx.modelRegistry.resolveModelHeaders(model, headerSignal),
				};
				const resolvedCarrier = carrier
					? {
							...carrier,
							resolveHeaders: (headerSignal?: AbortSignal) =>
								ctx.modelRegistry.resolveModelHeaders(carrier, headerSignal),
						}
					: undefined;
				try {
					const result = await generateImage(resolvedModel, request, {
						apiKey,
						fetch: ctx.fetch ?? fetch,
						signal: requestSignal,
						carrier: resolvedCarrier,
						sessionId,
					});
					return buildToolResult(model, result);
				} catch (error) {
					if (!(error instanceof ProviderHttpError) || requestSignal?.aborted) throw error;
					failures.push(error);
				}
			}

			const attempted = candidates.map(model => `${model.provider}/${model.id}`).join(", ");
			// 压平空白并截断每条 provider 错误；去掉句末句号，避免与统一补的标点叠成 `..`。
			const reasons =
				failures.length > 0
					? ` Failures: ${failures
							.map(error =>
								truncate(error.message.replace(/\s+/g, " ").trim(), MAX_FAILURE_REASON_CHARS).replace(
									/\.$/,
									"",
								),
							)
							.join("; ")}.`
					: "";
			const suffix = skipped.length > 0 ? ` Skipped: ${skipped.join(", ")}.` : "";
			throw new AggregateError(
				failures,
				`Image generation exhausted the resolved image chain${attempted ? `: ${attempted}` : ""}.${reasons}${suffix}`,
			);
		});
	},
};

export async function getImageGenTools(
	_modelRegistry?: ModelRegistry,
	_activeModel?: Model,
): Promise<Array<CustomTool<typeof imageGenSchema, ImageGenToolDetails>>> {
	return [imageGenTool];
}

export async function getImageGenToolsWithRegistry(
	_modelRegistry: ModelRegistry,
	_activeModel?: Model,
): Promise<Array<CustomTool<typeof imageGenSchema, ImageGenToolDetails>>> {
	return [imageGenTool];
}