import "@mariozechner/mini-lit/dist/ThemeToggle.js";
import { Agent, type AgentMessage } from "@jaeswift/jae-agent-core";
import { getModel } from "@jaeswift/jae-ai";
import {
	type AgentState,
	ApiKeyPromptDialog,
	AppStorage,
	ChatPanel,
	CustomProvidersStore,
	createJavaScriptReplTool,
	IndexedDBStorageBackend,
	ProviderKeysStore,
	ProvidersModelsTab,
	ProxyTab,
	SessionListDialog,
	SessionsStore,
	SettingsDialog,
	SettingsStore,
	setAppStorage,
} from "@jaeswift/jae-web-ui";
import { html, render } from "lit";
import { Brain, Download, History, Keyboard, Plus, Settings } from "lucide";
import "./app.css";
import { createImageGenTool, createTTSTool, createWebSearchTool } from "@jaeswift/jae-web-ui";
import { icon } from "@mariozechner/mini-lit";
import { Button } from "@mariozechner/mini-lit/dist/Button.js";
import { Input } from "@mariozechner/mini-lit/dist/Input.js";
import type { CommandPalette } from "./components/command-palette.js";
import type { CostTracker } from "./components/cost-tracker.js";
import type { KeyboardShortcuts } from "./components/keyboard-shortcuts.js";
import type { MemoryManager } from "./components/memory-manager.js";
import { exportSessionAsJson, exportSessionAsMarkdown } from "./components/session-export.js";
import { customConvertToLlm, registerCustomMessageRenderers } from "./custom-messages.js";
import "./components/command-palette.js";
import "./components/keyboard-shortcuts.js";
import "./components/memory-manager.js";
import "./components/cost-tracker.js";
import { JaeEmptyState } from "./components/empty-state.js";
import type { JaeUtilityToggle, UtilityVisibility } from "./components/utility-toggle.js";
import "./components/empty-state.js";
import "./components/utility-toggle.js";
import type { JaeSessionSidebar } from "./components/session-sidebar.js";
import "./components/session-sidebar.js";

import type { JaeBrowserPanel } from "./components/browser-panel.js";
import type { JaeTerminalPanel } from "./components/terminal-panel.js";
import "./components/terminal-panel.js";
import "./components/browser-panel.js";

registerCustomMessageRenderers();

const settings = new SettingsStore();
const providerKeys = new ProviderKeysStore();
const sessions = new SessionsStore();
const customProviders = new CustomProvidersStore();

const configs = [
	settings.getConfig(),
	SessionsStore.getMetadataConfig(),
	providerKeys.getConfig(),
	customProviders.getConfig(),
	sessions.getConfig(),
];

const backend = new IndexedDBStorageBackend({
	dbName: "jae-web-ui-example",
	version: 2,
	stores: configs,
});

settings.setBackend(backend);
providerKeys.setBackend(backend);
customProviders.setBackend(backend);
sessions.setBackend(backend);

const storage = new AppStorage(settings, providerKeys, sessions, customProviders, backend);
setAppStorage(storage);

let currentSessionId: string | undefined;
let currentTitle = "";
let isEditingTitle = false;
let agent: Agent;
let rightPanel: "none" | "terminal" | "browser" = "none";
let sidebarWidth = 220;
let rightPanelWidth = 480;
let hasStarted = false;
let terminalPanel: JaeTerminalPanel | null = null;
const browserPanel: JaeBrowserPanel | null = null;
let chatPanel: ChatPanel;
let agentUnsubscribe: (() => void) | undefined;

const commandPalette = document.createElement("command-palette") as CommandPalette;
const keyboardShortcuts = document.createElement("keyboard-shortcuts") as KeyboardShortcuts;
const memoryManager = document.createElement("memory-manager") as MemoryManager;
const costTracker = document.createElement("cost-tracker") as CostTracker;
const sidebar = document.createElement("jae-session-sidebar") as JaeSessionSidebar;

sidebar.onLoadSession = async (id: string) => {
	await loadSession(id);
};
sidebar.onNewSession = () => newSession();
sidebar.addEventListener("delete-session", async (e: Event) => {
	const id = (e as CustomEvent<string>).detail;
	if (storage.sessions) {
		await storage.sessions.delete(id);
		if (id === currentSessionId) newSession();
		await refreshSidebar();
	}
});

