#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "accelerate>=1.14.0",
#     "diffusers>=0.39.0",
#     "gguf>=0.19.0",
#     "llama-cpp-python>=0.3.34",
#     "numpy>=2.5.1",
#     "pillow>=12.3.0",
#     "requests>=2.34.2",
#     "spandrel>=0.4.1",
#     "suntimes>=1.1.2",
#     "torch>=2.13.0",
#     "torchvision>=0.28.0",
#     "transformers>=5.14.1",
# ]
# [tool.uv.sources]
# torch = [
#   { index = "pytorch-cpu" },
# ]
# torchvision = [
#   { index = "pytorch-cpu" },
# ]
# [[tool.uv.index]]
# name = "pytorch-cpu"
# url = "https://download.pytorch.org/whl/cpu"
# explicit = true
# ///

import os

OFFLINE = os.environ.get("OFFLINE", "1") != "0"

os.environ["HF_HUB_OFFLINE"] = "1" if OFFLINE else "0"
os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"

from pathlib import Path
from typing import Any
import argparse
import datetime as dt
import logging
import random
import subprocess
import sys
import gc

import numpy as np
import requests
import torch
from llama_cpp import Llama
from huggingface_hub import hf_hub_download, snapshot_download
from diffusers import (
    Flux2KleinPipeline,
    Flux2Transformer2DModel,
    GGUFQuantizationConfig,
)
from PIL import Image
from spandrel import ModelLoader
from suntimes import SunTimes

logger = logging.getLogger("wallpaper")


CAP = torch.backends.cpu.get_cpu_capability()
COMPUTE_DTYPE = torch.bfloat16 if CAP.startswith("AVX512") else torch.float32


def step_context(
    now: dt.datetime,
    latitude: float | None,
    longitude: float | None,
    elevation: int = 0,
) -> str:

    Y = 2000  # dummy leap year to allow input X-02-29 (leap day)
    seasons = [
        ("winter", (dt.date(Y, 1, 1), dt.date(Y, 3, 20))),
        ("spring", (dt.date(Y, 3, 21), dt.date(Y, 6, 20))),
        ("summer", (dt.date(Y, 6, 21), dt.date(Y, 9, 22))),
        ("autumn", (dt.date(Y, 9, 23), dt.date(Y, 12, 20))),
        ("winter", (dt.date(Y, 12, 21), dt.date(Y, 12, 31))),
    ]

    date = now.date()
    date = date.replace(year=Y)
    season = next(season for season, (start, end) in seasons if start <= date <= end)

    today = now.strftime(f"Date and time: %A %B %d, %Y at %H:%M %Z. Season: {season}.")

    if latitude is None or longitude is None:
        return today

    sun = SunTimes(longitude, latitude, elevation)

    today_sr = sun.riselocal(now)
    today_ss = sun.setlocal(now)
    sun_times = (
        f"sunrise {today_sr.strftime('%H:%M')}, sunset {today_ss.strftime('%H:%M')}"
    )

    detail = "unknown"
    _WMO_SHORT = {
        0: "clear sky",
        1: "mostly clear",
        2: "partly cloudy",
        3: "overcast",
        45: "foggy",
        48: "rime fog",
        51: "light drizzle",
        53: "drizzle",
        55: "heavy drizzle",
        61: "light rain",
        63: "rain",
        65: "heavy rain",
        71: "light snow",
        73: "snow",
        75: "heavy snow",
        80: "showers",
        81: "moderate showers",
        82: "heavy showers",
        95: "thunderstorm",
        96: "thunderstorm + hail",
        99: "severe thunderstorm",
    }
    try:
        url = (
            f"https://api.open-meteo.com/v1/forecast?"
            f"latitude={latitude}&longitude={longitude}"
            f"&current=weather_code,temperature_2m&timezone=auto"
        )
        data = requests.get(url, timeout=10).json()["current"]
        code = data["weather_code"]
        code_desc = _WMO_SHORT.get(code, "")
        temp = float(data["temperature_2m"])
        detail = f"{code_desc} {temp:.0f}°C"
    except Exception as exc:
        logger.warning("Weather fetch failed (%s), continuing anyway.", exc)

    return f"{today} Daylight: {sun_times}. Weather: {detail}" if detail else today


