GridSnap/extension.js

const St = imports.gi.St;
const Main = imports.ui.main;
const Meta = imports.gi.Meta;
const Clutter = imports.gi.Clutter;
const Lang = imports.lang;
const Mainloop = imports.mainloop;
const Settings = imports.ui.settings;
const GLib = imports.gi.GLib;
const Gio = imports.gi.Gio;

let zoneManager;

// Storage path for custom layouts
const STORAGE_DIR = GLib.get_user_data_dir() + '/gridsnap';
const LAYOUTS_FILE = STORAGE_DIR + '/layouts.json';

// Default zone layouts
const DEFAULT_LAYOUTS = {
    'grid-2x2': {
        name: '2x2 Grid',
        zones: [
            { x: 0, y: 0, width: 0.5, height: 0.5 },      // Top-left
            { x: 0.5, y: 0, width: 0.5, height: 0.5 },    // Top-right
            { x: 0, y: 0.5, width: 0.5, height: 0.5 },    // Bottom-left
            { x: 0.5, y: 0.5, width: 0.5, height: 0.5 }   // Bottom-right
        ]
    },
    'columns-3': {
        name: '3 Columns',
        zones: [
            { x: 0, y: 0, width: 0.33, height: 1 },
            { x: 0.33, y: 0, width: 0.34, height: 1 },
            { x: 0.67, y: 0, width: 0.33, height: 1 }
        ]
    },
    'focus-left': {
        name: 'Focus Left',
        zones: [
            { x: 0, y: 0, width: 0.7, height: 1 },        // Main left
            { x: 0.7, y: 0, width: 0.3, height: 0.5 },    // Top-right
            { x: 0.7, y: 0.5, width: 0.3, height: 0.5 }   // Bottom-right
        ]
    }
};

function ZoneEditor() {
    this._init();
}