const utilityToggle = document.createElement("jae-utility-toggle") as JaeUtilityToggle;
utilityToggle.addEventListener("visibility-change", (e: Event) => {
	const vis = (e as CustomEvent<UtilityVisibility>).detail;
	const chatEl = document.getElementById("chat-wrapper");
	if (chatEl) {
		chatEl.classList.toggle("hide-tool-calls", !vis.showToolCalls);
		chatEl.classList.toggle("hide-thinking", !vis.showThinking);
		chatEl.classList.toggle("hide-system-msgs", !vis.showSystemMessages);
		chatEl.classList.toggle("hide-timestamps", !vis.showTimestamps);
	}
});

document.body.appendChild(commandPalette);
document.body.appendChild(keyboardShortcuts);
document.body.appendChild(memoryManager);

const refreshSidebar = async () => {
	if (storage.sessions) {
		const all = await storage.sessions.getAllMetadata();
		sidebar.setSessions(all);
		sidebar.currentSessionId = currentSessionId;
	}
};

window.addEventListener("keydown", (e: KeyboardEvent) => {
	const meta = e.metaKey || e.ctrlKey;
	if (meta && e.key === "k") {
		e.preventDefault();
		commandPalette.show();
		return;
	}
	if (meta && e.key === "e") {
		e.preventDefault();
		handleExport();
		return;
	}
	if (meta && e.key === "n") {
		e.preventDefault();
		newSession();
		return;
	}
	if (e.key === "?" && !(e.target instanceof HTMLInputElement) && !(e.target instanceof HTMLTextAreaElement)) {
		keyboardShortcuts.toggle();
	}
});

function setupCommands() {
	commandPalette.setCommands([
		{
			id: "new-session",
			label: "New Session",
			description: "Start a fresh conversation",
			shortcut: "Ctrl+N",
			keywords: ["new", "fresh", "start"],
			action: newSession,
		},
		{
			id: "sessions",
			label: "Session History",
			description: "Browse and load past sessions",
			shortcut: "Ctrl+H",
			keywords: ["history", "sessions", "past"],
			action: () =>
				SessionListDialog.open(
					async (id) => await loadSession(id),
					(id) => {
						if (id === currentSessionId) newSession();
					},
				),
		},
		{
			id: "export-md",
			label: "Export as Markdown",
			description: "Download current session as .md",
			shortcut: "Ctrl+E",
			keywords: ["export", "download", "markdown"],
			action: () => handleExport("markdown"),
		},
		{
			id: "export-json",
			label: "Export as JSON",
			description: "Download current session as .json",
			keywords: ["export", "download", "json"],
			action: () => handleExport("json"),
		},
		{
			id: "memory",
			label: "Memory Manager",
			description: "Browse and manage stored memories",
			keywords: ["memory", "remember", "recall"],
			action: () => memoryManager.show(),
		},
		{
			id: "settings",
			label: "Settings",
			description: "Configure providers and models",
			keywords: ["settings", "config", "provider", "api", "model"],
			action: () => SettingsDialog.open([new ProvidersModelsTab(), new ProxyTab()]),
		},
		{
			id: "shortcuts",
			label: "Keyboard Shortcuts",
			description: "View all keyboard shortcuts",
			shortcut: "?",
			keywords: ["keyboard", "shortcuts", "help"],
			action: () => keyboardShortcuts.show(),
		},
		{
			id: "cost",
			label: "Token Usage & Cost",
			description: "View API usage stats for this session",
			keywords: ["tokens", "cost", "usage"],
			action: () => costTracker.dispatchEvent(new MouseEvent("click")),
		},
	]);
}

function handleExport(format: "markdown" | "json" = "markdown") {
	if (!agent) return;
	const messages = agent.state.messages;
	const title = currentTitle || "JAE Session";
	if (format === "markdown") exportSessionAsMarkdown(messages, title);
	else exportSessionAsJson(messages, title);
}

