import Gio from "gi://Gio";
import GLib from "gi://GLib";
import { createState } from "ags";

// Types (matching niri-ipc 26.4.0)

export interface NiriWorkspace {
    id: number;
    idx: number;
    name: string | null;
    output: string | null;
    is_urgent: boolean;
    is_active: boolean;
    is_focused: boolean;
    active_window_id: number | null;
}

/**
 * Discriminated union of all niri compositor events.
 * In practice we parse JSON dynamically so the union is for documentation;
 * consumers can narrow with `"WorkspacesChanged" in event` etc.
 */
export type NiriEvent =
    | { WorkspacesChanged: { workspaces: NiriWorkspace[] } }
    | { WorkspaceActivated: { id: number; focused: boolean } }

type EventHandler = (event: NiriEvent) => void;

// Client
export class NiriEventClient {
    private socket: Gio.SocketConnection | null = null;
    private input: Gio.DataInputStream | null = null;
    private cancel: Gio.Cancellable | null = null;
    private handlers: EventHandler[] = [];
    private reconnectBackoff = 1000;
    private reconnectTimer: number | null = null;
    private running = false;

    // Reactive state

    private setWorkspaces: (ws: NiriWorkspace[]) => void;
    private setActiveId: (id: number | null) => void;
    private setFocusedId: (id: number | null) => void;

    /** Reactive list of all workspaces (sort by `.idx` in the widget). */
    workspaces: () => NiriWorkspace[];

    /** Reactive ID of the currently active workspace on the focused output. */
    activeWorkspaceId: () => number | null;

    /** Reactive ID of the globally focused workspace (there is only one). */
    focusedWorkspaceId: () => number | null;

    constructor() {
        const [w, sw] = createState<NiriWorkspace[]>([]);
        const [a, sa] = createState<number | null>(null);
        const [f, sf] = createState<number | null>(null);
        this.workspaces = w;
        this.setWorkspaces = sw;
        this.activeWorkspaceId = a;
        this.setActiveId = sa;
        this.focusedWorkspaceId = f;
        this.setFocusedId = sf;
    }

    // Public API

    /**
     * Subscribe to all niri events.
     * Returns an unsubscribe function.
     */
    onEvent(handler: EventHandler): () => void {
        this.handlers.push(handler);
        return () => {
            const i = this.handlers.indexOf(handler);
            if (i >= 0) this.handlers.splice(i, 1);
        };
    }

    /** Open the UNIX socket and start the event stream. Idempotent. */
    connect(): void {
        if (this.running) return;
        this.running = true;
        this.reconnectBackoff = 1000;
        this.doConnect();
    }

    /** Close the socket and stop reconnection. */
    disconnect(): void {
        this.running = false;
        if (this.reconnectTimer !== null) {
            GLib.source_remove(this.reconnectTimer);
            this.reconnectTimer = null;
        }
        this.closeSocket();
    }

    // Internals

    private get socketPath(): string {
        return GLib.getenv("NIRI_SOCKET") ?? "";
    }

    private closeSocket(): void {
        this.cancel?.cancel();
        this.cancel = null;
        try {
            this.input?.close(null);
        } catch (_) {
            /* ignore */
        }
        try {
            this.socket?.close(null);
        } catch (_) {
            /* ignore */
        }
        this.input = null;
        this.socket = null;
    }

    private async doConnect(): Promise<void> {
        this.closeSocket();

        const path = this.socketPath;
        if (!path) {
            console.error("[NiriClient] NIRI_SOCKET not set");
            this.scheduleReconnect();
            return;
        }

        try {
            const client = new Gio.SocketClient();
            const address = Gio.UnixSocketAddress.new(path);
            this.cancel = new Gio.Cancellable();

            this.socket = await new Promise<Gio.SocketConnection>(
                (resolve, reject) => {
                    client.connect_async(address, this.cancel!, (src, result) => {
                        try {
                            resolve(src!.connect_finish(result));
                        } catch (e) {
                            reject(e);
                        }
                    });
                },
            );

            const istream = this.socket.get_input_stream();
            this.input = new Gio.DataInputStream({
                base_stream: istream,
                close_base_stream: false,
            });

            // Send EventStream request
            const ostream = this.socket.get_output_stream();
            const request = JSON.stringify({ EventStream: null }) + "\n";
            ostream.write_bytes(new GLib.Bytes(request), null);
            ostream.flush(null);

            // Read the initial Reply::Ok(Response::Handled)
            const replyLine = await this.readLine();
            if (!replyLine) throw new Error("No reply from niri to EventStream");

            const reply = JSON.parse(replyLine);
            if ("Ok" in reply) {
                console.log("[NiriClient] Event stream connected");
            } else {
                console.warn(
                    "[NiriClient] Unexpected reply to EventStream:",
                    replyLine,
                );
            }

            // Reset backoff on a clean connection
            this.reconnectBackoff = 1000;

            // Enter the read loop
            await this.readLoop();
        } catch (e: any) {
            if (!this.running) return; // intentional disconnect
            console.error(
                "[NiriClient] Connection failed:",
                e.message ?? String(e),
            );
            this.closeSocket();
            this.scheduleReconnect();
        }
    }

