import typer
import json
import shutil
import os
import sys
import re
import subprocess

from jinja2 import Environment, FileSystemLoader, StrictUndefined, Undefined
from typing import Any, Optional, Tuple
from zshell.utils.schemepalettes import get_palette, list_schemes, resolve_preset
from pathlib import Path
from PIL import Image
from materialyoucolor.quantize import QuantizeCelebi
from materialyoucolor.score.score import Score
from materialyoucolor.dynamiccolor.material_dynamic_colors import MaterialDynamicColors
from materialyoucolor.hct.hct import Hct
from materialyoucolor.utils.color_utils import argb_from_rgb
from materialyoucolor.utils.math_utils import (
    difference_degrees,
    rotation_direction,
    sanitize_degrees_double,
)

app = typer.Typer()


def _complete_scheme_name(incomplete):
    schemes = [
        "fruit-salad",
        "expressive",
        "monochrome",
        "rainbow",
        "tonal-spot",
        "neutral",
        "fidelity",
        "content",
        "vibrant",
    ]
    return [s for s in schemes if incomplete in s]


def _complete_preset(incomplete):
    results = []
    for sid, meta in list_schemes().items():
        for v in meta.variants:
            preset = f"{sid}:{v.id}"
            if incomplete in preset:
                results.append((preset, f"{meta.name} - {v.name}"))
    return results


def _complete_mode(incomplete):
    return [m for m in ("dark", "light") if incomplete in m]


def _complete_accent(ctx, incomplete):
    preset_val = ctx.params.get("preset")
    if preset_val:
        try:
            p_scheme, p_variant = resolve_preset(preset_val)
            for v in list_schemes()[p_scheme].variants:
                if v.id == p_variant:
                    return [a for a in v.accents if incomplete in a]
        except (ValueError, KeyError):
            pass
    all_accents = set()
    for meta in list_schemes().values():
        for v in meta.variants:
            all_accents.update(v.accents)
    return [a for a in sorted(all_accents) if incomplete in a]


@app.command()
def list_presets(
    json_format: bool = typer.Option(False, "--json", help="Output in JSON format"),
):
    schemes = list_schemes()
    if json_format:
        out = {}
        for sid, meta in sorted(schemes.items()):
            variants = {}
            for v in meta.variants:
                entry: dict[str, Any] = {"modes": sorted(v.modes)}
                if v.accents:
                    entry["accents"] = sorted(v.accents)
                    entry["default_accent"] = sorted(v.accents)[0]
                variants[v.id] = entry
            out[meta.name] = {
                "id": sid,
                "variants": variants,
            }
        print(json.dumps({"presets": out}, indent=2))
    else:
        for sid, meta in sorted(schemes.items()):
            var_list = []
            for v in meta.variants:
                parts = [f"{v.id} ({', '.join(sorted(v.modes))})"]
                if v.accents:
                    parts.append(f"accents: {', '.join(v.accents)}")
                var_list.append(" | ".join(parts))
            print(f"{meta.name} ({sid})")
            print(f"  Variants: {', '.join(var_list)}")
            print()