const generateTitle = (messages: AgentMessage[]): string => {
	const firstUserMsg = messages.find((m) => m.role === "user" || m.role === "user-with-attachments");
	if (!firstUserMsg || (firstUserMsg.role !== "user" && firstUserMsg.role !== "user-with-attachments")) return "";
	let text = "";
	const content = firstUserMsg.content;
	if (typeof content === "string") {
		text = content;
	} else {
		const textBlocks = content.filter((c: any) => c.type === "text");
		text = textBlocks.map((c: any) => c.text || "").join(" ");
	}
	text = text.trim();
	if (!text) return "";
	const sentenceEnd = text.search(/[.!?]/);
	if (sentenceEnd > 0 && sentenceEnd <= 50) return text.substring(0, sentenceEnd + 1);
	return text.length <= 50 ? text : text.substring(0, 47) + "...";
};

const shouldSaveSession = (messages: AgentMessage[]): boolean => {
	const hasUserMsg = messages.some((m: any) => m.role === "user" || m.role === "user-with-attachments");
	const hasAssistantMsg = messages.some((m: any) => m.role === "assistant");
	return hasUserMsg && hasAssistantMsg;
};

const saveSession = async () => {
	if (!storage.sessions || !currentSessionId || !agent || !currentTitle) return;
	const state = agent.state;
	if (!shouldSaveSession(state.messages)) return;
	try {
		const sessionData = {
			id: currentSessionId,
			title: currentTitle,
			model: state.model!,
			thinkingLevel: state.thinkingLevel,
			messages: state.messages,
			createdAt: new Date().toISOString(),
			lastModified: new Date().toISOString(),
		};
		const metadata = {
			id: currentSessionId,
			title: currentTitle,
			createdAt: sessionData.createdAt,
			lastModified: sessionData.lastModified,
			messageCount: state.messages.length,
			usage: {
				input: 0,
				output: 0,
				cacheRead: 0,
				cacheWrite: 0,
				totalTokens: 0,
				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
			},
			modelId: state.model?.id || null,
			thinkingLevel: state.thinkingLevel,
			preview: generateTitle(state.messages),
		};
		await storage.sessions.save(sessionData, metadata);
		await refreshSidebar();
	} catch (err) {
		console.error("Failed to save session:", err);
	}
};

const updateUrl = (sessionId: string) => {
	const url = new URL(window.location.href);
	url.searchParams.set("session", sessionId);
	window.history.replaceState({}, "", url);
};

const createAgent = async (initialState?: Partial<AgentState>) => {
	if (agentUnsubscribe) agentUnsubscribe();
	agent = new Agent({
		initialState: initialState || {
			systemPrompt:
				"You are JAE, a helpful AI assistant and coding agent with access to tools including web search, image generation, JavaScript REPL, text-to-speech, and artifact creation. Use these tools whenever helpful.",
			model: getModel("venice", "llama-3.3-70b"),
			thinkingLevel: "off",
			messages: [],
			tools: [],
		},
		convertToLlm: customConvertToLlm,
		onApiKeyRequired: async (provider: string) => {
			const key = await ApiKeyPromptDialog.prompt(provider);
			if (key) await providerKeys.set(provider, key);
			return key;
		},
		getProviderApiKey: async (provider: string) => providerKeys.get(provider),
		onStateChange: async (state: AgentState, prevState: AgentState | undefined) => {
			if (state.messages.length > 0) hasStarted = true;
			if (prevState?.messages.length !== state.messages.length) {
				if (!currentTitle) {
					const generated = generateTitle(state.messages);
					if (generated) {
						currentTitle = generated;
						if (!currentSessionId) currentSessionId = crypto.randomUUID();
						updateUrl(currentSessionId);
					}
				}
				await saveSession();
			}
			renderApp();
		},
		createTools: async (runtimeProvidersFactory: any) => {
			const replTool = createJavaScriptReplTool();
			replTool.runtimeProvidersFactory = runtimeProvidersFactory;
			return [replTool, createWebSearchTool(), createImageGenTool(), createTTSTool()];
		},
	});
	costTracker.bindAgent(agent);
	chatPanel?.setAgent(agent);
	if (!currentSessionId) currentSessionId = crypto.randomUUID();
};