ZoneEditor.prototype = {
    _init: function() {
        this.editorOverlay = null;
        this.isEditing = false;
        this.zones = [];
        this.currentZone = null;
        this.startX = 0;
        this.startY = 0;
        this.isDragging = false;
    },
    
    startEditor: function() {
        if (this.isEditing) return;
        
        this.isEditing = true;
        this.zones = [];
        
        let monitor = Main.layoutManager.primaryMonitor;
        
        // Create editor overlay
        this.editorOverlay = new St.Widget({
            style_class: 'gridsnap-editor-overlay',
            reactive: true,
            x: monitor.x,
            y: monitor.y,
            width: monitor.width,
            height: monitor.height
        });
        
        this.editorOverlay.set_style(
            'background-color: rgba(0, 0, 0, 0.5);'
        );
        
        // Add instruction label
        this.instructionLabel = new St.Label({
            text: 'GridSnap Zone Editor\n\nDrag to draw zones | Ctrl+S to save | Ctrl+C to cancel | Delete to clear last zone',
            style_class: 'gridsnap-editor-instructions',
            x: 20,
            y: 20
        });
        this.instructionLabel.set_style(
            'color: white;' +
            'font-size: 18px;' +
            'background-color: rgba(0, 0, 0, 0.8);' +
            'padding: 15px;' +
            'border-radius: 8px;'
        );
        this.editorOverlay.add_child(this.instructionLabel);
        
        // Connect mouse events
        this.buttonPressId = this.editorOverlay.connect('button-press-event',
            Lang.bind(this, this._onButtonPress));
        this.buttonReleaseId = this.editorOverlay.connect('button-release-event',
            Lang.bind(this, this._onButtonRelease));
        this.motionId = this.editorOverlay.connect('motion-event',
            Lang.bind(this, this._onMotion));
        
        Main.layoutManager.addChrome(this.editorOverlay);
        Main.pushModal(this.editorOverlay);
        
        // Setup key listener for save/cancel
        this._setupEditorKeys();
    },
    
    _setupEditorKeys: function() {
        this.keyPressId = this.editorOverlay.connect('key-press-event',
            Lang.bind(this, function(actor, event) {
                let symbol = event.get_key_symbol();
                let state = event.get_state();
                
                // Ctrl+S to save
                if (symbol === Clutter.KEY_s && (state & Clutter.ModifierType.CONTROL_MASK)) {
                    this._saveLayout();
                    return Clutter.EVENT_STOP;
                }
                
                // Ctrl+C to cancel
                if (symbol === Clutter.KEY_c && (state & Clutter.ModifierType.CONTROL_MASK)) {
                    this._cancelEditor();
                    return Clutter.EVENT_STOP;
                }
                
                // Delete to remove last zone
                if (symbol === Clutter.KEY_Delete || symbol === Clutter.KEY_BackSpace) {
                    this._removeLastZone();
                    return Clutter.EVENT_STOP;
                }
                
                // Escape to cancel
                if (symbol === Clutter.KEY_Escape) {
                    this._cancelEditor();
                    return Clutter.EVENT_STOP;
                }
                
                return Clutter.EVENT_PROPAGATE;
            })
        );
    },
    
    _onButtonPress: function(actor, event) {
        let [x, y] = event.get_coords();
        let monitor = Main.layoutManager.primaryMonitor;
        
        this.startX = x - monitor.x;
        this.startY = y - monitor.y;
        this.isDragging = true;
        
        // Create zone preview
        this.currentZone = new St.Widget({
            style_class: 'gridsnap-editor-zone',
            x: this.startX,
            y: this.startY,
            width: 1,
            height: 1
        });
        this.currentZone.set_style(
            'border: 3px solid rgba(100, 255, 100, 0.9);' +
            'background-color: rgba(100, 255, 100, 0.3);' +
            'border-radius: 4px;'
        );
        this.editorOverlay.add_child(this.currentZone);
        
        return Clutter.EVENT_STOP;
    },
    
    _onMotion: function(actor, event) {
        if (!this.isDragging || !this.currentZone) return Clutter.EVENT_PROPAGATE;
        
        let [x, y] = event.get_coords();
        let monitor = Main.layoutManager.primaryMonitor;
        
        let endX = x - monitor.x;
        let endY = y - monitor.y;
        
        let rectX = Math.min(this.startX, endX);
        let rectY = Math.min(this.startY, endY);
        let rectWidth = Math.abs(endX - this.startX);
        let rectHeight = Math.abs(endY - this.startY);
        
        this.currentZone.set_position(rectX, rectY);
        this.currentZone.set_size(rectWidth, rectHeight);
        
        return Clutter.EVENT_STOP;
    },
    
    _onButtonRelease: function(actor, event) {
        if (!this.isDragging) return Clutter.EVENT_PROPAGATE;
        
        this.isDragging = false;
        
        if (this.currentZone) {
            let monitor = Main.layoutManager.primaryMonitor;
            let width = this.currentZone.width;
            let height = this.currentZone.height;
            
            // Only save zones with reasonable size
            if (width > 20 && height > 20) {
                // Convert to relative coordinates
                let zone = {
                    x: this.currentZone.x / monitor.width,
                    y: this.currentZone.y / monitor.height,
                    width: width / monitor.width,
                    height: height / monitor.height
                };
                
                this.zones.push(zone);
                
                // Add zone number label
                let label = new St.Label({
                    text: String(this.zones.length),
                    style_class: 'gridsnap-editor-zone-label',
                    x: 10,
                    y: 10
                });
                label.set_style(
                    'color: white;' +
                    'font-size: 24px;' +
                    'font-weight: bold;' +
                    'text-shadow: 2px 2px 4px rgba(0,0,0,0.8);'
                );
                this.currentZone.add_child(label);
            } else {
                // Remove if too small
                this.currentZone.destroy();
            }
            
            this.currentZone = null;
        }
        
        return Clutter.EVENT_STOP;
    },
    
    _removeLastZone: function() {
        if (this.zones.length > 0) {
            this.zones.pop();
            // Redraw editor
            this._cancelEditor();
            this.startEditor();
            
            // Re-add existing zones
            let monitor = Main.layoutManager.primaryMonitor;
            this.zones.forEach((zone, index) => {
                let zoneActor = new St.Widget({
                    style_class: 'gridsnap-editor-zone',
                    x: zone.x * monitor.width,
                    y: zone.y * monitor.height,
                    width: zone.width * monitor.width,
                    height: zone.height * monitor.height
                });
                zoneActor.set_style(
                    'border: 3px solid rgba(100, 255, 100, 0.9);' +
                    'background-color: rgba(100, 255, 100, 0.3);' +
                    'border-radius: 4px;'
                );
                
                let label = new St.Label({
                    text: String(index + 1),
                    x: 10,
                    y: 10
                });
                label.set_style(
                    'color: white;' +
                    'font-size: 24px;' +
                    'font-weight: bold;' +
                    'text-shadow: 2px 2px 4px rgba(0,0,0,0.8);'
                );
                zoneActor.add_child(label);
                this.editorOverlay.add_child(zoneActor);
            });
        }
    },
    
    _saveLayout: function() {
        if (this.zones.length === 0) {
            Main.notify('GridSnap', 'No zones to save!');
            this._cancelEditor();
            return;
        }
        
        // Generate layout name
        let layoutId = 'custom-' + Date.now();
        let layoutName = 'Custom Layout (' + this.zones.length + ' zones)';
        
        // Add to zoneManager's layouts
        if (zoneManager) {
            zoneManager.layouts[layoutId] = {
                name: layoutName,
                zones: this.zones
            };
            zoneManager.currentLayout = layoutId;

            // Persist to file for permanent storage
            zoneManager._saveLayouts();

            Main.notify('GridSnap', 'Layout saved: ' + layoutName);
        }
        
        this._cancelEditor();
    },
    
    _cancelEditor: function() {
        if (!this.isEditing) return;
        
        if (this.editorOverlay) {
            if (this.buttonPressId) this.editorOverlay.disconnect(this.buttonPressId);
            if (this.buttonReleaseId) this.editorOverlay.disconnect(this.buttonReleaseId);
            if (this.motionId) this.editorOverlay.disconnect(this.motionId);
            if (this.keyPressId) this.editorOverlay.disconnect(this.keyPressId);
            
            Main.popModal(this.editorOverlay);
            Main.layoutManager.removeChrome(this.editorOverlay);
            this.editorOverlay.destroy();
            this.editorOverlay = null;
        }
        
        this.isEditing = false;
        this.zones = [];
        this.currentZone = null;
    }
};