@app.command()
def generate(
    image_path: Optional[Path] = typer.Option(
        None, help="Path to source image. Required for image mode."
    ),
    scheme: Optional[str] = typer.Option(
        None,
        help="Color scheme algorithm to use for image mode. Ignored in preset mode.",
        autocompletion=_complete_scheme_name,
    ),
    preset: Optional[str] = typer.Option(
        None,
        help="Name of a premade scheme in this format: <scheme>:<variant>",
        autocompletion=_complete_preset,
    ),
    mode: Optional[str] = typer.Option(
        None,
        help="Mode of the preset scheme (dark or light).",
        autocompletion=_complete_mode,
    ),
    accent: Optional[str] = typer.Option(
        None,
        help="Accent for schemes that support it (e.g. mauve).",
        autocompletion=_complete_accent,
    ),
):
    if not any([image_path, scheme, preset, mode, accent]):
        print(
            "Hint: use --preset <scheme>:<variant> or --image-path <path>",
            file=sys.stderr,
        )

    HOME = str(os.getenv("HOME"))
    OUTPUT = Path(HOME + "/.local/state/zshell/scheme.json")
    SEQ_STATE = Path(HOME + "/.local/state/zshell/sequences.txt")
    THUMB_PATH = Path(HOME + "/.cache/zshell/imagecache/thumbnail.jpg")
    WALL_DIR_PATH = Path(HOME + "/.local/state/zshell/wallpaper_path.json")

    TEMPLATE_DIR = Path(HOME + "/.config/zshell/templates")
    WALL_PATH = Path()
    CONFIG = Path(HOME + "/.config/zshell/config.json")

    if preset is not None and image_path is not None:
        raise typer.BadParameter("Use either --image-path or --preset, not both.")

    def get_scheme_class(scheme_name: str):
        match scheme_name:
            case "fruit-salad":
                from materialyoucolor.scheme.scheme_fruit_salad import SchemeFruitSalad

                return SchemeFruitSalad
            case "expressive":
                from materialyoucolor.scheme.scheme_expressive import SchemeExpressive

                return SchemeExpressive
            case "monochrome":
                from materialyoucolor.scheme.scheme_monochrome import SchemeMonochrome

                return SchemeMonochrome
            case "rainbow":
                from materialyoucolor.scheme.scheme_rainbow import SchemeRainbow

                return SchemeRainbow
            case "tonal-spot":
                from materialyoucolor.scheme.scheme_tonal_spot import SchemeTonalSpot

                return SchemeTonalSpot
            case "neutral":
                from materialyoucolor.scheme.scheme_neutral import SchemeNeutral

                return SchemeNeutral
            case "fidelity":
                from materialyoucolor.scheme.scheme_fidelity import SchemeFidelity

                return SchemeFidelity
            case "content":
                from materialyoucolor.scheme.scheme_content import SchemeContent

                return SchemeContent
            case "vibrant":
                from materialyoucolor.scheme.scheme_vibrant import SchemeVibrant

                return SchemeVibrant
            case _:
                from materialyoucolor.scheme.scheme_fruit_salad import SchemeFruitSalad

                return SchemeFruitSalad

    def hex_to_hct(hex_color: str) -> Hct:
        s = hex_color.strip()
        if s.startswith("#"):
            s = s[1:]
        if len(s) != 6:
            raise ValueError(f"Expected 6-digit hex color, got: {hex_color!r}")
        return Hct.from_int(int("0xFF" + s, 16))

    LIGHT_GRUVBOX = list(
        map(
            hex_to_hct,
            [
                "FDF9F3",
                "FF6188",
                "A9DC76",
                "FC9867",
                "FFD866",
                "F47FD4",
                "78DCE8",
                "333034",
                "121212",
                "FF6188",
                "A9DC76",
                "FC9867",
                "FFD866",
                "F47FD4",
                "78DCE8",
                "333034",
            ],
        )
    )

    DARK_GRUVBOX = list(
        map(
            hex_to_hct,
            [
                "282828",
                "CC241D",
                "98971A",
                "D79921",
                "458588",
                "B16286",
                "689D6A",
                "A89984",
                "928374",
                "FB4934",
                "B8BB26",
                "FABD2F",
                "83A598",
                "D3869B",
                "8EC07C",
                "EBDBB2",
            ],
        )
    )

    with WALL_DIR_PATH.open() as f:
        path = json.load(f)["currentWallpaperPath"]
        WALL_PATH = path

    def lighten(color: Hct, amount: float) -> Hct:
        diff = (100 - color.tone) * amount
        tone = max(0.0, min(100.0, color.tone + diff))
        chroma = max(0.0, color.chroma + diff / 5)
        return Hct.from_hct(color.hue, chroma, tone)

    def darken(color: Hct, amount: float) -> Hct:
        diff = color.tone * amount
        tone = max(0.0, min(100.0, color.tone - diff))
        chroma = max(0.0, color.chroma - diff / 5)
        return Hct.from_hct(color.hue, chroma, tone)

    def grayscale(color: Hct, light: bool) -> Hct:
        color = darken(color, 0.35) if light else lighten(color, 0.65)
        color.chroma = 0
        return color

    def harmonize(from_hct: Hct, to_hct: Hct, tone_boost: float) -> Hct:
        diff = difference_degrees(from_hct.hue, to_hct.hue)
        rotation = min(diff * 0.8, 100)
        output_hue = sanitize_degrees_double(
            from_hct.hue + rotation * rotation_direction(from_hct.hue, to_hct.hue)
        )
        tone = max(0.0, min(100.0, from_hct.tone * (1 + tone_boost)))
        return Hct.from_hct(output_hue, from_hct.chroma, tone)

    def terminal_palette(
        colors: dict[str, str], mode: str, variant: str
    ) -> dict[str, str]:
        light = mode.lower() == "light"

        key_hex = (
            colors.get("primary_paletteKeyColor")
            or colors.get("primaryPaletteKeyColor")
            or colors.get("primary")
            or int_to_hex(seed.to_int())
        )
        key_hct = hex_to_hct(key_hex)

        base = LIGHT_GRUVBOX if light else DARK_GRUVBOX
        out: dict[str, str] = {}

        is_mono = variant.lower() == "monochrome"

        for i, base_hct in enumerate(base):
            if is_mono:
                h = grayscale(base_hct, light)
            else:
                tone_boost = (0.35 if i < 8 else 0.2) * (-1 if light else 1)
                h = harmonize(base_hct, key_hct, tone_boost)

            out[f"term{i}"] = int_to_hex(h.to_int())

        return out

    def generate_thumbnail(image_path, thumbnail_path, size=(128, 128)):
        thumbnail_file = Path(thumbnail_path)

        image = Image.open(image_path)
        image = image.convert("RGB")
        image.thumbnail(size, Image.Resampling.NEAREST)

        thumbnail_file.parent.mkdir(parents=True, exist_ok=True)
        image.save(thumbnail_path, "JPEG")

    def apply_terms(sequences: str, sequences_tmux: str, state_path: Path) -> None:
        state_path.parent.mkdir(parents=True, exist_ok=True)
        state_path.write_text(sequences, encoding="utf-8")

        pts_path = Path("/dev/pts")
        if not pts_path.exists():
            return

        O_NOCTTY = getattr(os, "O_NOCTTY", 0)

        for pt in pts_path.iterdir():
            if not pt.name.isdigit():
                continue
            try:
                fd = os.open(str(pt), os.O_WRONLY | os.O_NONBLOCK | O_NOCTTY)
                try:
                    os.write(fd, sequences_tmux.encode())
                    os.write(fd, sequences.encode())
                finally:
                    os.close(fd)
            except (PermissionError, OSError, BlockingIOError):
                pass

    def smart_mode(image_path: Path) -> str:
        is_dark = ""

        with Image.open(image_path) as img:
            img.thumbnail((1, 1), Image.Resampling.LANCZOS)
            px = img.getpixel((0, 0))
            if isinstance(px, (int, float)):
                r = g = b = int(px)
            elif px is not None:
                r, g, b = int(px[0]), int(px[1]), int(px[2])
            else:
                r = g = b = 0
            hct = Hct.from_int(argb_from_rgb(r, g, b))
            is_dark = "light" if hct.tone > 50 else "dark"

        return is_dark

    def apply_gtk_mode(mode: str) -> None:
        mode = mode.lower()
        preference = "prefer-dark" if mode == "dark" else "prefer-light"

        try:
            subprocess.run(
                [
                    "gsettings",
                    "set",
                    "org.gnome.desktop.interface",
                    "color-scheme",
                    preference,
                ],
                check=False,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        except FileNotFoundError:
            pass

    def apply_qt_mode(mode: str, home: str) -> None:
        mode = mode.lower()
        qt_conf = Path(home) / ".config/qt6ct/qt6ct.conf"

        if not qt_conf.exists():
            return

        try:
            text = qt_conf.read_text(encoding="utf-8")
        except OSError:
            return

        target = "Dark.colors" if mode == "dark" else "Light.colors"

        new_text, count = re.subn(
            r"^(color_scheme_path=.*?)(?:Light|Dark)\.colors\s*$",
            rf"\1{target}",
            text,
            flags=re.MULTILINE,
        )

        if count > 0 and new_text != text:
            try:
                qt_conf.write_text(new_text, encoding="utf-8")
            except OSError:
                pass

    def build_template_context(
        *,
        colors: dict[str, str],
        seed: Hct,
        mode: str,
        wallpaper_path: str,
        name: str,
        flavor: str,
        variant: str,
    ) -> dict[str, Any]:
        ctx: dict[str, Any] = {
            "mode": mode,
            "wallpaper_path": wallpaper_path,
            "source_color": int_to_hex(seed.to_int()),
            "name": name,
            "seed": seed.to_int(),
            "flavor": flavor,
            "variant": variant,
            "colors": colors,
        }

        for k, v in colors.items():
            ctx[k] = v
            ctx[f"m3{k}"] = v

        term = terminal_palette(colors, mode, variant)
        ctx.update(term)
        ctx["term"] = [term[f"term{i}"] for i in range(16)]

        seq = make_sequences(
            term=term,
            foreground=ctx["m3onSurface"],
            background=ctx["m3surface"],
        )
        ctx["sequences"] = seq
        ctx["sequences_tmux"] = tmux_wrap_sequences(seq)

        return ctx

    def make_sequences(
        *,
        term: dict[str, str],
        foreground: str,
        background: str,
    ) -> str:
        ESC = "\x1b"
        ST = ESC + "\\"

        parts: list[str] = []

        for i in range(16):
            parts.append(f"{ESC}]4;{i};{term[f'term{i}']}{ST}")

        parts.append(f"{ESC}]10;{foreground}{ST}")
        parts.append(f"{ESC}]11;{background}{ST}")

        return "".join(parts)

    def tmux_wrap_sequences(seq: str) -> str:
        ESC = "\x1b"
        return f"{ESC}Ptmux;{seq.replace(ESC, ESC + ESC)}{ESC}\\"

    def parse_output_directive(first_line: str) -> Optional[Path]:
        s = first_line.strip()
        if not s.startswith("#") or s.startswith("#!"):
            return None

        target = s[1:].strip()
        if not target:
            return None

        expanded = os.path.expandvars(os.path.expanduser(target))
        return Path(expanded)

    def split_directive_and_body(text: str) -> Tuple[Optional[Path], str]:
        lines = text.splitlines(keepends=True)
        if not lines:
            return None, ""

        out_path = parse_output_directive(lines[0])
        if out_path is None:
            return None, text

        body = "".join(lines[1:])
        return out_path, body

    def render_all_templates(
        templates_dir: Path,
        context: dict[str, object],
        *,
        strict: bool = True,
    ) -> list[Path]:
        undefined_cls = StrictUndefined if strict else Undefined
        env = Environment(
            loader=FileSystemLoader(str(templates_dir)),
            autoescape=False,
            keep_trailing_newline=True,
            undefined=undefined_cls,
        )

        rendered_outputs: list[Path] = []

        for tpl_path in sorted(p for p in templates_dir.rglob("*") if p.is_file()):
            rel = tpl_path.relative_to(templates_dir)

            if any(part.startswith(".") for part in rel.parts):
                continue

            raw = tpl_path.read_text(encoding="utf-8")
            out_path, body = split_directive_and_body(raw)
            if out_path is None:
                continue

            out_path.parent.mkdir(parents=True, exist_ok=True)

            try:
                template = env.from_string(body)
                text = template.render(**context)
            except Exception as e:
                raise RuntimeError(f"Template render failed for '{rel}': {e}") from e

            out_path.write_text(text, encoding="utf-8")

            try:
                shutil.copymode(tpl_path, out_path)
            except OSError:
                pass

            rendered_outputs.append(out_path)

        return rendered_outputs

    def seed_from_image(image_path: Path) -> Hct:
        image = Image.open(image_path)
        pixel_len = image.width * image.height
        image_data = image.getdata()

        quality = 1
        pixel_array = [image_data[_] for _ in range(0, pixel_len, quality)]

        result = QuantizeCelebi(pixel_array, 128)
        return Hct.from_int(Score.score(result)[0])

    def generate_color_scheme(seed: Hct, mode: str, scheme_class) -> dict[str, str]:

        is_dark = mode.lower() == "dark"

        scheme = scheme_class(seed, is_dark, 0.0)

        color_dict = {}
        for color in vars(MaterialDynamicColors).keys():
            color_name = getattr(MaterialDynamicColors, color)
            if hasattr(color_name, "get_hct"):
                color_int = color_name.get_hct(scheme).to_int()
                color_dict[color] = int_to_hex(color_int)

        return color_dict

    def int_to_hex(argb_int):
        return "#{:06X}".format(argb_int & 0xFFFFFF)

    try:
        with CONFIG.open() as f:
            config = json.load(f)

        scheme_type = config["colors"].get("schemeType", "fruit-salad")
        scheme = scheme or scheme_type
        assert isinstance(scheme, str)
        config_mode = config["general"]["color"]["mode"]
        smart = bool(config["general"]["color"].get("smart", False))
        scheme_class = get_scheme_class(scheme)

        p_variant = "default"
        if preset:
            p_scheme, p_variant = resolve_preset(preset)
            schemes = list_schemes()
            if accent and p_scheme in schemes:
                meta = schemes[p_scheme]
                var_accents = next(
                    (v.accents for v in meta.variants if v.id == p_variant), ()
                )
                if accent not in var_accents:
                    available = ", ".join(var_accents) if var_accents else "none"
                    raise typer.BadParameter(
                        f"Accent '{accent}' not available for '{p_scheme}:{p_variant}'. Available accents: {available}"
                    )
            palette_obj = get_palette(
                p_scheme, p_variant, mode or config_mode, accent=accent
            )
            colors = palette_obj.colors
            effective_mode = palette_obj.mode
            name = palette_obj.scheme
            flavor = palette_obj.variant

            seed = hex_to_hct(colors.get("primary", "#000000").lstrip("#"))
        else:
            image_path = image_path or Path(WALL_PATH)
            generate_thumbnail(image_path, str(THUMB_PATH))
            seed = seed_from_image(THUMB_PATH)
            name = "dynamic"
            flavor = "default"

            if smart:
                effective_mode = smart_mode(THUMB_PATH)
            elif mode is not None:
                effective_mode = mode
            else:
                effective_mode = config_mode

            colors = generate_color_scheme(seed, effective_mode, scheme_class)

        variant_val = scheme if not preset else p_variant

        if smart and not preset:
            apply_gtk_mode(effective_mode)
            apply_qt_mode(effective_mode, HOME)

        output_dict = {
            "name": name,
            "flavor": flavor,
            "mode": effective_mode,
            "variant": variant_val,
            "colors": colors,
            "seed": seed.to_int(),
        }

        if TEMPLATE_DIR is not None:
            wp = str(WALL_PATH)
            ctx = build_template_context(
                colors=colors,
                seed=seed,
                mode=effective_mode,
                wallpaper_path=wp,
                name=name,
                flavor=flavor,
                variant=variant_val,
            )

            rendered = render_all_templates(
                templates_dir=TEMPLATE_DIR,
                context=ctx,
            )

            apply_terms(ctx["sequences"], ctx["sequences_tmux"], SEQ_STATE)

            for p in rendered:
                print(f"rendered: {p}")

        OUTPUT.parent.mkdir(parents=True, exist_ok=True)
        with open(OUTPUT, "w") as f:
            json.dump(output_dict, f, indent=4)
    except Exception as e:
        print(f"Error: {e}")