import contextlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any

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

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: Path | None = typer.Option(
		None,
		"--image-path",
		"-i",
		help="Path to source image. Required for image mode.",
	),
	scheme: str | None = typer.Option(
		None,
		"--scheme",
		"-s",
		help="Color scheme algorithm to use for image mode. Ignored in preset mode.",
		autocompletion=_complete_scheme_name,
	),
	preset: str | None = typer.Option(
		None,
		help="Name of a premade scheme in this format: <scheme>:<variant>",
		autocompletion=_complete_preset,
	),
	mode: str | None = typer.Option(
		None,
		"--mode",
		"-m",
		help="Mode of the preset scheme (dark or light).",
		autocompletion=_complete_mode,
	),
	accent: str | None = typer.Option(
		None,
		"--accent",
		"-a",
		help="Accent for schemes that support it (e.g. mauve).",
		autocompletion=_complete_accent,
	),
	preview: bool = typer.Option(
		False,
		"--preview",
		"-p",
		help="Preview scheme - internal use",
	),
):
	if not any([image_path, scheme, preset, mode, accent, preview]):
		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_DIR = Path(HOME + "/.cache/zshell/imagecache/thumbnails")
	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 thumbnail_cache_path(image_path: Path, thumb_dir: Path) -> Path:
		stat = image_path.stat()
		key = f"{image_path.stem}_{stat.st_size}_{int(stat.st_mtime)}"
		safe_key = re.sub(r"[^A-Za-z0-9._-]", "_", key)
		return thumb_dir / f"{safe_key}_thumbnail.jpg"

	def generate_thumbnail(
		image_path: Path, thumb_dir: Path, size=(128, 128)
	) -> Path:
		thumb_dir.mkdir(parents=True, exist_ok=True)
		cache_path = thumbnail_cache_path(image_path, thumb_dir)

		if cache_path.exists():
			return cache_path

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

		return cache_path

	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"

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

	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:
			with contextlib.suppress(OSError):
				qt_conf.write_text(new_text, encoding="utf-8")

	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) -> Path | None:
		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[Path | None, 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")

			with contextlib.suppress(OSError):
				shutil.copymode(tpl_path, out_path)

			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):
			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 f"#{argb_int & 0xFFFFFF:06X}"

	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}"
					)

			requested_mode = mode or config_mode
			resolved_mode = requested_mode
			if p_scheme in schemes:
				meta = schemes[p_scheme]
				variant = next(
					(vari for vari in meta.variants if vari.id == p_variant),
					None,
				)
				if (
					variant
					and requested_mode not in variant.modes
					and variant.modes
				):
					resolved_mode = sorted(variant.modes)[0]

			palette_obj = get_palette(
				p_scheme, p_variant, resolved_mode, accent=accent
			)
			colors = palette_obj.colors
			effective_mode = palette_obj.mode
			name = palette_obj.scheme
			flavor = palette_obj.variant

			display_name = schemes[p_scheme].name
			config["colors"]["presets"] = {
				"name": display_name,
				"variant": p_variant,
				"accent": accent or "",
			}
			tmp = CONFIG.with_suffix(".json.tmp")
			with tmp.open("w") as f:
				json.dump(config, f, indent=4)
			os.replace(tmp, CONFIG)

			seed = hex_to_hct(colors.get("primary", "#000000").lstrip("#"))
		else:
			image_path = image_path or Path(WALL_PATH)
			thumb_path = generate_thumbnail(image_path, THUMB_DIR)
			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

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

		if preview:
			print(json.dumps(output_dict))
			return

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

		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)
		tmp_output = OUTPUT.with_suffix(".json.tmp")
		with open(tmp_output, "w") as f:
			json.dump(output_dict, f, indent=4)
		os.replace(tmp_output, OUTPUT)
	except Exception as e:
		print(f"Error: {e}")