#!/usr/bin/env python3

import argparse
import logging
import os
from pathlib import Path
import sys

try:
    from PIL import Image
    import numpy as np
    from skimage.color import rgb2lab, lab2rgb
    import pytoml
except ImportError:
    print("Dependencies not found. Please install them with 'pip install Pillow scikit-image numpy pytoml'", file=sys.stderr)
    sys.exit(1)

# --- Configuration & Setup ---

APP_NAME = "paintr"
LOG_LEVELS = ["FATAL", "ERROR", "WARN", "INFO", "DEBUG"]

# Default theme data embedded in the script for first-run setup.
DEFAULT_THEME_NAME = "catppuccin"
DEFAULT_THEME_CONTENT = """[data]
name = "Catppuccin"
version = "1.0"
variant = "dark"

[colors]
background = "#1e1e2e"
foreground = "#cdd6f4"
primary = "#cba6f7"
secondary = "#89b4fa"
accent = "#f38ba8"
muted = "#6c7086"
"""

# --- Logging Setup ---

def setup_logging(level):
    """Sets up the logging configuration based on the provided level."""
    log_level = getattr(logging, level.upper(), logging.INFO)
    logging.basicConfig(
        level=log_level,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )

# --- Helper Functions ---

def get_config_dir():
    """Returns the XDG config directory for the application, creating it if necessary."""
    xdg_config_home = os.environ.get("XDG_CONFIG_HOME")
    config_dir = Path(xdg_config_home) / APP_NAME if xdg_config_home else Path.home() / ".config" / APP_NAME

    # Create the directory if it doesn't exist
    if not config_dir.exists():
        try:
            config_dir.mkdir(parents=True)
            logging.info(f"Created configuration directory at {config_dir}")
        except Exception as e:
            logging.error(f"Failed to create configuration directory: {e}")

    return config_dir

def create_default_theme_file(config_dir):
    """Creates a default theme file if one doesn't exist."""
    theme_file_path = config_dir / f"{DEFAULT_THEME_NAME}.toml"
    if not theme_file_path.exists():
        try:
            with open(theme_file_path, "w") as f:
                f.write(DEFAULT_THEME_CONTENT)
            logging.info(f"Created default theme file '{DEFAULT_THEME_NAME}.toml'")
        except Exception as e:
            logging.error(f"Failed to create default theme file: {e}")

def hex_to_rgb(hex_color):
    """Converts a hex color string to an RGB tuple."""
    hex_color = hex_color.lstrip('#')
    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))

def load_palette(theme_name, variant_name):
    """
    Loads a color palette from a TOML file based on its theme name and variant.
    """
    config_dir = get_config_dir()

    # Iterate through all TOML files to find a matching theme and variant
    for theme_file in config_dir.glob("*.toml"):
        try:
            with open(theme_file, 'rb') as f:
                palette_data = pytoml.load(f)

            # Check for name and variant match
            data = palette_data.get('data', {})
            if data.get('name') == theme_name and data.get('variant') == variant_name:
                if "colors" not in palette_data:
                    logging.fatal(f"Invalid theme file: '{theme_file.name}' is missing the [colors] table.")
                    sys.exit(1)

                logging.info(f"Loaded theme '{theme_name}' variant '{variant_name}' from '{theme_file.name}'")

                # Filter and convert colors from the [colors] table
                hex_colors = [v for k, v in palette_data['colors'].items() if isinstance(v, str)]
                return [hex_to_rgb(h) for h in hex_colors]

        except Exception as e:
            logging.debug(f"Could not parse theme file {theme_file.name}: {e}")
            continue

    # If loop completes without finding a match
    logging.fatal(f"Color scheme '{theme_name}:{variant_name}' not found.")
    sys.exit(1)

