import { Gtk } from "ags/gtk4";
import { createState, createEffect } from "ags";
import Pango from "gi://Pango";
import { streamCompletion, type Message, type Backend, getConfig, MODELS } from "../../lib/llm";

function MessageLabel(msg: Message) {
    let displayContent = msg.content;
    if (msg.role === "assistant" && !displayContent) {
        displayContent = "⟳";
    }
    const label = new Gtk.Label({
        label: displayContent,
        wrap: true,
        wrap_mode: Pango.WrapMode.WORD_CHAR,
        selectable: true,
        use_markup: false,
        max_width_chars: 70,
    });
    label.add_css_class("message");
    label.add_css_class(msg.role);
    if (msg.role === "assistant" && !msg.content) {
        label.add_css_class("thinking");
    }
    return label;
}

function EmptyStateContent() {
    const box = new Gtk.Box({
        orientation: Gtk.Orientation.VERTICAL,
        halign: Gtk.Align.CENTER,
        valign: Gtk.Align.CENTER,
        hexpand: true,
        vexpand: true,
        cssClasses: ["empty-state"],
    });

    const picture = Gtk.Picture.new_for_filename("assets/ai-assistant.svg");

    const title = new Gtk.Label({ label: "", cssClasses: ["empty-title"] });
    const subtitle = new Gtk.Label({ label: "What do you need?", cssClasses: ["empty-subtitle"] });
    box.append(picture);
    box.append(title);
    box.append(subtitle);
    return box;
}

