#!/usr/bin/env python3

import socket
import json
import subprocess
import threading
from sys import exit
from time import sleep
from os import environ
from pathlib import Path

NORMAL_WALLPAPER_DIR = Path("/home/kolkas/.local/share/wallpaper/current")
BLURRED_WALLPAPER_DIR = Path("/home/kolkas/.local/share/wallpaper/blurred")


def _get_first_file(dir: Path, pattern: str = "*") -> Path | None:
    return next(dir.glob(pattern), None)


def get_niri_socket():
    return environ['NIRI_SOCKET']


def _log(msg: str):
    # print(msg)

    # logFIle = Path("/tmp/niri-autoblur.log")
    # try:
    #     with logFIle.open("a") as f:
    #         f.write(msg + "\n")
    # except Exception:
    #     pass

    pass


def swwwLoadImg(namespace: str, wallpaper: Path):
    cmd = [
        "swww",
        "img",
        "-n",
        namespace,
        str(wallpaper),
        "--transition-type",
        "fade",
        "--transition-duration",
        "0.5",
    ]
    _log(" ".join(cmd))
    ret = 0
    try:
        ret = subprocess.run(cmd, check=True).returncode
    except Exception as e:
        _log(f"[SWWW] failed to set wallpaper: {e}")
        return False

    if ret != 0:
        _log(f"[SWWW] failed to set wallpaper, exit code: {ret}")
        return False

    return True


