/**
 * My Rainbow Editor Extension for pi
 *
 * This extension provides a custom TUI editor component for the pi coding agent.
 * It takes the rendered lines of the default editor and applies a continuous 
 * rainbow color pattern to all characters. To prevent the TUI from breaking, 
 * it safely ignores and preserves existing ANSI escape codes (such as cursor 
 * placement and text selection sequences) during the colorization process.
 *
 * Usage: pi --extension ./my-rainbow-editor.ts
 */
import { CustomEditor, type ExtensionAPI } from "@mariozechner/pi-coding-agent";

const COLORS: [number, number, number][] = [
	[233, 137, 115], // coral
	[228, 186, 103], // yellow
	[141, 192, 122], // green
	[102, 194, 179], // teal
	[121, 157, 207], // blue
	[157, 134, 195], // purple
	[206, 130, 172], // pink
];
const RESET = "\x1b[0m";
// Regex to split string on ANSI escape codes, keeping them intact
const ANSI_REGEX = /(\x1b(?:\[[0-9;?]*[a-zA-Z]|(?:\]|_)[^\x07\x1b]*(?:\x07|\x1b\\)))/;

function colorizeSafe(line: string): string {
	const parts = line.split(ANSI_REGEX);
	let out = "";
	let charIndex = 0;
	
	for (const part of parts) {
		if (!part) continue;
		// If this is an ANSI sequence, pass it through without colorizing
		if (part.startsWith("\x1b")) {
			out += part;
		} else {
			// Colorize regular characters individually
			for (const char of part) {
				const [r, g, b] = COLORS[charIndex % COLORS.length]!;
				out += `\x1b[38;2;${r};${g};${b}m${char}`;
				charIndex++;
			}
		}
	}
	// Append a single reset at the end of the line just to be safe
	return out + RESET;
}

class MyRainbowEditor extends CustomEditor {
	render(width: number): string[] {
		// Render the original editor line, but safely add our rainbow colors 
		// on top without breaking the terminal cursor/selection ANSI sequences
		return super.render(width).map(line => colorizeSafe(line));
	}
}

export default function (pi: ExtensionAPI) {
	pi.on("session_start", (_event, ctx) => {
		ctx.ui.setEditorComponent((tui, theme, kb) => new MyRainbowEditor(tui, theme, kb));
	});
}