def step_llm(
    repo_id: str,
    filename: str,
    system_prompt: str,
    max_tokens: int,
    taste_prompts: list[str],
    ctx: str,
) -> str:
    logger.info("Loading llm model...")

    model_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
    )

    llm = Llama(
        model_path=model_path,
        n_ctx=2048,
        verbose=False,
    )
    taste_prompt: list[str] = list(taste_prompts)
    random.shuffle(taste_prompt)
    taste_prompt_str = ", ".join(taste_prompt)

    prompt = (
        f"Taste reference:\n{taste_prompt_str}\n\n"
        f"Context: {ctx}\n\n"
        f"Reply with one text-to-image prompt only."
    )
    logger.info("Running llm...")
    output = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": prompt,
            },
        ],
        max_tokens=max_tokens,
    )

    del llm
    gc.collect()
    return output["choices"][0]["message"]["content"].strip()


def step_diffuse(
    repo_id: str,
    filename: str,
    gen_width: int,
    gen_height: int,
    steps: int,
    prompt: str,
) -> list[Any]:

    ckpt_path = f"https://huggingface.co/{repo_id}/blob/main/{filename}"

    # populate cache
    if not OFFLINE:
        snapshot_download("black-forest-labs/FLUX.2-klein-4B")

    # Phase 1: text encoding only
    logger.info("Encoding prompt...")
    pipe = Flux2KleinPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-klein-4B",
        transformer=None,
        vae=None,
        torch_dtype=torch.bfloat16,  # always, for memory purposes
        local_files_only=OFFLINE,
    )
    with torch.inference_mode():
        prompt_embeds, _text_ids = pipe.encode_prompt(prompt=prompt)
    del pipe
    gc.collect()
    prompt_embeds = prompt_embeds.to(COMPUTE_DTYPE)

    # Phase 2: denoise + decode, without added memory load
    logger.info("Denoising...")
    transformer = Flux2Transformer2DModel.from_single_file(
        ckpt_path,
        quantization_config=GGUFQuantizationConfig(compute_dtype=COMPUTE_DTYPE),
        torch_dtype=COMPUTE_DTYPE,
        config="black-forest-labs/FLUX.2-klein-4B",
        subfolder="transformer",
        device_map="cpu",
        local_files_only=OFFLINE,
    )
    pipe = Flux2KleinPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-klein-4B",
        text_encoder=None,
        tokenizer=None,
        transformer=transformer,
        torch_dtype=COMPUTE_DTYPE,
        local_files_only=OFFLINE,
    )

    out = pipe(
        prompt_embeds=prompt_embeds,
        width=gen_width,
        height=gen_height,
        num_inference_steps=steps,
    )
    del pipe, transformer, prompt_embeds
    gc.collect()
    return out.images


def upscale_tiled(
    model: Any,
    x: torch.Tensor,
    scale: int,
    tile: int = 256,
    overlap: int = 16,
) -> torch.Tensor:
    # x: [1, C, H, W] in [0, 1]. Peak memory is bounded by the tile size,
    # not the image size. Each tile is padded with `overlap` px of context,
    # which is cropped away on write-back, so tiles blend seamlessly.
    _, c, h, w = x.shape
    out = torch.zeros(1, c, h * scale, w * scale)
    stride = tile - 2 * overlap
    for y0 in range(0, h, stride):
        for x0 in range(0, w, stride):
            iy0, ix0 = max(y0 - overlap, 0), max(x0 - overlap, 0)
            iy1 = min(y0 + stride + overlap, h)
            ix1 = min(x0 + stride + overlap, w)
            with torch.inference_mode():
                up = model(x[:, :, iy0:iy1, ix0:ix1])
            oy0, ox0 = y0 * scale, x0 * scale
            oy1 = min(y0 + stride, h) * scale
            ox1 = min(x0 + stride, w) * scale
            py0, px0 = (y0 - iy0) * scale, (x0 - ix0) * scale
            out[:, :, oy0:oy1, ox0:ox1] = up[
                :, :, py0 : py0 + (oy1 - oy0), px0 : px0 + (ox1 - ox0)
            ]
            del up
    return out


