from __future__ import annotations

import argparse
import json
import os
import random
import subprocess
from importlib.resources import files
from pathlib import Path

from PIL import Image

PALETTES = json.loads(files("bilta_gilata1").joinpath("palettes.json").read_text())["palettes"]
DEFAULT_CACHE = Path(os.environ.get("BILTA_GILATA_CACHE", Path.home() / ".cache" / "bilta-gilata1"))


def quantize(source, target, size, palette_name):
    image = Image.open(source).convert("RGB").resize((size, size), Image.Resampling.BOX)
    colors = PALETTES[palette_name]["colors"][:256]
    palette = Image.new("P", (1, 1))
    values = [component for color in colors for component in bytes.fromhex(color[1:])]
    palette.putpalette(values + [0] * (768 - len(values)))
    image.quantize(palette=palette, dither=Image.Dither.NONE).save(target, "PNG")


def main():
    parser = argparse.ArgumentParser(prog="bilta-gilata1", description="Generate local pixel art on Apple Silicon")
    parser.add_argument("prompt", nargs="?")
    parser.add_argument("--palette", default="pico-8", choices=sorted(PALETTES))
    parser.add_argument("--size", type=int, default=128, choices=(16, 32, 48, 64, 128, 256))
    parser.add_argument("--reference", type=Path)
    parser.add_argument("--output", type=Path, default=Path("bilta-gilata1.png"))
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--list-palettes", action="store_true")
    args = parser.parse_args()
    if args.list_palettes:
        print("\n".join(sorted(PALETTES)))
        return 0
    if not args.prompt:
        parser.error("a prompt is required")
    executable = Path(os.environ.get("BILTA_GILATA_MFLUX", "mflux-generate-flux2"))
    DEFAULT_CACHE.mkdir(parents=True, exist_ok=True)
    work = DEFAULT_CACHE / "jobs" / str(random.randint(100000, 999999))
    work.mkdir(parents=True)
    master = work / "master.png"
    prompt = (f"pixel art sprite, {args.prompt}, one centered readable game asset, crisp square pixels, "
              f"limited {args.palette} palette, plain background, no text, no interface")
    command = [str(executable), "--model", "ar9av/FLUX.2-klein-4B-mflux-4bit", "--base-model", "flux2-klein-4b",
               "--lora", "Limbicnation/pixel-art-lora:pytorch_lora_weights.safetensors", "1.0", "--low-ram",
               "--mlx-cache-limit-gb", "1", "--prompt", prompt, "--steps", "4", "--guidance", "1.0",
               "--width", "512", "--height", "512", "--seed", str(args.seed or random.randint(1, 999999999)),
               "--output", str(master)]
    if args.reference:
        if not args.reference.is_file():
            raise SystemExit(f"Reference not found: {args.reference}")
        command += ["--image", str(args.reference.resolve()), "0.58"]
    env = dict(os.environ, HF_HOME=str(DEFAULT_CACHE / "models"))
    try:
        subprocess.run(command, check=True, env=env)
    except FileNotFoundError:
        raise SystemExit("mflux executable not found; reinstall with: pipx install bilta-gilata1")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    quantize(master, args.output, args.size, args.palette)
    print(args.output.resolve())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())