#!/usr/bin/env python3
import argparse
import datetime
import json
import os
import pathlib
import random as _random
import re
import shutil
import subprocess
import urllib.parse
import urllib.request
import uuid

WALLPAPERS_DIR = str(
    pathlib.Path(
        os.environ.get("WALLPAPERS_DIR")
        or pathlib.Path.home() / "Pictures" / "Wallpapers"
    ).expanduser()
)
SYMLINK_PATH = str(pathlib.Path("~/Pictures/wallpaper.png").expanduser())

# old_print = print


# def new_print(*args, **kwargs):
#     old_print(*args, **kwargs)
#     old_print("\r", end="")


# print = new_print


def check_setup() -> None:
    # Check required programs are available
    for prog in ["magick", "curl"]:
        try:
            subprocess.run([prog, "--version"], capture_output=True)
        except FileNotFoundError:
            print(f"Warning: Required program '{prog}' not found in PATH.")
            exit(1)

    for prog in ["zenity", "yad"]:
        try:
            subprocess.run([prog, "--version"], capture_output=True)
            break
        except FileNotFoundError:
            continue
    else:
        print("Warning: No GUI file chooser (zenity or yad) found in PATH.")
        exit(1)

    for prog in ["swww", "hyprpanel", "wal"]:
        try:
            subprocess.run([prog, "--version"], capture_output=True)
        except FileNotFoundError:
            print(f"Note: Optional program '{prog}' not found in PATH.")

    wallpaper_dir = pathlib.Path(WALLPAPERS_DIR)
    if not wallpaper_dir.is_dir():
        print(f"Creating wallpapers directory: {WALLPAPERS_DIR}")
        wallpaper_dir.mkdir(parents=True, exist_ok=True)

    return


def get_path(file: bool = False, multiple: bool = False) -> list[str]:
    # Try common Linux GUI choosers (zenity, yad).
    # Returns a list of selected paths (absolute).

    def _normalize_output(raw: str) -> list[str]:
        if not raw:
            return []
        # zenity/yad often use a separator, kdialog/newlines, some use '|' or NUL
        for sep in ("\x00", "::", "|"):
            raw = raw.replace(sep, "\n")
        parts = [p.strip() for p in raw.splitlines() if p.strip()]
        return [os.path.abspath(os.path.expanduser(p)) for p in parts]

    chooser_commands = []

    if file:
        if multiple:
            chooser_commands = [
                ["zenity", "--file-selection", "--multiple", "--separator=::"],
                ["yad", "--file", "--multiple", "--separator=::"],
            ]
        else:
            chooser_commands = [
                ["zenity", "--file-selection"],
                ["yad", "--file"],
            ]
    else:
        if multiple:
            chooser_commands = [
                [
                    "zenity",
                    "--file-selection",
                    "--directory",
                    "--multiple",
                    "--separator=::",
                ],
                ["yad", "--file", "--directory", "--multiple", "--separator=::"],
            ]
        else:
            chooser_commands = [
                ["zenity", "--file-selection", "--directory"],
                ["yad", "--file", "--directory"],
            ]

    for cmd in chooser_commands:
        try:
            p = subprocess.run(cmd, capture_output=True, text=True)
        except FileNotFoundError:
            continue

        out = p.stdout.strip()
        if p.returncode == 0 and out:
            return _normalize_output(out) if multiple else _normalize_output(out)[:1]

    return []


def add_wallpaper(path: list[str], remove_original: bool = False) -> None:
    for pic in path:
        pic_path = pathlib.Path(pic)
        print(f"Adding wallpaper: {pic}")
        filename = f"wp{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}_{pic_path.name[:10]}.png"
        # ensure target directory exists
        target_dir = pathlib.Path(WALLPAPERS_DIR)
        target_dir.mkdir(parents=True, exist_ok=True)

        dest = target_dir / filename

        try:
            # convert/source -> PNG using ImageMagick
            # TODO format to res?
            subprocess.run(["magick", str(pic_path), str(dest)], check=True)
            print(f"Saved converted wallpaper to: {dest}")
        except FileNotFoundError:
            print(
                "ImageMagick 'magick' not found. Please install ImageMagick or ensure it's in PATH."
            )
        except subprocess.CalledProcessError as e:
            print(f"ImageMagick conversion failed: {e}")

        if remove_original:
            print(f"Removing original file: {pic}")
            try:
                pic_path.unlink()
            except Exception as e:
                print(f"Failed to remove original file: {e}")

        return


