import { StringEnum } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ExtensionContext, Theme } from "@earendil-works/pi-coding-agent";
import { Text, truncateToWidth } from "@earendil-works/pi-tui";
import { Type } from "typebox";

type TaskStatus = "pending" | "in_progress" | "completed";

interface Task {
	content: string;
	status: TaskStatus;
}

interface TodoDetails {
	tasks: Task[];
}

const TOOL_NAME = "todo";
const WIDGET_KEY = "native-todo";
const MAX_WIDGET_LINES = 12;
const PROGRESS_SEGMENTS = 10;

const PROMPT_SNIPPET = "Proactively manage the current todo list to track steps and progress for complex, ambiguous, or multi-phase tasks";
const PROMPT_GUIDELINES = [
	"Call `todo` for non-trivial tasks that require multiple actions over a long time horizon, have logical phases or dependencies, or have ambiguity that benefits from outlining goals.",
	"Do NOT use `todo` for padding simple work with filler steps or stating the obvious. Do not use for simple or single-step queries that you can just do or answer immediately.",
	"Break the task into meaningful, logically ordered steps that are easy to verify as you go. Each step should be a short sentence of at most 5-7 words.",
	"High-quality steps are concrete and specific (e.g., 'Add CLI entry with file args', 'Parse Markdown via CommonMark library', 'Apply semantic HTML template'). Low-quality steps are vague (e.g., 'Create CLI tool', 'Convert to HTML').",
	"Mark the current step `in_progress` before starting work, and mark it `completed` immediately after finishing. Do NOT batch-complete multiple items after the fact.",
	"Exactly ONE task `in_progress` at a time. Do not jump from `pending` to `completed`; always set it to `in_progress` first. Complete current steps before starting new ones.",
	"Before moving to the next step, mark the previous step as `completed`.",
	"If understanding changes (split/merge/reorder items), update the plan before continuing. Do not let the plan go stale while coding.",
	"Do not repeat the full contents of the plan after updating it — summarize the change made and highlight any important context or next step.",
	"Finish with all items `completed` or explicitly removed before ending the turn.",
	"Remove tasks that are no longer relevant.",
	"Todo list content must be written in Chinese.",
];

const ParamsSchema = Type.Object({
	tasks: Type.Optional(
		Type.Array(
			Type.Object({
				content: Type.String({ description: "Short actionable task step. At most 5-7 words. Concrete and specific, not vague." }),
				status: StringEnum(["pending", "in_progress", "completed"] as const, {
					description: "Task status: pending, in_progress, or completed",
				}),
			}),
			{
				description:
					"Complete current todo list in execution order. Omit a previously existing task to remove it. Pass an empty array to clear the list.",
			},
		),
	),
});

let cachedTasks: Task[] = [];
let cacheHydrated = false;
let widgetRegistered = false;
let widgetTui: { requestRender: () => void } | undefined;

function cloneTasks(tasks: readonly Task[]): Task[] {
	return tasks.map((task) => ({ ...task }));
}

function isTodoDetails(value: unknown): value is TodoDetails {
	if (!value || typeof value !== "object") return false;
	const record = value as Record<string, unknown>;
	return Array.isArray(record.tasks);
}

function replayFromBranch(ctx: { sessionManager: { getBranch(): Iterable<unknown> } }): Task[] {
	let tasks: Task[] = [];
	for (const entry of ctx.sessionManager.getBranch()) {
		const e = entry as { type?: string; message?: { role?: string; toolName?: string; details?: unknown } };
		if (e.type !== "message") continue;
		const msg = e.message;
		if (!msg || msg.role !== "toolResult" || msg.toolName !== TOOL_NAME) continue;
		if (!isTodoDetails(msg.details)) continue;
		tasks = cloneTasks(
			msg.details.tasks.filter(
				(task): task is Task =>
					!!task &&
					typeof task === "object" &&
					typeof (task as Record<string, unknown>).content === "string" &&
					((task as Record<string, unknown>).status === "pending" ||
						(task as Record<string, unknown>).status === "in_progress" ||
						(task as Record<string, unknown>).status === "completed"),
			),
		);
	}
	return tasks;
}

function syncTasks(tasks: readonly Task[]): void {
	cachedTasks = cloneTasks(tasks);
	cacheHydrated = true;
	widgetTui?.requestRender();
}

function formatList(tasks: readonly Task[]): string {
	if (tasks.length === 0) return "No tasks";
	const done = tasks.filter((task) => task.status === "completed").length;
	const lines = [`${done}/${tasks.length} completed`];
	for (const task of tasks) {
		lines.push(`[${task.status}] ${task.content}`);
	}
	return lines.join("\n");
}

