import type {
	AgentToolResult,
	ExtensionAPI,
	ExtensionCommandContext,
} from "@earendil-works/pi-coding-agent";
import {
	type Component,
	Key,
	matchesKey,
	type TUI,
	truncateToWidth,
	visibleWidth,
} from "@earendil-works/pi-tui";
import { Type } from "typebox";
import { boardPrompt, type GameState, TicTacToeGame } from "./game.ts";

const SAVE_TYPE = "tic-tac-toe-save";
const moveSchema = Type.Object({
	square: Type.Integer({
		minimum: 1,
		maximum: 9,
		description: "Board square from 1 to 9.",
	}),
});

type MoveDetails = { square: number; state: GameState; status: string };

class TicTacToeComponent implements Component {
	cursor = 4;
	error = "";
	private tui: TUI;
	private game: TicTacToeGame;
	private onClose: () => void;
	private onMove: (cell: number) => void;
	private onNew: () => void;

	constructor(
		tui: TUI,
		game: TicTacToeGame,
		onClose: () => void,
		onMove: (cell: number) => void,
		onNew: () => void,
	) {
		this.tui = tui;
		this.game = game;
		this.onClose = onClose;
		this.onMove = onMove;
		this.onNew = onNew;
	}

	handleInput(data: string): void {
		if (matchesKey(data, "q")) return this.onClose();
		if (matchesKey(data, "n")) return this.onNew();
		if (this.game.winner() || this.game.turn !== "X") return;

		const digit = Number(data);
		if (Number.isInteger(digit) && digit >= 1 && digit <= 9)
			return this.pick(digit - 1);
		if (matchesKey(data, Key.enter) || matchesKey(data, Key.return))
			return this.pick(this.cursor);
		if (matchesKey(data, Key.up)) this.cursor = Math.max(0, this.cursor - 3);
		if (matchesKey(data, Key.down)) this.cursor = Math.min(8, this.cursor + 3);
		if (matchesKey(data, Key.left) && this.cursor % 3) this.cursor--;
		if (matchesKey(data, Key.right) && this.cursor % 3 < 2) this.cursor++;
		this.tui.requestRender();
	}

	pick(cell: number): void {
		this.cursor = cell;
		this.onMove(cell);
	}

	render(width: number): string[] {
		return [
			this.center("TIC TAC TOE - You(X) vs LLM(O)", width),
			"",
			...this.boardLines(width),
			"",
			this.center(`Status: ${this.game.status()}`, width),
			this.error ? this.center(`Error: ${this.error}`, width) : "",
			this.center("Arrows/1-9=move  Enter=place  n=new  q=quit", width),
		].map((line) => truncateToWidth(line, width));
	}

	boardLines(width: number): string[] {
		const rows = [0, 3, 6].map((start) =>
			[0, 1, 2].map((i) => this.cell(start + i)).join("|"),
		);
		return rows.flatMap((row, i) =>
			i < 2
				? [this.center(row, width), this.center("---+---+---", width)]
				: [this.center(row, width)],
		);
	}

	cell(i: number): string {
		const mark = this.game.board[i] ?? String(i + 1);
		return i === this.cursor && this.game.turn === "X" && !this.game.winner()
			? `[${mark}]`
			: ` ${mark} `;
	}

	center(text: string, width: number): string {
		return (
			" ".repeat(Math.max(0, Math.floor((width - visibleWidth(text)) / 2))) +
			text
		);
	}

	setGame(game: TicTacToeGame): void {
		this.game = game;
		this.cursor = game.legalMoves()[0] ?? 4;
		this.error = "";
		this.tui.requestRender();
	}

	setError(error: string): void {
		this.error = error;
		this.tui.requestRender();
	}

	invalidate(): void {}
}

export default function ticTacToe(pi: ExtensionAPI): void {
	let game: TicTacToeGame | null = null;
	let active: TicTacToeComponent | null = null;

	pi.on("session_start", (_event, ctx) => {
		game = loadGame(ctx.sessionManager.getBranch());
	});

	pi.registerCommand("tic-tac-toe", {
		description: "Play tic tac toe against the LLM",
		handler: async (args, ctx) => {
			if (ctx.mode !== "tui") {
				ctx.ui.notify("Tic tac toe overlay only works in TUI mode", "error");
				return;
			}
			if (!game || game.winner() || args.trim() === "new") {
				game = new TicTacToeGame();
				save();
			}
			await show(ctx);
		},
	});

	pi.registerTool<typeof moveSchema, MoveDetails>({
		name: "make_tic_tac_toe_move",
		label: "Tic Tac Toe Move",
		description: "Place O on the tic tac toe board for the LLM.",
		parameters: moveSchema,
		async execute(_id, params) {
			if (!game || game.turn !== "O") throw new Error("Not the LLM turn.");
			const cell = params.square - 1;
			if (!game.play(cell, "O"))
				throw new Error(`Invalid square. Legal: ${legal()}`);
			save();
			active?.setGame(game);
			return result(
				`Moved to ${params.square}. ${game.winner() ? game.status() : "X to move."}`,
				params.square,
			);
		},
	});

	function handleHumanMove(cell: number): void {
		if (!game?.play(cell, "X"))
			return active?.setError(`Invalid square: ${cell + 1}`);
		save();
		active?.setGame(game);
		if (game.winner()) return;
		pi.sendUserMessage(
			`You are O in tic tac toe. Pick the best legal square and call make_tic_tac_toe_move.\n\nBoard:\n${boardPrompt(game)}\n\nLegal squares: ${legal()}`,
			{ deliverAs: "steer" },
		);
	}

	async function show(ctx: ExtensionCommandContext): Promise<void> {
		await ctx.ui.custom(
			(tui, _theme, _keybindings, done) => {
				active = new TicTacToeComponent(
					tui,
					game!,
					() => {
						active = null;
						done(undefined);
					},
					handleHumanMove,
					() => {
						game = new TicTacToeGame();
						save();
						active?.setGame(game);
					},
				);
				return active;
			},
			{
				overlay: true,
				overlayOptions: {
					width: "44%",
					maxHeight: "70%",
					anchor: "center",
					margin: { top: 1 },
				},
			},
		);
	}

	function save(): void {
		if (game) pi.appendEntry(SAVE_TYPE, game.toState());
	}

	function legal(): string {
		return (
			game
				?.legalMoves()
				.map((cell) => cell + 1)
				.join(", ") ?? ""
		);
	}

	function result(text: string, square: number): AgentToolResult<MoveDetails> {
		return {
			content: [{ type: "text", text }],
			details: { square, state: game!.toState(), status: game!.status() },
		};
	}
}

function loadGame(
	entries: Array<{ type: string; customType?: string; data?: unknown }>,
): TicTacToeGame | null {
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i];
		if (entry.type === "custom" && entry.customType === SAVE_TYPE) {
			try {
				return TicTacToeGame.fromState(entry.data);
			} catch {
				return null;
			}
		}
	}
	return null;
}