const loadSession = async (sessionId: string): Promise<boolean> => {
	if (!storage.sessions) return false;
	const sessionData = await storage.sessions.get(sessionId);
	if (!sessionData) return false;
	currentSessionId = sessionId;
	hasStarted = sessionData.messages.length > 0;
	const metadata = await storage.sessions.getMetadata(sessionId);
	currentTitle = metadata?.title || "";
	await createAgent({
		model: sessionData.model,
		thinkingLevel: sessionData.thinkingLevel,
		messages: sessionData.messages,
		tools: [],
	});
	sidebar.currentSessionId = currentSessionId;
	updateUrl(sessionId);
	renderApp();
	return true;
};

const newSession = () => {
	currentSessionId = undefined;
	currentTitle = "";
	isEditingTitle = false;
	hasStarted = false;
	createAgent().then(() => renderApp());
};

const handleSuggestion = (e: Event) => {
	const text = (e as CustomEvent<string>).detail;
	if (!text) return;
	// Try ChatPanel.agentInterface.setInput first
	if (chatPanel?.agentInterface) {
		chatPanel.agentInterface.setInput(text);
		// Focus the textarea after injection
		requestAnimationFrame(() => {
			const ta =
				(document.querySelector("message-editor textarea") as HTMLTextAreaElement) ||
				(document.querySelector("textarea") as HTMLTextAreaElement);
			if (ta) ta.focus();
		});
	} else {
		const ta =
			(document.querySelector("message-editor textarea") as HTMLTextAreaElement) ||
			(document.querySelector("textarea") as HTMLTextAreaElement);
		if (ta) {
			ta.value = text;
			ta.dispatchEvent(new Event("input", { bubbles: true }));
			ta.focus();
		}
	}
};

const getModelLabel = (): string | null => {
	if (!agent?.state?.model) return null;
	const m = agent.state.model as any;
	return m.name || m.id || null;
};

