import type { AgentTool } from "@jaeswift/jae-agent-core";
import type { ToolResultMessage } from "@jaeswift/jae-ai";
import { Type } from "@sinclair/typebox";
import { html } from "lit";
import { GitBranch } from "lucide";
import { registerToolRenderer, renderHeader } from "./renderer-registry.js";
import type { ToolRenderer, ToolRenderResult } from "./types.js";

export interface MermaidDetails {
	diagram: string;
	rendered: boolean;
	error?: string;
}

interface MermaidParams {
	diagram: string;
	title?: string;
}

const mermaidSchema = Type.Object({
	diagram: Type.String({ description: "Mermaid diagram source code" }),
	title: Type.Optional(Type.String({ description: "Optional title for the diagram" })),
});

export const mermaidTool: AgentTool<typeof mermaidSchema, MermaidDetails> = {
	name: "render_diagram",
	label: "Render Diagram",
	description: "Render a Mermaid diagram (flowchart, sequence, gantt, class diagram, etc.)",
	parameters: mermaidSchema,
	async execute(toolCallId, params, signal) {
		return {
			content: [{ type: "text", text: `Diagram rendered: ${params.title || "Untitled"}` }],
			details: { diagram: params.diagram, rendered: true },
		};
	},
};

let mermaidLoaded = false;
async function loadMermaid(): Promise<any> {
	if ((window as any).mermaid) return (window as any).mermaid;
	return new Promise((resolve, reject) => {
		const script = document.createElement("script");
		script.src = "https://cdn.jsdelivr.net/npm/mermaid@11/dist/mermaid.min.js";
		script.onload = () => {
			const m = (window as any).mermaid;
			m.initialize({
				startOnLoad: false,
				theme: document.documentElement.classList.contains("dark") ? "dark" : "default",
			});
			mermaidLoaded = true;
			resolve(m);
		};
		script.onerror = reject;
		document.head.appendChild(script);
	});
}

class MermaidRenderer implements ToolRenderer<MermaidParams, MermaidDetails> {
	render(params: MermaidParams | undefined, result: ToolResultMessage<MermaidDetails> | undefined): ToolRenderResult {
		const state = result ? (result.isError ? "error" : "complete") : "inprogress";
		const diagram = result?.details?.diagram || params?.diagram || "";
		const title = params?.title || "Diagram";

		if (!diagram) {
			return { content: renderHeader(state, GitBranch, "Rendering diagram..."), isCustom: false };
		}

		const containerId = `mermaid-${Math.random().toString(36).slice(2)}`;

		const renderDiagram = async (container: HTMLElement) => {
			try {
				const mermaid = await loadMermaid();
				const { svg } = await mermaid.render(containerId + "-svg", diagram);
				container.innerHTML = svg;
				container.style.maxWidth = "100%";
				const svgEl = container.querySelector("svg");
				if (svgEl) {
					svgEl.style.maxWidth = "100%";
					svgEl.style.height = "auto";
				}
			} catch (err: any) {
				container.innerHTML = `<div class="text-destructive text-sm p-2">Diagram error: ${err.message}</div>`;
			}
		};

		return {
			content: html`
        <div class="flex flex-col gap-3">
          ${renderHeader(state, GitBranch, `Diagram: ${title}`)}
          <div
            class="rounded border border-border bg-background p-4 overflow-x-auto"
            ${(el: HTMLElement | undefined) => el && renderDiagram(el)}
          ></div>
        </div>
      `,
			isCustom: false,
		};
	}
}

registerToolRenderer("render_diagram", new MermaidRenderer());

export function createMermaidTool(): AgentTool<typeof mermaidSchema, MermaidDetails> {
	return mermaidTool;
}
