/**
 * Plan Mode Extension
 *
 * Original source: https://github.com/earendil-works/pi/tree/main/packages/coding-agent/examples/extensions/plan-mode
 * Modified by Andrea Richiardi
 *
 * MIT License
 *
 * Copyright (c) 2025 Mario Zechner
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 * Read-only exploration mode for safe code analysis.
 * When enabled, only read-only tools are available.
 *
 * Features:
 * - /plan command or Ctrl+Alt+P to toggle
 * - Bash restricted to allowlisted read-only commands
 * - Extracts numbered plan steps from "Plan:" sections
 * - [DONE:n] markers to complete steps during execution
 * - Progress tracking widget during execution
 *
 * Configuration:
 * Create ~/.config/pi/agent/ar-llm/plan-mode.json to customize:
 * - Command allowlists (safe/destructive patterns)
 * - Tool restrictions
 * - UI settings
 * - Plan extraction patterns
 * - System prompts
 */

import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { AssistantMessage, TextContent } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Key } from "@earendil-works/pi-tui";
import { isSafeCommand } from "./commands.ts";
import { extractTodoItems, markCompletedSteps, type TodoItem } from "./extraction.ts";
import { planFormatPromptValues, renderPromptTemplate } from "./prompt.ts";
import { resolveConfig } from "./config.ts";
import type { PlanModeConfig, UIConfig } from "./types.ts";

// Type guard for assistant messages
function isAssistantMessage(m: AgentMessage): m is AssistantMessage {
	return m.role === "assistant" && Array.isArray(m.content);
}

// Extract text content from an assistant message
function getTextContent(message: AssistantMessage): string {
	return message.content
		.filter((block): block is TextContent => block.type === "text")
		.map((block) => block.text)
		.join("\n");
}