def step_upscale(
    images: list[Any],
    target_width: int,
    target_height: int,
    repo_id: str,
    filename: str,
    tile: int,
) -> list[Any]:
    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
    model = ModelLoader().load_from_file(model_path)
    model.to("cpu").eval()

    results: list[Any] = []
    for im in images:
        arr = np.array(im.convert("RGB"), dtype=np.uint8)
        x = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).float().div(255)

        if tile > 0:
            up = upscale_tiled(model, x, model.scale, tile=tile)
        else:
            with torch.inference_mode():
                up = model(x)

        # Tensor is [1, C, H', W'] -> Permute to [H', W', C], scale from
        # [0, 1] to uint8 and clamp
        img_array = up[0].permute(1, 2, 0).mul(255).clamp(0, 255).byte().numpy()

        del x, up
        gc.collect()
        final_img = Image.fromarray(img_array)

        # by default 4x lands above the target (1280x720 -> 5120x2880), then
        # Lanczos downscales to the target, averaging away upscaler noise.
        target_size = (target_width, target_height)
        if (final_img.width, final_img.height) != target_size:
            final_img = final_img.resize(target_size, Image.LANCZOS)

        results.append(final_img)

    del model
    gc.collect()
    return results


def step_save(
    images: list[Any], output_dir: Path, now: dt.datetime, prompt_text: str
) -> list[Path]:
    output_dir.mkdir(parents=True, exist_ok=True)
    saved = []
    tag = now.strftime("%Y%m%d-%H%M%S")
    for i, img in enumerate(images):
        fname = f"wallpaper-{tag}" + (f"-{i}" if len(images) > 1 else "") + ".jpg"
        path = output_dir / fname
        img.save(path, "JPEG", quality=95)
        saved.append(path)

    txt_path = output_dir / f"prompt-{tag}.txt"
    txt_path.write_text(prompt_text)
    return saved


