# Thanks to https://github.com/muhchaudhary for the original code. You are a legend.
import json

import cairo
import gi
from fabric.hyprland.service import Hyprland
from fabric.utils.helpers import get_desktop_applications
from fabric.widgets.box import Box
from fabric.widgets.button import Button
from fabric.widgets.eventbox import EventBox
from fabric.widgets.image import Image
from fabric.widgets.label import Label
from fabric.widgets.overlay import Overlay
from loguru import logger

import config.data as data
import modules.icons as icons
# WIP icon resolver (app_id to guessing the icon name)
from utils.icon_resolver import IconResolver

gi.require_version("Gtk", "3.0")
from gi.repository import Gdk, Gtk

screen = Gdk.Screen.get_default()
CURRENT_WIDTH = screen.get_width()
CURRENT_HEIGHT = screen.get_height()

icon_resolver = IconResolver()
connection = Hyprland()
BASE_SCALE = 0.1  # Base scale factor for overview

# Credit to Aylur for the drag and drop code
TARGET = [Gtk.TargetEntry.new("text/plain", Gtk.TargetFlags.SAME_APP, 0)]

# Credit to Aylur for the createSurfaceFromWidget code
def createSurfaceFromWidget(widget: Gtk.Widget) -> cairo.ImageSurface:
    alloc = widget.get_allocation()
    surface = cairo.ImageSurface(
        cairo.Format.ARGB32,
        alloc.width,
        alloc.height,
    )
    cr = cairo.Context(surface)
    cr.set_source_rgba(255, 255, 255, 0)
    cr.rectangle(0, 0, alloc.width, alloc.height)
    cr.fill()
    widget.draw(cr)
    return surface


class HyprlandWindowButton(Button):
    def __init__(
        self,
        window: Box,
        title: str,
        address: str,
        app_id: str,
        size,
        transform: int = 0,
    ):
        self.transform = transform % 4
        self.size = size if transform in [0, 2] else (size[1], size[0])
        self.address = address
        self.app_id = app_id
        self.title = title
        self.window: Box = window

        # Compute dynamic icon sizes based on the button size.
        # Using the minimum dimension of the button for scaling.
        icon_size_main = int(min(self.size) * 0.5)  # adjust factor as needed

        # Enhanced icon resolution using desktop apps
        desktop_app = window.find_app(app_id)
        
        # Get icon using improved method with fallbacks
        icon_pixbuf = None
        if desktop_app:
            icon_pixbuf = desktop_app.get_icon_pixbuf(size=icon_size_main)
        
        if not icon_pixbuf:
            # Fallback to IconResolver
            icon_pixbuf = icon_resolver.get_icon_pixbuf(app_id, icon_size_main)
        
        if not icon_pixbuf:
            # Additional fallbacks for common apps
            icon_pixbuf = icon_resolver.get_icon_pixbuf("application-x-executable-symbolic", icon_size_main)
            if not icon_pixbuf:
                icon_pixbuf = icon_resolver.get_icon_pixbuf("image-missing", icon_size_main)
                
        # Ensure icon is scaled to the correct size
        if icon_pixbuf and (icon_pixbuf.get_width() != icon_size_main or icon_pixbuf.get_height() != icon_size_main):
            icon_pixbuf = icon_pixbuf.scale_simple(
                icon_size_main, 
                icon_size_main, 
                gi.repository.GdkPixbuf.InterpType.BILINEAR
            )

        super().__init__(
            name="overview-client-box",
            image=Image(pixbuf=icon_pixbuf),
            tooltip_text=title,
            size=size,
            on_clicked=self.on_button_click,
            on_button_press_event=lambda _, event: connection.send_command(
                f"/dispatch closewindow address:{address}"
            )
            if event.button == 3
            else None,
            on_drag_data_get=lambda _s, _c, data, *_: data.set_text(
                address, len(address)
            ),
            on_drag_begin=lambda _, context: Gtk.drag_set_icon_surface(
                context, createSurfaceFromWidget(self)
            ),
        )

        # Store the desktop_app for later use
        self.desktop_app = desktop_app

        self.drag_source_set(
            start_button_mask=Gdk.ModifierType.BUTTON1_MASK,
            targets=TARGET,
            actions=Gdk.DragAction.COPY,
        )

        self.connect("key_press_event", self.on_key_press_event)

    def on_key_press_event(self, widget, event):
        if event.get_state() & Gdk.ModifierType.SHIFT_MASK:
            if event.keyval in (Gdk.KEY_Return, Gdk.KEY_KP_Enter, Gdk.KEY_space):
                connection.send_command(f"/dispatch closewindow address:{self.address}")
                return True
        return False

    def update_image(self, image):
        # Compute overlay icon size dynamically.
        icon_size_overlay = int(min(self.size) * 0.5)  # adjust factor as needed
        
        # Enhanced icon resolution for overlay
        icon_pixbuf = None
        if hasattr(self, 'desktop_app') and self.desktop_app:
            icon_pixbuf = self.desktop_app.get_icon_pixbuf(size=icon_size_overlay)
            
        if not icon_pixbuf:
            icon_pixbuf = icon_resolver.get_icon_pixbuf(self.app_id, icon_size_overlay)
            
        if not icon_pixbuf:
            icon_pixbuf = icon_resolver.get_icon_pixbuf("application-x-executable-symbolic", icon_size_overlay)
            if not icon_pixbuf:
                icon_pixbuf = icon_resolver.get_icon_pixbuf("image-missing", icon_size_overlay)
        
        # Ensure icon is scaled to the correct size
        if icon_pixbuf and (icon_pixbuf.get_width() != icon_size_overlay or icon_pixbuf.get_height() != icon_size_overlay):
            icon_pixbuf = icon_pixbuf.scale_simple(
                icon_size_overlay, 
                icon_size_overlay, 
                gi.repository.GdkPixbuf.InterpType.BILINEAR
            )
                
        self.set_image(
            Overlay(
                child=image,
                overlays=Image(
                    name="overview-icon",
                    pixbuf=icon_pixbuf,
                    h_align="center",
                    v_align="end",
                    tooltip_text=self.title,
                ),
            )
        )

    def on_button_click(self, *_):
        connection.send_command(f"/dispatch focuswindow address:{self.address}")