def fetch_wallpaper(urls_or_ids: list[str] | None = None) -> None:

    TEMP_DIR = "/tmp/wallhaven_downloads"
    TARGET_RESOLUTION = os.environ.get("WALLPAPER_TARGET_RESOLUTION", "2560x1600")

    pathlib.Path(TEMP_DIR).mkdir(parents=True, exist_ok=True)

    def _get_api_key() -> str:
        key = os.environ.get("WALLHAVEN_API_KEY") or os.environ.get("WALLPAPER_API_KEY")
        if key:
            return key.strip()
        try:
            p = subprocess.run(
                ["pass", "api/wallhaven"], capture_output=True, text=True
            )
            if p.returncode == 0:
                return p.stdout.strip()
        except FileNotFoundError:
            pass

        return ""

    api_key = _get_api_key()

    items = urls_or_ids or []
    if not items:
        try:
            p = subprocess.run(
                ["wl-paste", "--no-newline"], capture_output=True, text=True
            )
            clip = p.stdout.strip() if p.returncode == 0 else ""
        except FileNotFoundError:
            clip = ""
        if not clip:
            print("No URL/id in clipboard")
            return
        items = [clip]

    for item in items:
        # Extract ID (last path segment, strip trailing slash and query)
        parsed = urllib.parse.urlparse(item)
        path = parsed.path or item
        wallpaper_id = re.sub(r".*/", "", path)
        wallpaper_id = re.split(r"[/?]", wallpaper_id)[0].strip()

        if not wallpaper_id:
            print(f"Could not extract id from URL: {item}")
            continue

        if not re.match(r"^[A-Za-z0-9]{4,8}$", wallpaper_id):
            print(f"Invalid wallhaven ID format: {wallpaper_id}")
            continue

        api_url = f"https://wallhaven.cc/api/v1/w/{wallpaper_id}"
        if api_key:
            api_url += f"?apikey={api_key}"

        print(f"Fetching wallpaper info for ID: {wallpaper_id}")
        try:
            p = subprocess.run(
                ["curl", "-s", "-f", api_url], capture_output=True, text=True
            )
            if p.returncode != 0:
                print(
                    f"Failed to fetch data from wallhaven API (curl exit: {p.returncode})"
                )
                continue
            response = p.stdout
        except FileNotFoundError:
            print("curl not found; cannot fetch wallpaper metadata")
            return

        try:
            resp_json = json.loads(response)
        except json.JSONDecodeError:
            print("Failed to parse API response as JSON")
            continue

        if resp_json.get("error"):
            print(f"API Error: {resp_json.get('error')}")
            continue

        data = resp_json.get("data")
        if not data:
            print("No wallpaper data found in API response")
            continue

        image_url = data.get("path")
        resolution = data.get("resolution")
        file_type = data.get("file_type")
        file_size = data.get("file_size")
        purity = data.get("purity")

        if not image_url:
            print("Could not extract image URL from API response")
            continue

        print("Wallpaper details:")
        print(f"  Resolution: {resolution}")
        print(f"  File type: {file_type}")
        print(f"  File size: {((file_size or 0) // 1024)} KB")
        print(f"  Purity: {purity}")
        print(f"  URL: {image_url}")

        if resolution and resolution != TARGET_RESOLUTION:
            print(
                f"Warning: Wallpaper resolution is {resolution} (target: {TARGET_RESOLUTION})"
            )

        filename = f"{wallpaper_id}"
        full_path = pathlib.Path(TEMP_DIR) / filename

        try:
            with (
                urllib.request.urlopen(image_url) as r,
                open(full_path, "wb") as out_f,
            ):
                shutil.copyfileobj(r, out_f)
        except Exception as e:
            print(f"Error: Failed to download the wallpaper from: {image_url} ({e})")
            continue

        if not full_path.is_file():
            print("Error: Downloaded file is missing")
            continue

        try:
            size = full_path.stat().st_size
        except OSError as e:
            print(f"Error: Could not stat downloaded file: {e}")
            continue

        if size == 0:
            print("Error: Downloaded file is empty")
            continue

        print(f"Download successful: {full_path} ({size // 1024} KB)")

        add_wallpaper([str(full_path)], remove_original=True)

        try:
            subprocess.run(
                ["zenity", "--notification", "--text", "Wallpaper downloaded!"]
            )
        except FileNotFoundError:
            pass

    return