def step_wallpaper(command: list[str], saved: list[Path], store: Path) -> None:
    idx = random.randrange(len(saved))
    try:
        perma = store / "generated_wallpaper.jpg"
        _ = perma.open("wb").write(Path(saved[idx]).read_bytes())
        logger.info(f"Setting wallpaper to {perma}")
        _ = subprocess.run(command + [perma], check=True)
    except Exception as e:
        logger.warning(f"Could not set wallpaper: {e}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate wallpaper. Use OFFLINE=0 env var for first run.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    wall = parser.add_argument_group("wallpaper")
    wall.add_argument(
        "--taste-prompt",
        nargs="+",
        default=[
            "natural scenes",
            "serene",
            "trees",
            "plants",
            "landscape",
            "hypermodern concrete architecture",
            "moss",
            "rocks",
            "abstract concrete sculptures",
            "precise water channel",
            "possible distant structures",
        ],
        help="Taste keywords for the LLM art-director prompt.",
    )
    wall.add_argument(
        "--ephemeral-dir", default="/tmp/wallpaper/", help="Output directory."
    )
    wall.add_argument(
        "--single-file-persist-path",
        default="~/media/pictures/",
        help="Single target for persistance.",
    )
    wall.add_argument(
        "--target-width", type=int, default=3840, help="Final wallpaper width."
    )
    wall.add_argument(
        "--target-height", type=int, default=2160, help="Final wallpaper height."
    )
    wall.add_argument(
        "--set-command",
        nargs="+",
        default=["awww", "img"],
        help="Command to run when setting wallpaper.",
    )

    llm = parser.add_argument_group("llm")
    llm.add_argument(
        "--llm-repo-id",
        default="unsloth/gemma-4-E4B-it-GGUF",
        help="HuggingFace repo for LLM to use to prompt image model.",
    )
    llm.add_argument(
        "--llm-filename",
        default="gemma-4-E4B-it-Q4_K_M.gguf",
        help="File name of LLM in the repo to use.",
    )
    llm.add_argument(
        "--llm-max-tokens",
        type=int,
        default=200,
        help="Max tokens of prompt to generate.",
    )
    llm.add_argument(
        "--llm-system-prompt",
        default=(
            "You are a wallpaper art director. "
            "Given the user's taste and the context, write one detailed text-to-image prompt for a diffusion model. "
            "Not all of the user's taste should be present, it's a suggestion and you should choose only a subset and improvise on the themes. "
            "Times after sunset and before sunrise should generate dark images for night time. "
            "Times during the day should track the brightness of the hour. "
            "The weather should be reflected in your prompt. "
            "The season should set the overall mood of the prompt. "
            "Do not include the time, date, or location in your prompt. "
            "Output ONLY the image prompt — no explanation or preamble."
        ),
        help="System prompt for LLM.",
    )

    flux2 = parser.add_argument_group("Flux Klein 4B settings")
    flux2.add_argument(
        "--flux2-klein-repo-id",
        default="unsloth/FLUX.2-klein-4B-GGUF",
        help="HuggingFace repo for quantised model.",
    )
    flux2.add_argument(
        "--flux2-klein-filename",
        default="flux-2-klein-4b-Q8_0.gguf",
        help="File name of quantised mode in the repo.",
    )
    flux2.add_argument(
        "--flux2-klein-steps",
        type=int,
        default=4,
        help="Steps to use when generating image.",
    )
    flux2.add_argument(
        "--flux2-klein-gen-width", type=int, default=1280, help="Generated image width."
    )
    flux2.add_argument(
        "--flux2-klein-gen-height",
        type=int,
        default=720,
        help="Generated image height.",
    )

    sr = parser.add_argument_group("super-resolution")
    sr.add_argument(
        "--sr-repo-id",
        default="Comfy-Org/Real-ESRGAN_repackaged",
        help="HuggingFace repo hosting the ESRGAN-family checkpoint.",
    )
    sr.add_argument(
        "--sr-filename",
        default="RealESRGAN_x4plus.safetensors",
        help="Checkpoint filename in the repo (any spandrel-supported "
        ".pth/.safetensors works).",
    )
    sr.add_argument(
        "--sr-tile",
        type=int,
        default=256,
        help="Tile size for tiled upscaling; bounds peak memory on "
        "constrained machines. 0 = full-frame (best quality).",
    )

    wx = parser.add_argument_group("weather")
    wx.add_argument(
        "--latitude", type=float, default=51.75, help="Latitude for weather."
    )
    wx.add_argument(
        "--longitude", type=float, default=-1.25, help="Longitude for weather."
    )
    wx.add_argument(
        "--elevation",
        type=int,
        default=61,
        help="Altitude (m) for sunrise/sunset times.",
    )

    return parser.parse_args()


def run() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s  %(message)s",
        datefmt="%H:%M:%S",
        stream=sys.stderr,
    )

    args = parse_args()

    # crazy
    now = dt.datetime.now(dt.datetime.now().astimezone().tzinfo)

    logger.info(f"Capability {CAP}, dtype {COMPUTE_DTYPE}")

    # 1. context
    ctx = step_context(now, args.latitude, args.longitude, args.elevation)
    logger.info(f"Context: {ctx}")

    # 2. prompt rewrite via local LLM
    prompt_text = step_llm(
        args.llm_repo_id,
        args.llm_filename,
        args.llm_system_prompt,
        args.llm_max_tokens,
        list(args.taste_prompt),
        ctx,
    )
    logger.info(f"Prompt: {prompt_text}")

    # 3. diffusion at base resolution
    logger.info(
        "Diffuse: %dx%d @ %d steps",
        args.flux2_klein_gen_width,
        args.flux2_klein_gen_height,
        args.flux2_klein_steps,
    )
    raw_images = step_diffuse(
        args.flux2_klein_repo_id,
        args.flux2_klein_filename,
        args.flux2_klein_gen_width,
        args.flux2_klein_gen_height,
        args.flux2_klein_steps,
        prompt_text,
    )

    # 4. super-resolution then scale
    logger.info(
        "Upscale then rescale: %dx%d -> %dx%d -> %dx%d via %s with tiling %d",
        args.flux2_klein_gen_width,
        args.flux2_klein_gen_height,
        args.flux2_klein_gen_width * 4,
        args.flux2_klein_gen_height * 4,
        args.target_width,
        args.target_height,
        args.sr_filename,
        args.sr_tile,
    )
    upscaled = step_upscale(
        raw_images,
        args.target_width,
        args.target_height,
        args.sr_repo_id,
        args.sr_filename,
        args.sr_tile,
    )

    # 5. save to configured output directory
    saved = step_save(upscaled, Path(args.ephemeral_dir).expanduser(), now, prompt_text)

    for p in saved:
        logger.info("Saved: %s", p.resolve())

    # 6. set wallpaper
    step_wallpaper(
        args.set_command,
        saved,
        Path(args.single_file_persist_path).expanduser(),
    )


if __name__ == "__main__":
    run()