const renderApp = () => {
	const app = document.getElementById("app");
	if (!app) return;
	const hasMessages = hasStarted || !!agent?.state?.messages?.length;
	render(
		html`
<div class="w-full h-screen flex flex-col bg-background text-foreground overflow-hidden">
  <div class="flex items-center justify-between border-b border-border shrink-0" style="height:44px">
    <div class="flex items-center gap-1 px-2">
      ${Button({
			variant: "ghost",
			size: "sm",
			children: icon(History, "sm"),
			onClick: () =>
				SessionListDialog.open(
					async (id) => {
						await loadSession(id);
					},
					(id) => {
						if (id === currentSessionId) newSession();
					},
				),
			title: "Sessions",
		})}
      ${Button({ variant: "ghost", size: "sm", children: icon(Plus, "sm"), onClick: newSession, title: "New Session (Ctrl+N)" })}
        <div class="flex items-center gap-2">
          <img src="/mascot/jae-default.png" alt="JAE" class="w-7 h-auto header-logo cursor-pointer" />
          <span class="text-base font-semibold text-foreground">JAE</span>
          ${getModelLabel() ? html`<span class="ml-1 text-[11px] text-muted-foreground bg-muted/80 px-1.5 py-0.5 rounded font-mono truncate max-w-[180px] cursor-pointer hover:bg-muted" title="${getModelLabel()}">${getModelLabel()}</span>` : html``}
        </div>
        ${
				currentTitle
					? isEditingTitle
						? html`<div class="flex items-center gap-2">${Input({
								type: "text",
								value: currentTitle,
								className: "text-sm w-64",
								onChange: async (e: Event) => {
									const v = (e.target as HTMLInputElement).value.trim();
									if (v && v !== currentTitle && storage.sessions && currentSessionId) {
										await storage.sessions.updateTitle(currentSessionId, v);
										currentTitle = v;
										await refreshSidebar();
									}
									isEditingTitle = false;
									renderApp();
								},
								onKeyDown: async (e: KeyboardEvent) => {
									if (e.key === "Enter") {
										const v = (e.target as HTMLInputElement).value.trim();
										if (v && v !== currentTitle && storage.sessions && currentSessionId) {
											await storage.sessions.updateTitle(currentSessionId, v);
											currentTitle = v;
											await refreshSidebar();
										}
										isEditingTitle = false;
										renderApp();
									} else if (e.key === "Escape") {
										isEditingTitle = false;
										renderApp();
									}
								},
							})}</div>`
						: html`<button class="px-2 py-1 text-sm text-foreground hover:bg-secondary rounded transition-colors max-w-xs truncate" @click=${() => {
								isEditingTitle = true;
								renderApp();
								requestAnimationFrame(() => {
									const inp = app.querySelector('input[type="text"]') as HTMLInputElement;
									if (inp) {
										inp.focus();
										inp.select();
									}
								});
							}} title="Click to edit">${currentTitle}</button>`
					: html``
			}
    </div>
    <div class="flex items-center gap-1 px-2">
      ${costTracker}
      ${Button({ variant: "ghost", size: "sm", children: icon(Brain, "sm"), onClick: () => memoryManager.show(), title: "Memory Manager" })}
      ${Button({ variant: "ghost", size: "sm", children: icon(Download, "sm"), onClick: () => handleExport(), title: "Export Session (Ctrl+E)" })}
      ${Button({ variant: "ghost", size: "sm", children: icon(Keyboard, "sm"), onClick: () => keyboardShortcuts.show(), title: "Keyboard Shortcuts (?)" })}
      ${Button({ variant: "ghost", size: "sm", children: html`<span class="text-xs font-mono px-1">&#x2318;K</span>`, onClick: () => commandPalette.show(), title: "Command Palette (Ctrl+K)" })}
      ${utilityToggle}
      <theme-toggle></theme-toggle>
      ${Button({
			variant: "ghost",
			size: "sm",
			children: html`<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><polyline points="4 17 10 11 4 5"/><line x1="12" y1="19" x2="20" y2="19"/></svg>`,
			onClick: () => {
				rightPanel = rightPanel === "terminal" ? "none" : "terminal";
				renderApp();
				if (rightPanel === "terminal")
					requestAnimationFrame(() => {
						terminalPanel = document.querySelector("jae-terminal-panel") as JaeTerminalPanel;
						terminalPanel?.connect();
					});
			},
			title: "Toggle Terminal",
		})}
${Button({
	variant: "ghost",
	size: "sm",
	children: html`<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><circle cx="12" cy="12" r="10"/><line x1="2" y1="12" x2="22" y2="12"/><path d="M12 2a15.3 15.3 0 0 1 4 10 15.3 15.3 0 0 1-4 10 15.3 15.3 0 0 1-4-10 15.3 15.3 0 0 1 4-10z"/></svg>`,
	onClick: () => {
		rightPanel = rightPanel === "browser" ? "none" : "browser";
		renderApp();
	},
	title: "Toggle Browser",
})}
${Button({ variant: "ghost", size: "sm", children: icon(Settings, "sm"), onClick: () => SettingsDialog.open([new ProvidersModelsTab(), new ProxyTab()]), title: "Settings" })}
    </div>
  </div>
  <div class="flex flex-1 min-h-0 overflow-hidden">
<div id="sidebar-wrap" style="width:${sidebarWidth}px;min-width:150px;max-width:420px;flex-shrink:0;display:flex;flex-direction:column;overflow:hidden;transition:width 0.05s">
${sidebar}
</div>
<div id="sb-resize" style="width:5px;cursor:col-resize;flex-shrink:0;background:transparent;z-index:10;transition:background 0.15s"
  @mousedown=${(e: MouseEvent) => {
		e.preventDefault();
		const sx = e.clientX,
			sw = sidebarWidth;
		const mv = (me: MouseEvent) => {
			sidebarWidth = Math.max(150, Math.min(420, sw + me.clientX - sx));
			const w = document.getElementById("sidebar-wrap");
			if (w) w.style.width = sidebarWidth + "px";
		};
		const up = () => {
			document.removeEventListener("mousemove", mv);
			document.removeEventListener("mouseup", up);
			renderApp();
		};
		document.addEventListener("mousemove", mv);
		document.addEventListener("mouseup", up);
  }}
  @mouseenter=${(e: Event) => {
		(e.currentTarget as HTMLElement).style.background = "rgba(128,128,128,0.4)";
  }}
  @mouseleave=${(e: Event) => {
		(e.currentTarget as HTMLElement).style.background = "transparent";
  }}
></div>
<div class="flex flex-col flex-1 min-w-0 min-h-0 relative">
<div class="absolute inset-x-0 top-0 z-10 flex flex-col" style="bottom:130px;pointer-events:${hasMessages ? "none" : "auto"}" @suggestion=${handleSuggestion}>
<jae-empty-state .faded=${hasMessages} style="display:flex;flex-direction:column;flex:1;width:100%;min-height:0"></jae-empty-state>
</div>
<div id="chat-wrapper" class="flex flex-col flex-1 min-h-0" >
${chatPanel}
</div>
</div>
${
	rightPanel !== "none"
		? html`
<div id="rp-resize" style="width:5px;cursor:col-resize;flex-shrink:0;background:transparent;z-index:10;transition:background 0.15s"
  @mousedown=${(e: MouseEvent) => {
		e.preventDefault();
		const sx = e.clientX,
			sw = rightPanelWidth;
		const mv = (me: MouseEvent) => {
			rightPanelWidth = Math.max(280, Math.min(800, sw - (me.clientX - sx)));
			const p = document.getElementById("right-panel");
			if (p) p.style.width = rightPanelWidth + "px";
		};
		const up = () => {
			document.removeEventListener("mousemove", mv);
			document.removeEventListener("mouseup", up);
			renderApp();
		};
		document.addEventListener("mousemove", mv);
		document.addEventListener("mouseup", up);
  }}
  @mouseenter=${(e: Event) => {
		(e.currentTarget as HTMLElement).style.background = "rgba(128,128,128,0.4)";
  }}
  @mouseleave=${(e: Event) => {
		(e.currentTarget as HTMLElement).style.background = "transparent";
  }}
></div>
<div id="right-panel" class="flex flex-col border-l border-border" style="width:${rightPanelWidth}px;min-width:280px;max-width:800px;flex-shrink:0">
<div class="flex items-center gap-1 px-2 shrink-0 border-b border-border bg-muted/20" style="height:36px">
<button class="text-xs px-2 py-1 rounded ${
				rightPanel === "terminal"
					? "bg-primary text-primary-foreground"
					: "hover:bg-secondary text-muted-foreground"
			}" @click=${() => {
				rightPanel = "terminal";
				renderApp();
				requestAnimationFrame(() => {
					if (!terminalPanel) terminalPanel = document.querySelector("jae-terminal-panel") as JaeTerminalPanel;
					terminalPanel?.connect();
				});
			}}>Terminal</button>
<button class="text-xs px-2 py-1 rounded ${
				rightPanel === "browser" ? "bg-primary text-primary-foreground" : "hover:bg-secondary text-muted-foreground"
			}" @click=${() => {
				rightPanel = "browser";
				renderApp();
			}}>Browser</button>
<div class="flex-1"></div>
<button class="text-xs px-2 py-1 rounded hover:bg-secondary text-muted-foreground" @click=${() => {
				rightPanel = "none";
				renderApp();
			}} title="Close panel">✕</button>
</div>
${rightPanel === "terminal" ? html`<jae-terminal-panel class="flex-1 min-h-0"></jae-terminal-panel>` : html``}
${rightPanel === "browser" ? html`<jae-browser-panel  class="flex-1 min-h-0"></jae-browser-panel>` : html``}
</div>
`
		: html``
}
</div>
`,
		app,
	);
};

async function initApp() {
	const app = document.getElementById("app");
	if (!app) throw new Error("App container not found");
	render(
		html`<div class="w-full h-screen flex items-center justify-center bg-background text-foreground"><div class="text-muted-foreground">Loading...</div></div>`,
		app,
	);
	chatPanel = new ChatPanel();
	setupCommands();
	await refreshSidebar();
	const sessionIdFromUrl = new URLSearchParams(window.location.search).get("session");
	if (sessionIdFromUrl) {
		const loaded = await loadSession(sessionIdFromUrl);
		if (!loaded) {
			newSession();
			return;
		}
	} else {
		await createAgent();
	}
	renderApp();
}

initApp();