def list_themes():
    """Lists available theme files, grouping them by name and showing variants."""
    config_dir = get_config_dir()
    if not config_dir.is_dir():
        logging.info(f"Theme directory not found at {config_dir}")
        return

    # Dictionary to store themes, grouped by name
    themes = {}

    for theme_file in sorted(config_dir.glob("*.toml")):
        file_name = theme_file.stem
        try:
            with open(theme_file, 'rb') as f:
                data = pytoml.load(f)

            display_name = data.get('data', {}).get('name', file_name)
            variant = data.get('data', {}).get('variant', 'default')

            # Generate the color blocks
            color_blocks = ""
            colors = data.get('colors', {})
            for hex_color in colors.values():
                if isinstance(hex_color, str) and hex_color.startswith('#'):
                    rgb = hex_to_rgb(hex_color)
                    color_blocks += f"\033[48;2;{rgb[0]};{rgb[1]};{rgb[2]}m \033[0m"

            if display_name not in themes:
                themes[display_name] = []
            themes[display_name].append({
                'variant': variant,
                'file_name': f"{file_name}.toml",
                'color_blocks': color_blocks
            })

        except Exception as e:
            logging.debug(f"Error parsing {theme_file}: {e}")
            if 'unknown' not in themes:
                themes['unknown'] = []
            themes['unknown'].append({
                'variant': 'N/A',
                'file_name': f"{file_name}.toml",
                'color_blocks': ''
            })

    print("Available color schemes:")
    for display_name in sorted(themes.keys()):
        print(f"- {display_name}")
        for variant_info in themes[display_name]:
            print(f"   - {variant_info['color_blocks']} {variant_info['variant']} ({variant_info['file_name']})")

def clamp(value, min_val=0, max_val=255):
    """Clamps a value within a specified range."""
    return max(min_val, min(max_val, value))

def adjust_brightness(r, g, b, brightness):
    """
    Adjusts the brightness of an RGB color.
    Brightness is a value from -100 (darken) to 100 (lighten).
    """
    adjustment = (brightness / 100) * 255
    return (
        clamp(r + adjustment),
        clamp(g + adjustment),
        clamp(b + adjustment)
    )

def rgb_to_hsl(r, g, b):
    r /= 255.0
    g /= 255.0
    b /= 255.0

    max_val = max(r, g, b)
    min_val = min(r, g, b)
    h, s, l = 0, 0, (max_val + min_val) / 2.0

    if max_val != min_val:
        diff = max_val - min_val
        s = diff / (2.0 - max_val - min_val) if l > 0.5 else diff / (max_val + min_val)
        if max_val == r:
            h = (g - b) / diff + (6 if g < b else 0)
        elif max_val == g:
            h = (b - r) / diff + 2
        elif max_val == b:
            h = (r - g) / diff + 4
        h /= 6.0

    return [h * 360, s, l]

def hsl_to_rgb(h, s, l):
    h /= 360.0

    def hue_to_rgb(p, q, t):
        if t < 0: t += 1
        if t > 1: t -= 1
        if t < 1/6: return p + (q - p) * 6 * t
        if t < 1/2: return q
        if t < 2/3: return p + (q - p) * (2/3 - t) * 6
        return p

    if s == 0:
        r = g = b = l
    else:
        q = l * (1 + s) if l < 0.5 else l + s - l * s
        p = 2 * l - q
        r = hue_to_rgb(p, q, h + 1/3)
        g = hue_to_rgb(p, q, h)
        b = hue_to_rgb(p, q, h - 1/3)

    return [int(r * 255), int(g * 255), int(b * 255)]

def adjust_saturation(r, g, b, saturation):
    """Adjusts the saturation of an RGB color."""
    h, s, l = rgb_to_hsl(r, g, b)
    new_s = max(0, min(1, s * saturation))
    return hsl_to_rgb(h, new_s, l)


# --- Core Colorization Logic (Translated from imageColorizer.ts) ---