class WorkspaceEventBox(EventBox):
    def __init__(self, workspace_id: int, fixed: Gtk.Fixed | None = None, monitor_width: int = None, monitor_height: int = None, monitor_scale: float = 1.0):
        self.fixed = fixed
        
        # Use provided monitor dimensions or fallback to current screen
        width = monitor_width or CURRENT_WIDTH
        height = monitor_height or CURRENT_HEIGHT
        
        # Workspace containers should maintain consistent size across monitors
        # Only use BASE_SCALE, don't multiply by monitor_scale for the container
        container_scale = BASE_SCALE
        
        super().__init__(
            name="overview-workspace-bg",
            h_expand=True,
            v_expand=True,
            size=(int(width * container_scale), int(height * container_scale)),
            child=fixed
            if fixed
            else Label(
                name="overview-add-label",
                h_expand=True,
                v_expand=True,
                markup=icons.circle_plus,
            ),
            on_drag_data_received=lambda _w, _c, _x, _y, data, *_: connection.send_command(
                f"/dispatch movetoworkspacesilent {workspace_id},address:{data.get_data().decode()}"
            ),
        )
        self.drag_dest_set(
            Gtk.DestDefaults.ALL,
            TARGET,
            Gdk.DragAction.COPY,
        )
        if fixed:
            fixed.show_all()