export default function planModeExtension(pi: ExtensionAPI): void {
	let planModeEnabled = false;
	let executionMode = false;
	let todoItems: TodoItem[] = [];
	let currentConfig: PlanModeConfig | null = null;

	pi.registerFlag("plan", {
		description: "Start in plan mode (read-only exploration)",
		type: "boolean",
		default: false,
	});

	function getConfig(): PlanModeConfig {
		// Resolved once on session_start; fall back to defaults before that
		return currentConfig ?? resolveConfig();
	}

	function getUiConfig(): UIConfig {
		return getConfig().ui;
	}

	function updateStatus(ctx: ExtensionContext): void {
		const uiConfig = getUiConfig();

		// Footer status
		if (uiConfig.showStatusBar) {
			if (executionMode && todoItems.length > 0) {
				const completed = todoItems.filter((t) => t.completed).length;
				const format = uiConfig.statusBarFormat
					.replace("{completed}", String(completed))
					.replace("{total}", String(todoItems.length))
					.replace("{mode}", executionMode ? "exec" : "plan");
				ctx.ui.setStatus("plan-mode", ctx.ui.theme.fg("accent", format));
			} else if (planModeEnabled) {
				ctx.ui.setStatus("plan-mode", ctx.ui.theme.fg("warning", "⏸ plan"));
			} else {
				ctx.ui.setStatus("plan-mode", undefined);
			}
		}

		// Widget showing todo list
		if (uiConfig.showProgressWidget && executionMode && todoItems.length > 0) {
			const lines = todoItems.map((item) => {
				if (item.completed) {
					return (
						ctx.ui.theme.fg("success", "☑ ") + ctx.ui.theme.fg("muted", ctx.ui.theme.strikethrough(item.text))
					);
				}
				return `${ctx.ui.theme.fg("muted", "☐ ")}${item.text}`;
			});
			ctx.ui.setWidget("plan-todos", lines);
		} else {
			ctx.ui.setWidget("plan-todos", undefined);
		}
	}

	function togglePlanMode(ctx: ExtensionContext): void {
		const config = getConfig();
		const uiConfig = getUiConfig();

		planModeEnabled = !planModeEnabled;
		executionMode = false;
		todoItems = [];

		const tools = planModeEnabled ? config.tools.planModeTools : config.tools.normalModeTools;

		if (planModeEnabled) {
			pi.setActiveTools(config.tools.planModeTools);
			const msg = uiConfig.notifications.planModeEnabled.replace("{tools}", tools.join(", "));
			ctx.ui.notify(msg);
		} else {
			pi.setActiveTools(config.tools.normalModeTools);
			ctx.ui.notify(uiConfig.notifications.planModeDisabled);
		}
		updateStatus(ctx);
		persistState();
	}

	function persistState(): void {
		pi.appendEntry("plan-mode", {
			enabled: planModeEnabled,
			todos: todoItems,
			executing: executionMode,
		});
	}

	pi.registerCommand("plan", {
		description: "Toggle plan mode (read-only exploration)",
		handler: async (_args, ctx) => togglePlanMode(ctx),
	});

	pi.registerCommand("todos", {
		description: "Show current plan todo list",
		handler: async (_args, ctx) => {
			if (todoItems.length === 0) {
				ctx.ui.notify(getUiConfig().notifications.noTodos, "info");
				return;
			}
			const list = todoItems.map((item, i) => `${i + 1}. ${item.completed ? "✓" : "○"} ${item.text}`).join("\n");
			ctx.ui.notify(`Plan Progress:\n${list}`, "info");
		},
	});

	pi.registerShortcut(Key.ctrlAlt("p"), {
		description: "Toggle plan mode",
		handler: async (ctx) => togglePlanMode(ctx),
	});

	// Block destructive bash commands in plan mode
	pi.on("tool_call", async (event) => {
		if (!planModeEnabled || event.toolName !== "bash") return;

		const command = event.input.command as string;
		if (!isSafeCommand(command, getConfig().commands)) {
			return {
				block: true,
				reason: `Plan mode: command blocked (not allowlisted). Use /plan to disable plan mode first.\nCommand: ${command}`,
			};
		}
	});

	// Filter out stale plan mode context when not in plan mode
	pi.on("context", async (event) => {
		if (planModeEnabled) return;

		return {
			messages: event.messages.filter((m) => {
				const msg = m as AgentMessage & { customType?: string };
				if (msg.customType === "plan-mode-context") return false;
				if (msg.role !== "user") return true;

				const content = msg.content;
				if (typeof content === "string") {
					return !content.includes("[PLAN MODE ACTIVE]");
				}
				if (Array.isArray(content)) {
					return !content.some(
						(c) => c.type === "text" && (c as TextContent).text?.includes("[PLAN MODE ACTIVE]"),
					);
				}
				return true;
			}),
		};
	});

	// Inject plan/execution context before agent starts
	pi.on("before_agent_start", async () => {
		const config = getConfig();

		if (planModeEnabled) {
			const values = { ...planFormatPromptValues(config.planFormat), tools: config.tools.planModeTools.join(", ") };
			const prompt = renderPromptTemplate(config.prompts.planModeContext, values);
			return {
				message: {
					customType: "plan-mode-context",
					content: prompt,
					display: false,
				},
			};
		}

		if (executionMode && todoItems.length > 0) {
			const remaining = todoItems.filter((t) => !t.completed);
			const todoList = remaining.map((t) => `${t.step}. ${t.text}`).join("\n");
			const values = { ...planFormatPromptValues(config.planFormat), todoList };
			const prompt = renderPromptTemplate(config.prompts.executionContext, values);
			return {
				message: {
					customType: "plan-execution-context",
					content: prompt,
					display: false,
				},
			};
		}
	});

	// Track progress after each turn
	pi.on("turn_end", async (event, ctx) => {
		if (!executionMode || todoItems.length === 0) return;
		if (!isAssistantMessage(event.message)) return;

		const text = getTextContent(event.message);
		if (markCompletedSteps(text, todoItems, getConfig().planFormat) > 0) {
			updateStatus(ctx);
		}
		persistState();
	});

	// Handle plan completion and plan mode UI
	pi.on("agent_end", async (event, ctx) => {
		const config = getConfig();

		// Check if execution is complete
		if (executionMode && todoItems.length > 0) {
			if (todoItems.every((t) => t.completed)) {
				const completedList = todoItems.map((t) => `~~${t.text}~~`).join("\n");
				pi.sendMessage(
					{ customType: "plan-complete", content: `**Plan Complete!** ✓\n\n${completedList}`, display: true },
					{ triggerTurn: false },
				);
				executionMode = false;
				todoItems = [];
				pi.setActiveTools(config.tools.normalModeTools);
				updateStatus(ctx);
				persistState(); // Save cleared state so resume doesn't restore old execution mode
			}
			return;
		}

		if (!planModeEnabled || !ctx.hasUI) return;

		const uiConfig = getUiConfig();

		// Extract todos from last assistant message
		const lastAssistant = [...event.messages].reverse().find(isAssistantMessage);
		let extracted = false;
		if (lastAssistant) {
			const items = extractTodoItems(getTextContent(lastAssistant), config.planFormat);
			if (items.length > 0) {
				todoItems = items;
				extracted = true;
			}
		}

		persistState();

		// Show plan steps when the last message contained an extractable plan
		const hasTodos = extracted && todoItems.length > 0;
		let planTodoListMessage: { customType: string; content: string; display: boolean } | null = null;
		if (hasTodos) {
			const todoListText = todoItems.map((t, i) => `${i + 1}. ☐ ${t.text}`).join("\n");
			planTodoListMessage = {
				customType: "plan-todo-list",
				content: `**Plan Steps (${todoItems.length}):**\n\n${todoListText}`,
				display: true,
			};
		} else {
			ctx.ui.notify(uiConfig.notifications.planNotDetected, "info");
		}

		const choice = await ctx.ui.select("Plan mode - what next?", [
			hasTodos ? uiConfig.choices.executeWithTodos : uiConfig.choices.createPlan,
			uiConfig.choices.stayInPlanMode,
			uiConfig.choices.refinePlan,
		]);

		if (hasTodos && choice === uiConfig.choices.executeWithTodos) {
			const firstTodoItem = todoItems[0];
			if (!firstTodoItem) return;

			planModeEnabled = false;
			executionMode = true;
			pi.setActiveTools(config.tools.normalModeTools);
			updateStatus(ctx);
			persistState();

			const remainingList = todoItems.map((t) => `${t.step}. ${t.text}`).join("\n");
			const execMessage = `Execute the plan.\n\nRemaining steps:\n${remainingList}\n\nStart with: ${firstTodoItem.text}\nAfter completing a step, include a [DONE:n] tag in your response.`;
			if (planTodoListMessage) pi.sendMessage(planTodoListMessage, { deliverAs: "followUp" });
			pi.sendMessage(
				{ customType: "plan-mode-execute", content: execMessage, display: true },
				{ triggerTurn: true, deliverAs: "followUp" },
			);
		} else if (choice === uiConfig.choices.refinePlan) {
			const refinement = await ctx.ui.editor("Refine the plan:", "");
			if (refinement?.trim()) {
				if (planTodoListMessage) pi.sendMessage(planTodoListMessage, { deliverAs: "followUp" });
				pi.sendUserMessage(refinement.trim(), { deliverAs: "followUp" });
			}
		} else if (choice === uiConfig.choices.createPlan) {
			const values = planFormatPromptValues(config.planFormat);
			const prompt = renderPromptTemplate(config.prompts.planCreationPrompt, values);
			pi.sendUserMessage(prompt, { deliverAs: "followUp" });
		}
	});

	// Restore state on session start/resume
	pi.on("session_start", async (_event, ctx) => {
		// Load configuration
		currentConfig = resolveConfig();
		const config = currentConfig;

		if (pi.getFlag("plan") === true) {
			planModeEnabled = true;
		}

		const entries = ctx.sessionManager.getEntries();

		// Restore persisted state
		const planModeEntry = entries
			.filter((e: { type: string; customType?: string }) => e.type === "custom" && e.customType === "plan-mode")
			.pop() as { data?: { enabled: boolean; todos?: TodoItem[]; executing?: boolean } } | undefined;

		if (planModeEntry?.data) {
			planModeEnabled = planModeEntry.data.enabled ?? planModeEnabled;
			todoItems = planModeEntry.data.todos ?? todoItems;
			executionMode = planModeEntry.data.executing ?? executionMode;
		}

		// On resume: re-scan messages to rebuild completion state
		// Only scan messages AFTER the last "plan-mode-execute" to avoid picking up [DONE:n] from previous plans
		const isResume = planModeEntry !== undefined;
		if (isResume && executionMode && todoItems.length > 0) {
			// Find the index of the last plan-mode-execute entry (marks when current execution started)
			let executeIndex = -1;
			for (let i = entries.length - 1; i >= 0; i--) {
				const entry = entries[i] as { type: string; customType?: string };
				if (entry.customType === "plan-mode-execute") {
					executeIndex = i;
					break;
				}
			}

			// Only scan messages after the execute marker
			const messages: AssistantMessage[] = [];
			for (let i = executeIndex + 1; i < entries.length; i++) {
				const entry = entries[i];
				if (entry.type === "message" && "message" in entry && isAssistantMessage(entry.message as AgentMessage)) {
					messages.push(entry.message as AssistantMessage);
				}
			}
			const allText = messages.map(getTextContent).join("\n");
			markCompletedSteps(allText, todoItems, config.planFormat);
		}

		if (planModeEnabled) {
			pi.setActiveTools(config.tools.planModeTools);
		}
		updateStatus(ctx);
	});
}