    private async readLoop(): Promise<void> {
        while (this.running && this.input) {
            try {
                const line = await this.readLine();
                if (line === null) break; // EOF

                const event = JSON.parse(line) as NiriEvent;
                this.processEvent(event);
                this.emit(event);
            } catch (e: any) {
                if (!this.running) return;
                if (
                    e.matches?.(Gio.IOErrorEnum, Gio.IOErrorEnum.CANCELLED)
                ) {
                    return;
                }
                console.warn(
                    "[NiriClient] Event parse error:",
                    e.message ?? String(e),
                );
            }
        }

        if (this.running) {
            console.log("[NiriClient] Stream ended, reconnecting...");
            this.closeSocket();
            this.scheduleReconnect();
        }
    }

    /**
     * Update reactive workspace state from an event.
     *
     * WorkspacesChanged is the ground truth — full replacement.
     * WorkspaceActivated is an incremental update for instant UI.
     */
    private processEvent(event: NiriEvent): void {
        if ("WorkspacesChanged" in event) {
            const ws = event.WorkspacesChanged.workspaces;
            this.setWorkspaces(ws);
            const focused = ws.find(w => w.is_focused);
            this.setFocusedId(focused?.id ?? null);
        } else if ("WorkspaceActivated" in event) {
            const { id, focused } = event.WorkspaceActivated;
            const current = this.workspaces();

            // Find the activated workspace to learn its output
            const activated = current.find(w => w.id === id);
            const output = activated?.output ?? null;

            const updated = current.map(w => {
                if (w.id === id) {
                    return {
                        ...w,
                        is_active: true,
                        is_focused: focused || w.is_focused,
                    };
                }
                // Clear is_active for same-output siblings
                if (output && w.output === output && w.is_active) {
                    return { ...w, is_active: false };
                }
                // Clear is_focused globally when this activation is a focus change
                if (focused && w.is_focused) {
                    return { ...w, is_focused: false };
                }
                return w;
            });

            this.setWorkspaces(updated);
            if (focused) this.setFocusedId(id);
        }
    }

    private emit(event: NiriEvent): void {
        for (const h of this.handlers) {
            try {
                h(event);
            } catch (e) {
                console.error("[NiriClient] Handler error:", e);
            }
        }
    }

    private readLine(): Promise<string | null> {
        return new Promise((resolve, reject) => {
            if (!this.input) return resolve(null);
            this.input.read_line_async(
                GLib.PRIORITY_DEFAULT,
                this.cancel,
                (src, result) => {
                    try {
                        if (!src || !this.running) return resolve(null);
                        const [bytes] = src.read_line_finish(result);
                        if (bytes === null) return resolve(null);
                        resolve(new TextDecoder().decode(bytes).trim());
                    } catch (e: any) {
                        if (
                            e.matches?.(
                                Gio.IOErrorEnum,
                                Gio.IOErrorEnum.CANCELLED,
                            )
                        ) {
                            resolve(null);
                        } else {
                            reject(e);
                        }
                    }
                },
            );
        });
    }

    private scheduleReconnect(): void {
        if (!this.running) return;
        console.log(
            `[NiriClient] Reconnecting in ${this.reconnectBackoff}ms…`,
        );
        this.reconnectTimer = GLib.timeout_add(
            GLib.PRIORITY_DEFAULT,
            this.reconnectBackoff,
            () => {
                this.reconnectTimer = null;
                if (this.running) this.doConnect();
                return GLib.SOURCE_REMOVE;
            },
        );
        this.reconnectBackoff = Math.min(this.reconnectBackoff * 2, 30_000);
    }
}