function ZoneManager() {
    this._init();
}

ZoneManager.prototype = {
    _init: function() {
        this.overlay = null;
        this.currentLayout = 'grid-2x2';
        this.layouts = {};
        this.isShowing = false;
        this.dragInProgress = false;
        this.draggedWindow = null;
        this.editor = new ZoneEditor();

        // Initialize settings
        this.settings = new Settings.ExtensionSettings(this, 'gridsnap@cinnamon-extension');

        // Load layouts (default + saved custom layouts)
        this._loadLayouts();

        // Connect to window grab events
        this._connectSignals();
        this._setupKeybindings();
    },
    
    _connectSignals: function() {
        // Monitor for window grab begin/end
        this.grabOpBeginId = global.display.connect('grab-op-begin',
            Lang.bind(this, this._onGrabBegin));
        this.grabOpEndId = global.display.connect('grab-op-end',
            Lang.bind(this, this._onGrabEnd));

        // Monitor for modifier key changes
        this.modifierPollTimeoutId = null;
    },
    
    _setupKeybindings: function() {
        // Add keybinding to show zones overlay (Super+Z)
        Main.keybindingManager.addHotKey(
            'show-zones',
            '<Super>z',
            Lang.bind(this, this._toggleZonesOverlay)
        );

        // Add keybinding to open zone editor (Super+Shift+E)
        Main.keybindingManager.addHotKey(
            'open-zone-editor',
            '<Super><Shift>e',
            Lang.bind(this, function() { this.editor.startEditor(); })
        );

        // Add keybindings for quick snap to zones (Super+Ctrl+1-9)
        for (let i = 1; i <= 9; i++) {
            const zoneIndex = i - 1;  // Capture the value for this iteration
            Main.keybindingManager.addHotKey(
                'snap-to-zone-' + i,
                '<Super><Ctrl>' + i,
                Lang.bind(this, function() {
                    this._snapToZone(zoneIndex);
                })
            );
        }
    },

    _loadLayouts: function() {
        // Start with default layouts
        for (let layoutId in DEFAULT_LAYOUTS) {
            this.layouts[layoutId] = DEFAULT_LAYOUTS[layoutId];
        }

        // Try to load custom layouts from file
        try {
            let file = Gio.File.new_for_path(LAYOUTS_FILE);
            if (file.query_exists(null)) {
                let [success, contents] = file.load_contents(null);
                if (success) {
                    let customLayouts = JSON.parse(contents.toString());
                    // Merge custom layouts with defaults
                    for (let layoutId in customLayouts) {
                        this.layouts[layoutId] = customLayouts[layoutId];
                    }
                    global.log('GridSnap: Loaded ' + Object.keys(customLayouts).length + ' custom layout(s)');
                }
            }
        } catch(e) {
            global.logError('GridSnap: Error loading custom layouts - ' + e);
        }
    },

    _saveLayouts: function() {
        try {
            // Create directory if it doesn't exist
            let dir = Gio.File.new_for_path(STORAGE_DIR);
            if (!dir.query_exists(null)) {
                dir.make_directory_with_parents(null);
            }

            // Extract only custom layouts (those starting with 'custom-')
            let customLayouts = {};
            for (let layoutId in this.layouts) {
                if (layoutId.startsWith('custom-')) {
                    customLayouts[layoutId] = this.layouts[layoutId];
                }
            }

            // Save to file
            let file = Gio.File.new_for_path(LAYOUTS_FILE);
            let contents = JSON.stringify(customLayouts, null, 2);
            file.replace_contents(contents, null, false, Gio.FileCreateFlags.REPLACE_DESTINATION, null);

            global.log('GridSnap: Saved ' + Object.keys(customLayouts).length + ' custom layout(s) to ' + LAYOUTS_FILE);
        } catch(e) {
            global.logError('GridSnap: Error saving custom layouts - ' + e);
        }
    },
    
    _checkModifiersDuringDrag: function() {
        if (!this.dragInProgress) {
            return false; // Stop polling
        }

        // Check if shift-drag is enabled
        if (!this.settings.getValue('enable-shift-drag')) {
            return true; // Continue polling but don't show overlay
        }

        let pointerInfo = global.get_pointer();
        let mods = pointerInfo[2];
        let shiftPressed = !!(mods & Clutter.ModifierType.SHIFT_MASK);

        if (shiftPressed && !this.isShowing) {
            this._showZonesOverlay();
        } else if (!shiftPressed && this.isShowing) {
            this._hideZonesOverlay();
        }

        return true; // Continue polling
    },

    _onGrabBegin: function(display, window) {
        // Window drag has begun
        this.dragInProgress = true;

        // Get the actual window being dragged - use focus_window as it's the dragged window
        this.draggedWindow = global.display.focus_window;

        // Start polling for modifier key changes during drag
        if (this.modifierPollTimeoutId) {
            Mainloop.source_remove(this.modifierPollTimeoutId);
        }
        this.modifierPollTimeoutId = Mainloop.timeout_add(50, Lang.bind(this, this._checkModifiersDuringDrag));
    },

    _onGrabEnd: function(display, window) {
        if (this.dragInProgress) {
            this.dragInProgress = false;

            // Stop polling for modifiers
            if (this.modifierPollTimeoutId) {
                Mainloop.source_remove(this.modifierPollTimeoutId);
                this.modifierPollTimeoutId = null;
            }

            if (this.isShowing) {
                // Snap to zone if cursor is over one
                this._snapWindowToZone(this.draggedWindow);
                this._hideZonesOverlay();
            }

            this.draggedWindow = null;
        }
    },
    
    _toggleZonesOverlay: function() {
        if (this.isShowing) {
            this._hideZonesOverlay();
        } else {
            this._showZonesOverlay();
        }
    },
    
    _showZonesOverlay: function() {
        if (this.isShowing) return;
        
        let monitor = Main.layoutManager.primaryMonitor;
        
        // Create overlay container
        this.overlay = new St.Widget({
            style_class: 'gridsnap-overlay',
            reactive: false,
            x: monitor.x,
            y: monitor.y,
            width: monitor.width,
            height: monitor.height
        });
        
        // Add semi-transparent background
        this.overlay.set_style(
            'background-color: rgba(0, 0, 0, 0.3);'
        );
        
        // Draw zones
        let layout = this.layouts[this.currentLayout];
        layout.zones.forEach((zone, index) => {
            let zoneActor = this._createZoneActor(zone, monitor, index);
            this.overlay.add_child(zoneActor);
        });
        
        Main.layoutManager.addChrome(this.overlay);
        this.isShowing = true;
    },
    
    _createZoneActor: function(zone, monitor, index) {
        let x = monitor.width * zone.x;
        let y = monitor.height * zone.y;
        let width = monitor.width * zone.width;
        let height = monitor.height * zone.height;

        // Get settings
        let borderWidth = this.settings.getValue('zone-border-width');
        let borderColor = this.settings.getValue('zone-border-color');
        let fillColor = this.settings.getValue('zone-fill-color');
        let opacity = this.settings.getValue('zone-opacity');
        let showNumbers = this.settings.getValue('show-zone-numbers');

        // Adjust fill color opacity
        let fillColorWithOpacity = fillColor.replace(/[\d.]+\)$/g, (opacity / 100) + ')');

        let actor = new St.Widget({
            style_class: 'gridsnap-zone',
            x: x,
            y: y,
            width: width,
            height: height
        });

        actor.set_style(
            'border: ' + borderWidth + 'px solid ' + borderColor + ';' +
            'background-color: ' + fillColorWithOpacity + ';' +
            'border-radius: 4px;'
        );

        // Add zone number label if enabled
        if (showNumbers) {
            let label = new St.Label({
                text: String(index + 1),
                style_class: 'gridsnap-zone-label',
                x: 10,
                y: 10
            });
            label.set_style(
                'color: white;' +
                'font-size: 24px;' +
                'font-weight: bold;' +
                'text-shadow: 2px 2px 4px rgba(0,0,0,0.8);'
            );
            actor.add_child(label);
        }

        return actor;
    },
    
    _hideZonesOverlay: function() {
        if (!this.isShowing) return;
        
        if (this.overlay) {
            Main.layoutManager.removeChrome(this.overlay);
            this.overlay.destroy();
            this.overlay = null;
        }
        
        this.isShowing = false;
    },
    
    _snapWindowToZone: function(window) {
        if (!window) return;

        let [x, y] = global.get_pointer();
        let monitor = Main.layoutManager.primaryMonitor;

        // Convert to relative coordinates
        let relX = (x - monitor.x) / monitor.width;
        let relY = (y - monitor.y) / monitor.height;

        // Find which zone the cursor is in
        let layout = this.layouts[this.currentLayout];
        let targetZone = null;

        for (let zone of layout.zones) {
            if (relX >= zone.x && relX < zone.x + zone.width &&
                relY >= zone.y && relY < zone.y + zone.height) {
                targetZone = zone;
                break;
            }
        }

        if (targetZone) {
            this._moveWindowToZone(window, targetZone, monitor);

            // Show notification if enabled
            if (this.settings.getValue('notification-on-snap')) {
                Main.notify('GridSnap', 'Window snapped to zone');
            }
        }
    },
    
    _snapToZone: function(zoneIndex) {
        // Check if keyboard snap is enabled
        if (!this.settings.getValue('enable-keyboard-snap')) {
            return;
        }

        // Snap the focused window to a specific zone
        let window = global.display.focus_window;
        if (!window) return;

        let layout = this.layouts[this.currentLayout];
        if (zoneIndex >= layout.zones.length) return;

        let monitor = Main.layoutManager.primaryMonitor;
        this._moveWindowToZone(window, layout.zones[zoneIndex], monitor);

        // Show notification if enabled
        if (this.settings.getValue('notification-on-snap')) {
            Main.notify('GridSnap', 'Window snapped to zone ' + (zoneIndex + 1));
        }
    },
    
    _moveWindowToZone: function(window, zone, monitor) {
        if (!window || !zone || !monitor) return;

        // Unmaximize if maximized
        if (window.maximized_horizontally || window.maximized_vertically) {
            window.unmaximize(Meta.MaximizeFlags.BOTH);
        }

        // Calculate absolute coordinates
        let x = monitor.x + Math.round(monitor.width * zone.x);
        let y = monitor.y + Math.round(monitor.height * zone.y);
        let width = Math.round(monitor.width * zone.width);
        let height = Math.round(monitor.height * zone.height);

        // Use Mainloop.idle_add to defer the operation
        // This ensures the unmaximize completes before move/resize
        Mainloop.idle_add(Lang.bind(this, function() {
            // Move and resize window (false = programmatic operation, not user-initiated)
            window.move_resize_frame(false, x, y, width, height);
            return false; // Don't repeat
        }));
    },
    
    cycleLayout: function() {
        let layouts = Object.keys(this.layouts);
        let currentIndex = layouts.indexOf(this.currentLayout);
        let nextIndex = (currentIndex + 1) % layouts.length;
        this.currentLayout = layouts[nextIndex];
        
        Main.notify('GridSnap', 'Layout: ' + this.layouts[this.currentLayout].name);
        
        // Refresh overlay if showing
        if (this.isShowing) {
            this._hideZonesOverlay();
            this._showZonesOverlay();
        }
    },
    
    destroy: function() {
        // Clean up overlay FIRST to prevent corrupt display
        this._hideZonesOverlay();

        // Stop modifier polling if active
        if (this.modifierPollTimeoutId) {
            Mainloop.source_remove(this.modifierPollTimeoutId);
            this.modifierPollTimeoutId = null;
        }

        // Clean up editor
        if (this.editor) {
            this.editor._cancelEditor();
        }

        // Clean up settings
        if (this.settings) {
            this.settings.finalize();
            this.settings = null;
        }

        // Disconnect signals
        if (this.grabOpBeginId) {
            global.display.disconnect(this.grabOpBeginId);
            this.grabOpBeginId = null;
        }
        if (this.grabOpEndId) {
            global.display.disconnect(this.grabOpEndId);
            this.grabOpEndId = null;
        }

        // Remove keybindings
        try {
            Main.keybindingManager.removeHotKey('show-zones');
            Main.keybindingManager.removeHotKey('open-zone-editor');
            for (let i = 1; i <= 9; i++) {
                Main.keybindingManager.removeHotKey('snap-to-zone-' + i);
            }
        } catch(e) {
            global.logError('GridSnap: Error removing keybindings - ' + e);
        }
    }
};

function init(metadata) {
    // Extension initialization
}

function enable() {
    zoneManager = new ZoneManager();
    
    // Add keybinding to cycle layouts
    Main.keybindingManager.addHotKey(
        'cycle-layout',
        '<Super><Shift>z',
        Lang.bind(zoneManager, zoneManager.cycleLayout)
    );
}

function disable() {
    if (zoneManager) {
        Main.keybindingManager.removeHotKey('cycle-layout');
        zoneManager.destroy();
        zoneManager = null;
    }
}