def colorize_image(image_path, selected_colors, brightness_adjustment=0, saturation_adjustment=1.0):
    """
    Colorizes an image using the provided colors based on their declared order,
    with options for brightness and saturation adjustment.
    """
    logging.info(f"Colorizing image '{image_path}' with {len(selected_colors)} colors...")

    try:
        img = Image.open(image_path).convert('RGB')
        pixels = np.array(img)
    except FileNotFoundError:
        logging.fatal(f"Input image not found: {image_path}")
        sys.exit(1)
    except Exception as e:
        logging.fatal(f"Error opening image: {e}")
        sys.exit(1)

    # Convert image and palette to Lab color space
    image_lab = rgb2lab(pixels / 255.0)
    palette_lab = rgb2lab(np.array(selected_colors) / 255.0)

    new_pixels_lab = np.zeros_like(image_lab)

    for i in range(image_lab.shape[0]):
        for j in range(image_lab.shape[1]):
            pixel_lab = image_lab[i, j]
            original_L = pixel_lab[0]

            # Map the original lightness (0-100) to an index in the palette.
            # This ensures the original order of colors is used.
            palette_index = (original_L / 100) * (len(palette_lab) - 1)
            lower_index = int(palette_index)
            upper_index = min(lower_index + 1, len(palette_lab) - 1)

            # Get the colors to interpolate between
            lower_color = palette_lab[lower_index]
            upper_color = palette_lab[upper_index]

            amount = palette_index - lower_index

            # Interpolate the color in Lab space
            interpolated_L = original_L
            interpolated_a = lower_color[1] + (upper_color[1] - lower_color[1]) * amount
            interpolated_b = lower_color[2] + (upper_color[2] - lower_color[2]) * amount

            new_pixels_lab[i, j] = [interpolated_L, interpolated_a, interpolated_b]

    # Convert back to RGB
    new_pixels_rgb = (lab2rgb(new_pixels_lab) * 255).astype(np.uint8)

    # Apply saturation and brightness adjustments
    for i in range(new_pixels_rgb.shape[0]):
        for j in range(new_pixels_rgb.shape[1]):
            r, g, b = new_pixels_rgb[i, j]

            # Apply saturation adjustment first
            if saturation_adjustment != 1.0:
                r, g, b = adjust_saturation(r, g, b, saturation_adjustment)

            # Apply brightness adjustment second
            if brightness_adjustment != 0:
                r, g, b = adjust_brightness(r, g, b, brightness_adjustment)

            new_pixels_rgb[i, j] = [r, g, b]

    return Image.fromarray(new_pixels_rgb)

# --- Main Execution ---

def main():
    parser = argparse.ArgumentParser(
        description="A command-line utility to colorize wallpapers with color palettes."
    )

    parser.add_argument("-i", "--input", help="Path to the input image file.", required=False)
    parser.add_argument("-t", "--theme", help="The theme name and variant to use, e.g., 'Catppuccin:dark'.", required=False)
    parser.add_argument("-s", "--saturation", type=float, default=1.0, help="Adjust saturation. 0.0 is desaturated, 1.0 is default, >1.0 is more saturated.")
    parser.add_argument("-d", "--darken", type=float, default=0, help="Darken (-100) or lighten (100) the image.")
    parser.add_argument("-l", "--list", action="store_true", help="Lists available color schemes.")
    parser.add_argument("output", nargs="?", help="Path for the output image file.")

    log_group = parser.add_mutually_exclusive_group()
    log_group.add_argument("-L", "--log-level", choices=LOG_LEVELS, default="INFO", help="Set the logging level.")
    log_group.add_argument("-v", "--verbose", action="store_true", help="Alias for -L DEBUG.")

    args = parser.parse_args()

    # Handle logging level first
    log_level = "DEBUG" if args.verbose else args.log_level
    setup_logging(log_level)

    # Prepare configuration directory and default theme
    config_dir = get_config_dir()
    create_default_theme_file(config_dir)

    # Handle list command
    if args.list:
        list_themes()
        sys.exit(0)

    # Validate required arguments for colorization
    if not all([args.input, args.theme, args.output]):
        parser.error("The -i, -t, and output arguments are required unless -l is used.")

    # Parse theme name and variant
    if ':' not in args.theme:
        logging.fatal("Theme selection must be in the format 'name:variant'.")
        sys.exit(1)

    theme_name, variant_name = args.theme.split(':', 1)

    # Main colorization process
    selected_colors_rgb = load_palette(theme_name, variant_name)
    output_image = colorize_image(args.input, selected_colors_rgb, args.darken, args.saturation)

    try:
        output_image.save(args.output)
        logging.info(f"Successfully saved colorized image to '{args.output}'")
    except Exception as e:
        logging.fatal(f"Error saving output image: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()