def swwwStartDaemon(namespace: str):
    # Check if daemon is already running
    cmd = ["pgrep", "-f", f"swww daemon -n {namespace}"]
    try:
        output = subprocess.check_output(cmd, text=True)
        pids = output.strip().splitlines()
        if pids:
            _log(f"[SWWW] daemon already running with PIDs: {', '.join(pids)}")
            return True
    except subprocess.CalledProcessError:
        # pgrep returns non-zero exit code if no process is found
        pass
    except Exception as e:
        _log(f"[SWWW] failed to check if daemon is running: {e}")
        pass

    try:
        subprocess.Popen(["swww-daemon", "-n", namespace], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        _log(f"[SWWW] daemon started for namespace: {namespace}")
        return True
    except Exception as e:
        _log(f"[SWWW] failed to start daemon non-blockingly: {e}")
        return False


class AutoBlur:
    _interval: float
    _normalDir: Path
    _blurredDir: Path
    _isBlurred = threading.Event()
    _thread: threading.Thread | None = None
    _lastWallpaer: Path | None = None
    _isFirst = True

    def __init__(self, normalDir, blurredDir, interval=0.2):
        self._interval = interval
        self._normalDir = normalDir
        self._blurredDir = blurredDir

        # Niri will send "WindowsChanged" event on connect, so no need to init here
        # init state
        # self.setBlurred(AutoBlur.initIsBlurred())

    @staticmethod
    def initIsBlurred() -> bool:
        '''[ $(niri msg focused-window | wc -l) -gt 1 ]'''
        cmd = ["niri", "msg", "focused-window"]
        try:
            output = subprocess.check_output(cmd, text=True)
            lines = output.strip().splitlines()
            return len(lines) > 1
        except Exception as e:
            _log(f"[initIsBlurred] failed to check focused window, assuming none: {e}")
            return False

    def setBlurred(self, isBlurred: bool) -> None:
        # Cache state, avoid starting thread unnecessarily
        # if not self._isFirst and self._isBlurred.is_set() == isBlurred:
        #     _log("[AutoBlur] state unchanged")
        #     return
        self._isFirst = False

        if isBlurred:
            self._isBlurred.set()
            _log("[AutoBlur] set to blurred")
        else:
            self._isBlurred.clear()
            _log("[AutoBlur] set to normal")

        if self._thread is None or not self._thread.is_alive():
            self._thread = threading.Thread(target=self._run, daemon=True)
            self._thread.start()

    def _run(self) -> None:
        '''Wait until wallpapers are ready & apply the correct one according to the current state'''
        while True:
            if self._isBlurred.is_set():
                wallpaper = _get_first_file(self._blurredDir)
            else:
                wallpaper = _get_first_file(self._normalDir)

            if wallpaper is not None and wallpaper.exists():
                if self._apply(wallpaper):
                    break

            sleep(self._interval)

    def _apply(self, wallpaper: Path) -> bool:
        if wallpaper == self._lastWallpaer:
            return True

        if not swwwLoadImg("background", wallpaper):
            return False

        self._lastWallpaer = wallpaper
        return True


autoBlurInst = AutoBlur(NORMAL_WALLPAPER_DIR, BLURRED_WALLPAPER_DIR)


def handleEvent(event_name, payload):
    if event_name == "WindowFocusChanged":
        _log(f"[EventHandler] WindowFocusChanged event received")
        id = payload.get("id", "")
        if isinstance(id, int):
            _log(f"[EventHandler] focused window id: {id}")
            autoBlurInst.setBlurred(True)
        elif isinstance(id, str) and id == "None":
            _log("[EventHandler] no focused window")
            autoBlurInst.setBlurred(False)
        else:
            _log(f"[EventHandler] unknown id: {id}, assuming no focused window")
            autoBlurInst.setBlurred(False)
    elif event_name == "WindowsChanged":
        _log(f"[EventHandler] WindowsChanged event received")
        windows = payload.get("windows", [])
        for window in windows:
            if window.get("is_focused", False):
                _log(f"[EventHandler] found focused window")
                autoBlurInst.setBlurred(True)
                return
        _log("[EventHandler] no focused window found")
        autoBlurInst.setBlurred(False)
    elif event_name == "WindowOpenedOrChanged":
        _log(f"[EventHandler] WindowOpenedOrChanged event received")
        window = payload.get("window", {})
        if window.get("is_focused", False):
            _log(f"[EventHandler] opened/changed focused window")
            autoBlurInst.setBlurred(True)
    else:
        _log(f"[EventHandler] unhandled event: {event_name}")


def print_event(eventName, payload):
    _log(f"[EventHandler] event: {eventName}, payload:\n{json.dumps(payload, indent=2, ensure_ascii=False)}")


def connect_niri(niriSocket: str, handler) -> bool:
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        sock.connect(niriSocket)
    except Exception as e:
        sock.close()
        _log(f"[Socket] failed to connect to {niriSocket}: {e}")
        return False

    f = sock.makefile("rwb")

    try:
        f.write(b'"EventStream"\n')
        f.flush()

        first = f.readline()
        if not first:
            # raise RuntimeError("connection closed by server before handshake")
            _log("[Socket] connection closed by server before handshake")
            return False
        try:
            resp = json.loads(first.decode())
        except Exception:
            resp = first.decode().strip()
        _log(f"[Socket] handshake response: {resp}")

        while True:
            line = f.readline()
            if not line:
                _log("[Socket] socket closed by server")
                break
            s = line.decode().strip()
            if s == "":
                continue
            try:
                obj = json.loads(s)
            except Exception as e:
                _log(f"[Socket] failed to parse line as JSON: {s}, error: {e}")
                continue

            keys = list(obj.keys())
            if keys:
                event_name = keys[0]
                payload = obj[event_name]
            else:
                event_name = "<unknown>"
                payload = obj
            handler(event_name, payload)
    finally:
        try:
            f.close()
        except Exception:
            pass
        try:
            sock.close()
        except Exception:
            return False
        return True


if __name__ == "__main__":
    desktop = environ.get("XDG_CURRENT_DESKTOP", "")
    if desktop == "niri":
        _log("[Main] running in Niri")
        _log("[Main] starting swww daemons")
        if not swwwStartDaemon("background"):
            exit(1)
        if not swwwStartDaemon("backdrop"):
            exit(1)
        sleep(1)  # give some time to start

        _log("[Main] loading initial wallpapers")
        # Init wallpaper for backdrop
        blurred = _get_first_file(BLURRED_WALLPAPER_DIR)
        if blurred:
            swwwLoadImg("backdrop", blurred)
        # Init wallpaper for background
        normal = _get_first_file(NORMAL_WALLPAPER_DIR)
        if normal:
            swwwLoadImg("background", normal)

        # Connect to Niri socket
        _log(f"[Main] connecting to Niri socket")
        niri_socket = get_niri_socket()
        if not niri_socket:
            _log("[Main] NIRI_SOCKET environment variable is not set.")
            exit(1)

        if not connect_niri(niri_socket, handleEvent):
            exit(1)
    elif desktop == "Hyprland":
        _log("[Main] running in Hyprland")
        _log("[Main] starting swww daemon")
        if not swwwStartDaemon("background"):
            exit(1)
        sleep(1)  # similarly

        _log("[Main] loading initial wallpaper")
        normal = _get_first_file(NORMAL_WALLPAPER_DIR)
        if normal:
            swwwLoadImg("background", normal)

        # Wait indefinitely
        while True:
            sleep(3600)