function renderWidget(theme: Theme, width: number): string[] {
	if (cachedTasks.length === 0) return [];
	const safeWidth = Math.max(1, width);
	const lines: string[] = [];
	const push = (line: string) => lines.push(truncateToWidth(line, safeWidth));
	const done = cachedTasks.filter((task) => task.status === "completed").length;
	const total = cachedTasks.length;
	const percentage = Math.round((done / total) * 100);
	const filled = Math.round((done / total) * PROGRESS_SEGMENTS);
	const progressBar =
		theme.fg("success", "▰".repeat(filled)) + theme.fg("dim", "▱".repeat(PROGRESS_SEGMENTS - filled));

	push(`${theme.fg("accent", "● Todos")} ${progressBar} ${theme.fg("muted", `${percentage}% · ${done}/${total}`)}`);

	const hasOverflow = cachedTasks.length > MAX_WIDGET_LINES - 1;
	const taskLineLimit = hasOverflow ? MAX_WIDGET_LINES - 2 : MAX_WIDGET_LINES - 1;
	const display = cachedTasks.slice(0, taskLineLimit);
	for (const task of display) {
		const glyph =
			task.status === "completed"
				? theme.fg("success", "✓")
				: task.status === "in_progress"
					? theme.fg("warning", "◐")
					: theme.fg("dim", "○");
		const text =
			task.status === "completed"
				? theme.fg("dim", theme.strikethrough(task.content))
				: task.status === "in_progress"
					? theme.fg("accent", task.content)
					: theme.fg("text", task.content);
		push(`${glyph} ${text}`);
	}
	if (hasOverflow) push(theme.fg("dim", `+${cachedTasks.length - display.length} more`));
	return lines;
}

function resetWidgetRegistration(): void {
	widgetRegistered = false;
	widgetTui = undefined;
}

function ensureWidget(ctx: ExtensionContext): void {
	if (!ctx.hasUI) return;
	if (cachedTasks.length === 0) {
		ctx.ui.setWidget(WIDGET_KEY, undefined);
		resetWidgetRegistration();
		return;
	}
	if (!widgetRegistered) {
		ctx.ui.setWidget(
			WIDGET_KEY,
			(tui, theme) => {
				widgetTui = tui;
				return {
					render: (width: number) => renderWidget(theme, width),
					invalidate: () => {},
				};
			},
			{ placement: "aboveEditor" },
		);
		widgetRegistered = true;
	} else {
		widgetTui?.requestRender();
	}
}

function validateTasks(raw: unknown): { tasks?: Task[]; error?: string } {
	if (raw === undefined) return { tasks: undefined };
	if (!Array.isArray(raw)) return { error: "tasks must be an array" };
	const tasks: Task[] = [];
	for (const item of raw) {
		if (!item || typeof item !== "object") return { error: "each task must be an object" };
		const record = item as Record<string, unknown>;
		if (typeof record.content !== "string" || record.content.trim() === "") {
			return { error: "each task requires non-empty content" };
		}
		if (record.status !== "pending" && record.status !== "in_progress" && record.status !== "completed") {
			return { error: "each task status must be pending, in_progress, or completed" };
		}
		tasks.push({ content: record.content.trim(), status: record.status });
	}
	const inProgressCount = tasks.filter((task) => task.status === "in_progress").length;
	if (inProgressCount > 1) return { error: "task lists allow at most one in_progress task" };
	return { tasks };
}

function buildResult(tasks: readonly Task[], error?: string) {
	return {
		content: [{ type: "text" as const, text: error ? `Error: ${error}` : formatList(tasks) }],
		details: {
			tasks: cloneTasks(tasks),
			...(error ? { error } : {}),
		},
	};
}

export default function nativeTodoExtension(pi: ExtensionAPI) {
	pi.on("session_start", async (_event, ctx) => {
		resetWidgetRegistration();
		syncTasks(replayFromBranch(ctx));
		ensureWidget(ctx);
	});

	pi.on("session_tree", async (_event, ctx) => {
		syncTasks(replayFromBranch(ctx));
		ensureWidget(ctx);
	});

	pi.on("session_shutdown", async () => {
		resetWidgetRegistration();
	});

	pi.registerTool({
		name: TOOL_NAME,
		label: "Todo",
		description:
			"Read or replace the current todo list for multi-step progress tracking. The list tracks steps and progress. Pass the full tasks array to update the list. Omit a previously existing task to remove it. Pass an empty array to clear the list.",
		promptSnippet: PROMPT_SNIPPET,
		promptGuidelines: PROMPT_GUIDELINES,
		parameters: ParamsSchema,
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const currentTasks = cacheHydrated ? cloneTasks(cachedTasks) : replayFromBranch(ctx);
			syncTasks(currentTasks);
			const validated = validateTasks((params as { tasks?: unknown }).tasks);
			if (validated.error) {
				ensureWidget(ctx);
				return buildResult(currentTasks, validated.error);
			}
			if (validated.tasks === undefined) {
				ensureWidget(ctx);
				return buildResult(currentTasks);
			}
			syncTasks(validated.tasks);
			ensureWidget(ctx);
			return buildResult(validated.tasks);
		},
		renderCall(args, theme) {
			const tasks = (args as { tasks?: Task[] }).tasks;
			const suffix = tasks === undefined ? "read" : `${tasks.length} task${tasks.length === 1 ? "" : "s"}`;
			return new Text(theme.fg("toolTitle", theme.bold("todo ")) + theme.fg("muted", suffix), 0, 0);
		},
		renderResult(result, _opts, theme) {
			const details = result.details as { tasks?: Task[]; error?: string } | undefined;
			if (details?.error) return new Text(theme.fg("error", `✗ ${details.error}`), 0, 0);
			if (details?.tasks && details.tasks.length > 0) {
				const done = details.tasks.filter((task) => task.status === "completed").length;
				return new Text(theme.fg("success", `✓ ${done}/${details.tasks.length}`), 0, 0);
			}
			return new Text(theme.fg("success", "✓"), 0, 0);
		},
	});
}