def next_wallpaper(
    path: str | None = None, random: bool = False, interactive: bool = False
) -> None:

    symlink = pathlib.Path(SYMLINK_PATH)
    wallpapers_dir = pathlib.Path(WALLPAPERS_DIR)  # use module-level constant

    def _atomic_symlink(target: str, link_path: pathlib.Path) -> None:
        tmp_name = link_path.parent / f".{link_path.name}.tmp.{uuid.uuid4().hex}"
        try:
            os.symlink(target, tmp_name)
            os.replace(tmp_name, link_path)
        finally:
            try:
                if tmp_name.exists() or tmp_name.is_symlink():
                    tmp_name.unlink()
            except Exception:
                pass

    def _apply_wallpaper(target: str) -> None:
        print(f"Setting wallpaper to: {target}")
        _atomic_symlink(target, symlink)

        # trigger external tools; ignore missing commands
        try:
            subprocess.run(
                ["swww", "img", target, "--transition-type", "fade"], check=False
            )
        except FileNotFoundError:
            pass
        try:
            subprocess.run(["hyprpanel", "-q"], check=False)
            # restart in background
            subprocess.Popen(
                ["hyprpanel"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
            )
        except FileNotFoundError:
            pass
        try:
            subprocess.run(["wal", "-c"], check=False)
            subprocess.run(["wal", "-n", "-i", target], check=False)
            pass
        except FileNotFoundError:
            pass

    # If interactive requested, ask user for a file
    if not path and interactive:
        selected = get_path(file=True, multiple=False)
        if not selected:
            print("No file selected interactively.")
            return
        path = selected[0]

    # If a file path was provided (force), validate and apply immediately
    if path:
        p = pathlib.Path(path).expanduser()
        if not p.is_file():
            print(f"Specified wallpaper not found: {p}")
            return

        target = str(p.resolve())
        _apply_wallpaper(target)
        return

    # No explicit path: choose from wallpapers dir
    if not wallpapers_dir.is_dir():
        print(f"Wallpapers directory not found: {wallpapers_dir}")
        return

    # Resolve current wallpaper target (if any)
    try:
        current = symlink.resolve(strict=False)
        current_str = str(current) if current.exists() else ""
    except Exception:
        current_str = ""

    def _collect_png_files(wallpapers_dir: pathlib.Path, files: list[pathlib.Path]):
        # If directory contains a ".off" file → skip entire directory
        # for item in wallpapers_dir.iterdir():
        # print(f"Checking item: {item}")
        if any(item.name == ".off" for item in wallpapers_dir.iterdir()):
            print(f"Skipping folder (contains .off): {wallpapers_dir}")
            return

        for item in wallpapers_dir.iterdir():
            if item.is_file():
                if item.suffix.lower() == ".png":
                    # print(f"Found wallpaper file: {item}")
                    files.append(item)

            elif item.is_dir():
                _collect_png_files(item, files)

    files = []
    _collect_png_files(wallpapers_dir, files)
    files.sort(key=lambda p: p.name.lower())
    files = [str(f) for f in files]

    if not files:
        print(f"No png wallpapers found in {wallpapers_dir}")
        return

    def choose_random_excluding_current() -> str:
        candidates = [
            f for f in files if os.path.realpath(f) != os.path.realpath(current_str)
        ]
        if not candidates:
            print("No other wallpaper to choose (only current exists).")
            raise SystemExit(1)
        return _random.choice(candidates)

    def choose_next_in_list() -> str:
        n = len(files)
        if n == 1:
            if os.path.realpath(files[0]) == os.path.realpath(current_str):
                print("No other wallpaper to choose (only current exists).")
                raise SystemExit(1)
            return files[0]
        # find current index
        idx = -1
        for i, f in enumerate(files):
            try:
                if os.path.realpath(f) == os.path.realpath(current_str):
                    idx = i
                    break
            except Exception:
                continue
        if idx == -1:
            return files[0]
        next_idx = (idx + 1) % n
        return files[next_idx]

    try:
        new = choose_random_excluding_current() if random else choose_next_in_list()
    except SystemExit:
        return

    new = str(pathlib.Path(new).resolve())
    print(f"Chosen wallpaper: {new}")

    _apply_wallpaper(new)

    return


def switch_folders(folders: list[str], on: bool = False, off: bool = False) -> None:
    if on and off:
        print("Cannot specify both --on and --off.")
        return

    for folder_name in folders:
        folder_path = pathlib.Path(WALLPAPERS_DIR) / folder_name
        if not folder_path.is_dir():
            print(f"Folder not found: {folder_path}")
            continue

        off_file = folder_path / ".off"
        if on:
            if off_file.is_file():
                off_file.unlink()
                print(f"Turned ON folder: {folder_name}")
            else:
                print(f"Folder already ON: {folder_name}")
        elif off:
            if not off_file.is_file():
                off_file.touch()
                print(f"Turned OFF folder: {folder_name}")
            else:
                print(f"Folder already OFF: {folder_name}")
        else:
            # Toggle state
            if off_file.is_file():
                off_file.unlink()
                print(f"Turned ON folder: {folder_name}")
            else:
                off_file.touch()
                print(f"Turned OFF folder: {folder_name}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="Wallpaper Setter",
        description="This script do basic actions with wallpaper for Hyprland",
        epilog="Designed by Kr0sH_512",
    )

    coparser = parser.add_subparsers(
        required=True, dest="command", title="commands", metavar="COMMAND"
    )
    # --- Add pic to wallpaper folder or run interactive mode ---
    add_coparser = coparser.add_parser(
        "add", aliases=["a", "move", "m"], help="Add wallpaper"
    )
    add_coparser.add_argument(
        "-r",
        "--remove",
        action="store_true",
        default=False,
        help="Remove original file(s) after adding",
    )
    add_coparser.add_argument(
        "-f",
        "--file",
        type=str,
        help="Path to wallpaper to add (interactive if not specified)",
        required=False,
        nargs="*",
    )
    # --- Download wallpaper from Wallhaven by URL or ID ---
    down_coparser = coparser.add_parser(
        "download", aliases=["d", "fetch", "f"], help="Download wallpaper"
    )
    down_coparser.add_argument(
        "-u",
        "--urls_or_ids",
        type=str,
        help="Wallhaven wallpaper URL(s) or ID(s) to download (clipboard if not specified)",
        required=False,
        nargs="*",
    )
    # --- Select wallpaper to set (interactive, by file or by order) ---
    next_coparser = coparser.add_parser("next", aliases=["n"], help="Next wallpaper")
    next_coparser.add_argument(
        "-r",
        "--random",
        action="store_true",
        default=False,
        help="Set a random wallpaper (if no file is specified)",
    )
    # TODO: if file is path to directory, pick from there
    next_coparser.add_argument(
        "-f",
        "--file",
        type=str,
        help="Set wallpaper by file path",
        required=False,
    )
    next_coparser.add_argument(
        "-i",
        "--interactive",
        action="store_true",
        default=False,
        help="Set wallpaper interactively",
    )
    # --- List wallpapers in directory (and info) ---
    list_coparser = coparser.add_parser(
        "list", aliases=["l"], help="List and count wallpapers"
    )
    list_coparser.add_argument(
        "-f",
        "--file",
        type=str,
        help="Path to list wallpapers from (defaults to wallpapers directory)",
        required=False,
    )
    list_coparser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        default=False,
        help="Show detailed information about wallpapers",
    )
    # --- Switch on/off wallpaper folders ---
    switch_coparser = coparser.add_parser(
        "switch", aliases=["sw"], help="Turn on/off folders in $WALLPAPERS_DIR"
    )
    switch_coparser.add_argument(
        "-0",
        "--off",
        action="store_true",
        default=False,
        help="Turn off specified folder(s) (switch state if not specified)",
    )
    switch_coparser.add_argument(
        "-1",
        "--on",
        action="store_true",
        default=False,
        help="Turn on specified folder(s) (switch state if not specified)",
    )
    switch_coparser.add_argument(
        "folders",
        help="Folder(s) name to switch (interactive if not specified)",
        nargs="+",
    )
    # --- Get info about current wallpaper and folders ---
    info_coparser = coparser.add_parser(
        "info", aliases=["i"], help="Get info about current wallpaper and folders"
    )
    info_coparser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        default=False,
        help="Show detailed information",
    )

    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    # print(args)

    check_setup()
    # return

    match args.command:
        case "add" | "a" | "move" | "m":
            path = args.file if args.file else get_path(file=True, multiple=True)
            add_wallpaper(path=path, remove_original=args.remove)
        case "download" | "d" | "fetch" | "f":
            fetch_wallpaper(urls_or_ids=args.urls_or_ids)
        case "next" | "n":
            next_wallpaper(
                path=args.file, random=args.random, interactive=args.interactive
            )
        case "list" | "l":
            pass
        case "switch" | "sw":
            switch_folders(args.folders, on=args.on, off=args.off)
        case "info" | "i":
            pass

    # print(args)


if __name__ == "__main__":
    main()