export default function LLMChat() {
    const [messages, setMessages] = createState<Message[]>([]);
    const [input, setInput] = createState("");
    const [loading, setLoading] = createState(false);
    const [backend, setBackend] = createState<Backend>("deepseek");
    const [model, setModel] = createState(MODELS["deepseek"][0]);

    let scrolled: Gtk.ScrolledWindow | undefined;
    let messagesBox: Gtk.Box | undefined;
    let stack: Gtk.Stack | undefined;
    let entryRef: Gtk.Entry | undefined;
    let modelPopover: Gtk.Popover | undefined;
    let modelBtn: Gtk.Button | undefined;
    let deepseekBtn: Gtk.Button | undefined;
    let ollamaBtn: Gtk.Button | undefined;

    // ── model popover ──────────────────────────────────────────────────────

    function buildModelPopover(): Gtk.Popover {
        const listBox = new Gtk.Box({
            orientation: Gtk.Orientation.VERTICAL,
            spacing: 4,
            cssClasses: ["model-list"],
        });

        for (const m of MODELS[backend()]) {
            const btn = new Gtk.Button({
                label: m,
                halign: Gtk.Align.FILL,
                cssClasses: m === model() ? ["model-item", "active"] : ["model-item"],
            });
            btn.connect("clicked", () => {
                setModel(m);
                modelPopover?.popdown();
            });
            listBox.append(btn);
        }

        const pop = new Gtk.Popover({
            has_arrow: false,
            cssClasses: ["model-popover"],
        });
        pop.set_child(listBox);
        return pop;
    }

    function openModelPopover() {
        if (!modelBtn) return;

        // Rebuild popover so it reflects current backend models + active state
        if (modelPopover) {
            modelPopover.unparent();
            modelPopover = undefined;
        }
        modelPopover = buildModelPopover();
        modelPopover.set_parent(modelBtn);
        modelPopover.popup();
    }

    // ── backend switch ─────────────────────────────────────────────────────

    function switchBackend(b: Backend) {
        setBackend(b);
        // Reset model to the first model of the new backend
        setModel(MODELS[b][0]);
    }

    // Update model button label whenever backend or model changes
    createEffect(() => {
        const _bk = backend();
        const _md = model();
        if (modelBtn) {
            modelBtn.label = _md;
        }
    });

    // Update backend toggle button CSS classes whenever backend changes
    createEffect(() => {
        const _bk = backend();
        if (deepseekBtn) {
            deepseekBtn.css_classes = _bk === "deepseek" ? ["toggle-btn", "active"] : ["toggle-btn"];
        }
        if (ollamaBtn) {
            ollamaBtn.css_classes = _bk === "ollama" ? ["toggle-btn", "active"] : ["toggle-btn"];
        }
    });

    // ── scroll helper ──────────────────────────────────────────────────────

    const scrollToBottom = () => {
        if (!scrolled) return;
        const adj = scrolled.vadjustment;
        adj.value = adj.upper - adj.page_size;
    };

    // ── message rendering ──────────────────────────────────────────────────

    createEffect(() => {
        const msgs = messages();
        if (!messagesBox) return;

        // Clear existing messages
        let child = messagesBox.get_first_child();
        while (child) {
            const next = child.get_next_sibling();
            messagesBox.remove(child);
            child = next;
        }

        // Add current messages
        for (const msg of msgs) {
            const rowBox = new Gtk.Box({
                orientation: Gtk.Orientation.HORIZONTAL,
                hexpand: true,
                cssClasses: ["message-row", msg.role],
            });

            const msgLabel = MessageLabel(msg);
            if (msg.role === "user") {
                rowBox.set_halign(Gtk.Align.END);
                rowBox.append(msgLabel);
            } else {
                rowBox.set_halign(Gtk.Align.START);
                rowBox.append(msgLabel);
            }
            messagesBox.append(rowBox);
        }
        scrollToBottom();

        // Switch stack page
        if (stack) {
            stack.set_visible_child_name(msgs.length === 0 ? "empty" : "chat");
        }
    });

    // ── send ───────────────────────────────────────────────────────────────

    const send = async () => {
        const text = input().trim();
        if (!text || loading()) return;

        const userMsgs: Message[] = [...messages(), { role: "user", content: text }];
        setMessages([...userMsgs, { role: "assistant", content: "" }]);
        setInput("");

        // Force‑clear the entry widget to work around the binding glitch
        entryRef?.set_text("");

        setLoading(true);

        const cfg = getConfig(backend(), model());

        await streamCompletion(
            cfg,
            userMsgs,
            (chunk) => {
                const all = messages();
                const last = all[all.length - 1];
                setMessages([...all.slice(0, -1), { ...last, content: last.content + chunk }]);
            },
            () => setLoading(false),
            (err) => {
                const all = messages();
                setMessages([...all.slice(0, -1), { role: "assistant", content: `⚠ ${err}` }]);
                setLoading(false);
            },
        );
    };

    // ── UI ─────────────────────────────────────────────────────────────────

    return (
        <box orientation={Gtk.Orientation.VERTICAL} vexpand={true} hexpand={true} cssClasses={["llm-chat", "grid-card"]}>
            <scrolledwindow
                vexpand={true}
                hexpand={true}
                cssClasses={["chat-scroll"]}
                onRealize={(self: Gtk.ScrolledWindow) => { scrolled = self; }}
            >
                <stack
                    transition_type={Gtk.StackTransitionType.CROSSFADE}
                    onRealize={(self: Gtk.Stack) => {
                        stack = self;
                        // Add empty state page
                        const emptyPage = EmptyStateContent();
                        self.add_named(emptyPage, "empty");
                        // Add chat page
                        const chatBox = new Gtk.Box({
                            orientation: Gtk.Orientation.VERTICAL,
                            spacing: 10,
                            cssClasses: ["messages"],
                        });
                        self.add_named(chatBox, "chat");
                        messagesBox = chatBox;
                        // Start with empty state visible
                        self.set_visible_child_name("empty");
                    }}
                />
            </scrolledwindow>

            {/* ── input row ── */}
            <box spacing={10} cssClasses={["input-row"]}>
                {/* ── toolbar row: backend toggle + model selector ── */}
                <box orientation={Gtk.Orientation.VERTICAL} spacing={4} cssClasses={["toolbar-row"]}>
                    {/* Backend toggle */}
                    <box cssClasses={["backend-toggle"]}>
                        <button
                            cssClasses={backend() === "deepseek" ? ["toggle-btn", "active"] : ["toggle-btn"]}
                            label="DeepSeek"
                            onRealize={(self) => { deepseekBtn = self; }}
                            onClicked={() => switchBackend("deepseek")}
                        />
                        <button
                            cssClasses={backend() === "ollama" ? ["toggle-btn", "active"] : ["toggle-btn"]}
                            label="Ollama"
                            onRealize={(self) => { ollamaBtn = self; }}
                            onClicked={() => switchBackend("ollama")}
                        />
                    </box>

                    {/* Model selector */}
                    <button
                        cssClasses={["model-selector"]}
                        label={model()}
                        onRealize={(self) => { modelBtn = self; }}
                        onClicked={openModelPopover}
                    />
                </box>

                <entry
                    hexpand={true}
                    placeholderText="Ask me anything..."
                    text={input()}
                    sensitive={!loading()}
                    cssClasses={["chat-input"]}
                    onChanged={(self) => setInput(self.text)}
                    onActivate={send}
                />
                <button
                    cssClasses={loading() ? ["send-button", "loading"] : ["send-button"]}
                    label={loading() ? "⏳" : "➤"}
                    sensitive={!loading()}
                    onClicked={send}
                />
            </box>
        </box>
    );
}