class Overview(Box):
    def __init__(self, monitor_id: int = 0, **kwargs):
        self.monitor_id = monitor_id
        self.monitor_manager = None
        self.workspace_start = 1
        self.workspace_end = 10
        
        # Get monitor manager and workspace range
        try:
            from utils.monitor_manager import get_monitor_manager
            self.monitor_manager = get_monitor_manager()
            self.workspace_start, self.workspace_end = self.monitor_manager.get_workspace_range_for_monitor(monitor_id)
        except ImportError:
            # Fallback if monitor manager not available
            pass
        
        # Get monitor dimensions
        monitor_width = CURRENT_WIDTH
        monitor_height = CURRENT_HEIGHT
        
        if self.monitor_manager:
            monitor_info = self.monitor_manager.get_monitor_by_id(monitor_id)
            if monitor_info:
                monitor_width = monitor_info['width']
                monitor_height = monitor_info['height']
        # Initialize as a Box instead of a PopupWindow.
        super().__init__(name="overview", orientation="v", spacing=8, **kwargs)
        self.workspace_boxes: dict[int, Box] = {}
        self.clients: dict[str, HyprlandWindowButton] = {}
        
        # Initialize app registry for better icon resolution
        self._all_apps = get_desktop_applications()
        self.app_identifiers = self._build_app_identifiers_map()
        
        # Remove the window_class_aliases dictionary completely

        connection.connect("event::openwindow", self.do_update)
        connection.connect("event::closewindow", self.do_update)
        connection.connect("event::movewindow", self.do_update)
        self.update()
        
    def _normalize_window_class(self, class_name):
        """Normalize window class by removing common suffixes and lowercase."""
        if not class_name:
            return ""
            
        normalized = class_name.lower()
        
        # Remove common suffixes
        suffixes = [".bin", ".exe", ".so", "-bin", "-gtk"]
        for suffix in suffixes:
            if normalized.endswith(suffix):
                normalized = normalized[:-len(suffix)]
                
        return normalized
    
    def _classes_match(self, class1, class2):
        """Check if two window class names match with stricter comparison."""
        if not class1 or not class2:
            return False
            
        # Normalize both classes
        norm1 = self._normalize_window_class(class1)
        norm2 = self._normalize_window_class(class2)
        
        # Direct match after normalization
        if norm1 == norm2:
            return True
            
        # Don't do substring matching as it's too error-prone
        # This avoids incorrectly matching flatpak apps and others
        return False
        
    def _build_app_identifiers_map(self):
        """Build a mapping of app identifiers (class names, executables, names) to DesktopApp objects"""
        identifiers = {}
        for app in self._all_apps:
            # Map by name (lowercase)
            if app.name:
                identifiers[app.name.lower()] = app
                
            # Map by display name
            if app.display_name:
                identifiers[app.display_name.lower()] = app
                
            # Map by window class if available
            if app.window_class:
                identifiers[app.window_class.lower()] = app
                
            # Map by executable name if available
            if app.executable:
                exe_basename = app.executable.split('/')[-1].lower()
                identifiers[exe_basename] = app
                
            # Map by command line if available (without parameters)
            if app.command_line:
                cmd_base = app.command_line.split()[0].split('/')[-1].lower()
                identifiers[cmd_base] = app
                
        return identifiers
        
    def find_app(self, app_identifier):
        """Return the DesktopApp object by matching any app identifier."""
        if not app_identifier:
            return None
        
        # Try direct lookup in our identifiers map
        normalized_id = str(app_identifier).lower()
        if normalized_id in self.app_identifiers:
            return self.app_identifiers[normalized_id]
            
        # Try with normalized class name
        norm_id = self._normalize_window_class(normalized_id)
        if norm_id in self.app_identifiers:
            return self.app_identifiers[norm_id]
            
        # More targeted matching with exact names only
        for app in self._all_apps:
            if app.name and app.name.lower() == normalized_id:
                return app
            if app.window_class and app.window_class.lower() == normalized_id:
                return app
            if app.display_name and app.display_name.lower() == normalized_id:
                return app
            # Try with executable basename
            if app.executable:
                exe_base = app.executable.split('/')[-1].lower()
                if exe_base == normalized_id:
                    return app
            # Try with command basename
            if app.command_line:
                cmd_base = app.command_line.split()[0].split('/')[-1].lower()
                if cmd_base == normalized_id:
                    return app
                
        return None

    def update(self, signal_update=False):
        self._all_apps = get_desktop_applications()
        self.app_identifiers = self._build_app_identifiers_map()
        for client in self.clients.values():
            client.destroy()
        self.clients.clear()
        for workspace in self.workspace_boxes.values():
            workspace.destroy()
        self.workspace_boxes.clear()

        if data.PANEL_THEME == "Panel" and data.BAR_POSITION in ["Left", "Right"]:
            rows = 5
            cols = 2
        else:
            rows = 2
            cols = 5

        self.children = [Box(spacing=8) for _ in range(rows)]

        # Get monitor dimensions and scale for scaling
        monitor_width = CURRENT_WIDTH
        monitor_height = CURRENT_HEIGHT
        monitor_scale = 1.0
        
        if self.monitor_manager:
            monitor_info = self.monitor_manager.get_monitor_by_id(self.monitor_id)
            if monitor_info:
                monitor_width = monitor_info['width']
                monitor_height = monitor_info['height']
                monitor_scale = monitor_info.get('scale', 1.0)
        
        # Calculate effective scale for this monitor
        # Higher scale monitors need larger overview elements to appear the same physical size
        effective_scale = BASE_SCALE * monitor_scale

        monitors = {
            monitor["id"]: (monitor["x"], monitor["y"], monitor["transform"])
            for monitor in json.loads(connection.send_command("j/monitors").reply.decode())
        }
        
        # Filter clients to only show those in this monitor's workspace range
        for client in json.loads(connection.send_command("j/clients").reply.decode()):
            workspace_id = client["workspace"]["id"]
            if workspace_id > 0 and self.workspace_start <= workspace_id <= self.workspace_end:
                btn = HyprlandWindowButton(
                    window=self,
                    title=client["title"],
                    address=client["address"],
                    app_id=client["initialClass"],
                    size=(client["size"][0] * effective_scale, client["size"][1] * effective_scale),
                    transform=monitors[client["monitor"]][2],
                )
                self.clients[client["address"]] = btn
                w_id = workspace_id
                if w_id not in self.workspace_boxes:
                    self.workspace_boxes[w_id] = Gtk.Fixed.new()
                self.workspace_boxes[w_id].put(
                    btn,
                    abs(client["at"][0] - monitors[client["monitor"]][0]) * effective_scale,
                    abs(client["at"][1] - monitors[client["monitor"]][1]) * effective_scale,
                )

        # Generate workspaces only for this monitor's range
        for w_id in range(self.workspace_start, self.workspace_end + 1):
            idx = w_id - self.workspace_start
            if rows == 2:
                row = 0 if idx < cols else 1
            else:
                row = idx // cols
            overview_row = self.children[row]
            overview_row.add(
                Box(
                    name="overview-workspace-box",
                    orientation="vertical",
                    children=[
                        Label(name="overview-workspace-label", label=f"Workspace {w_id}"),
                        WorkspaceEventBox(
                            w_id,
                            self.workspace_boxes.get(w_id),
                            monitor_width=monitor_width,
                            monitor_height=monitor_height,
                            monitor_scale=monitor_scale
                        ),
                    ],
                )
            )

    def do_update(self, *_):
        logger.info(f"[Overview] Updating for :{_[1].name}")
        self.update(signal_update=True)
