#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "requests>=2.34.2",
#     "suntimes>=1.1.2",
# ]
# ///

import argparse
import datetime as dt
import logging
import random
import re
import shutil
import subprocess
import sys
from pathlib import Path

import requests
from suntimes import SunTimes

logger = logging.getLogger("wallpaper")


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


def positive_int(value: str) -> int:
    parsed = int(value)
    if parsed <= 0:
        raise argparse.ArgumentTypeError("must be greater than zero")
    return parsed


def non_negative_int(value: str) -> int:
    parsed = int(value)
    if parsed < 0:
        raise argparse.ArgumentTypeError("must be zero or greater")
    return parsed


def non_negative_float(value: str) -> float:
    parsed = float(value)
    if parsed < 0:
        raise argparse.ArgumentTypeError("must be zero or greater")
    return parsed


def probability(value: str) -> float:
    parsed = float(value)
    if not 0.0 <= parsed <= 1.0:
        raise argparse.ArgumentTypeError("must be between zero and one")
    return parsed


def positive_fraction(value: str) -> float:
    parsed = float(value)
    if not 0.0 < parsed <= 1.0:
        raise argparse.ArgumentTypeError(
            "must be greater than zero and no greater than one"
        )
    return parsed


def step_context(
    now: dt.datetime,
    latitude: float | None,
    longitude: float | None,
    elevation: int = 0,
) -> str:
    year = 2000  # Dummy leap year to allow February 29.
    seasons = [
        ("winter", (dt.date(year, 1, 1), dt.date(year, 3, 20))),
        ("spring", (dt.date(year, 3, 21), dt.date(year, 6, 20))),
        ("summer", (dt.date(year, 6, 21), dt.date(year, 9, 22))),
        ("autumn", (dt.date(year, 9, 23), dt.date(year, 12, 20))),
        ("winter", (dt.date(year, 12, 21), dt.date(year, 12, 31))),
    ]

    date = now.date().replace(year=year)
    season = next(
        season_name for season_name, (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)
    sunrise = sun.riselocal(now)
    sunset = sun.setlocal(now)
    sun_times = (
        f"sunrise {sunrise.strftime('%H:%M')}, sunset {sunset.strftime('%H:%M')}"
    )

    weather = "unknown"
    try:
        response = requests.get(
            "https://api.open-meteo.com/v1/forecast",
            params={
                "latitude": latitude,
                "longitude": longitude,
                "current": "weather_code,temperature_2m",
                "timezone": "auto",
            },
            timeout=10,
        )
        response.raise_for_status()

        data = response.json()["current"]
        code = int(data["weather_code"])
        description = WMO_SHORT.get(code, f"weather code {code}")
        temperature = float(data["temperature_2m"])
        weather = f"{description}, {temperature:.0f}°C"
    except (requests.RequestException, KeyError, TypeError, ValueError) as exc:
        logger.warning("Weather fetch failed (%s), continuing anyway.", exc)

    return f"{today} Daylight: {sun_times}. Weather: {weather}."


def clean_llm_output(text: str) -> str:
    # llama-completion end marker.
    text = re.sub(
        r"\s*\[\s*end of text\s*\]\s*$",
        "",
        text,
        flags=re.IGNORECASE,
    )

    # Gemma 4 reasoning channel.
    text = re.sub(
        r"<\|channel>thought\s*.*?<channel\|>\s*",
        "",
        text,
        flags=re.DOTALL | re.IGNORECASE,
    )

    # DeepSeek/Qwen-style reasoning tags.
    text = re.sub(
        r"<think>.*?</think>",
        "",
        text,
        flags=re.DOTALL | re.IGNORECASE,
    )

    # Detect truncated reasoning rather than sending it to sd-cli.
    if re.match(
        r"^\s*(?:<\|channel>thought\b|<think>)",
        text,
        flags=re.IGNORECASE,
    ):
        raise RuntimeError(
            "The model generated reasoning but did not reach a final answer. Check that your llama.cpp build supports '--reasoning off', or increase --llm-max-tokens."
        )

    return text.strip()


def step_llm(
    system_prompt: str,
    model_path: Path,
    max_tokens: int,
    taste_prompts: list[str],
    taste_sample_frac: float,
    biomes: list[str],
    perspectives: list[str],
    temperature: float,
    top_p: float,
    min_p: float,
    context: str,
    explicit_instructions: str,
) -> str:
    if not taste_prompts:
        raise ValueError("At least one taste prompt is required")
    if not biomes:
        raise ValueError("At least one biome is required")
    if not perspectives:
        raise ValueError("At least one perspective is required")

    total = len(taste_prompts)
    sample_size = max(1, round(taste_sample_frac * total))
    sampled_tastes = random.sample(taste_prompts, sample_size)
    taste_prompt_text = ", ".join(sampled_tastes)

    biome = random.choice(biomes)
    perspective = random.choice(perspectives)

    logger.info(
        "Taste (%d/%d sampled): %s",
        sample_size,
        total,
        taste_prompt_text,
    )
    logger.info("Biome: %s", biome)
    logger.info("Perspective: %s", perspective)

    user_prompt = (
        f"Taste reference: {taste_prompt_text}. "
        f"Required biome: {biome}. "
        f"Required perspective: {perspective}. "
        f"Context: {context}. "
        "Reply with one text-to-image prompt only. /nothink"
    )

    cmd = [
        "llama-completion",
        "-m",
        str(model_path.expanduser()),
        "-c",
        "2048",
        "-n",
        str(max_tokens),
        "--temp",
        str(temperature),
        "--top-p",
        str(top_p),
        "--min-p",
        str(min_p),
        "--jinja",
        "--single-turn",
        "-sys",
        system_prompt,
        "-p",
        user_prompt,
        "--simple-io",
        "--no-display-prompt",
        "--no-perf",
        "--log-verbosity",
        "1",
    ]

    logger.info("Running llama.cpp ...")
    try:
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
    except FileNotFoundError as exc:
        raise RuntimeError("Could not find the llama-completion executable") from exc

    if result.returncode != 0:
        raise RuntimeError(
            f"llama-completion exited with {result.returncode}:\n"
            f"{result.stderr.strip()}"
        )

    generated_prompt = clean_llm_output(result.stdout)
    if not generated_prompt:
        raise RuntimeError(
            "llama-completion completed successfully but produced no prompt"
        )

    explicit_instructions = explicit_instructions.strip()
    if explicit_instructions:
        return f"{generated_prompt} {explicit_instructions}"
    return generated_prompt


def step_diffuse(
    *,
    diffusion_model: Path,
    vae: Path,
    llm: Path,
    cfg_scale: float,
    steps: int,
    flash_attention: bool,
    generation_width: int,
    generation_height: int,
    upscale_model: Path,
    upscale_repeats: int,
    upscale_tile_size: int,
    output_dir: Path,
    now: dt.datetime,
    prompt_text: str,
) -> Path:
    output_dir.mkdir(parents=True, exist_ok=True)
    tag = now.strftime("%Y%m%d-%H%M%S")

    output_path = output_dir / f"wallpaper-{tag}.png"
    prompt_path = output_dir / f"prompt-{tag}.txt"

    cmd = [
        "sd-cli",
        "--diffusion-model",
        str(diffusion_model.expanduser()),
        "--vae",
        str(vae.expanduser()),
        "--llm",
        str(llm.expanduser()),
        "--cfg-scale",
        str(cfg_scale),
        "--steps",
        str(steps),
        "-W",
        str(generation_width),
        "-H",
        str(generation_height),
        "--upscale-model",
        str(upscale_model.expanduser()),
        "--upscale-repeats",
        str(upscale_repeats),
        "--upscale-tile-size",
        str(upscale_tile_size),
        "-p",
        prompt_text,
        "-o",
        str(output_path),
    ]

    if flash_attention:
        cmd.append("--fa")

    logger.info(
        "Diffuse at %dx%d",
        generation_width,
        generation_height,
    )

    try:
        subprocess.run(cmd, check=True)
    except FileNotFoundError as exc:
        raise RuntimeError("Could not find the sd-cli executable") from exc
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(f"sd-cli exited with status {exc.returncode}") from exc

    if not output_path.is_file():
        raise RuntimeError(
            f"sd-cli completed successfully but did not create {output_path}"
        )

    prompt_path.write_text(prompt_text, encoding="utf-8")
    return output_path


def step_wallpaper(command: list[str], saved: Path, store: Path) -> None:
    try:
        store.mkdir(parents=True, exist_ok=True)
        permanent_path = store / "generated_wallpaper.png"
        shutil.copyfile(saved, permanent_path)

        logger.info("Setting wallpaper to %s", permanent_path)
        subprocess.run(
            [*command, str(permanent_path)],
            check=True,
        )
    except (OSError, subprocess.CalledProcessError) as exc:
        logger.warning("Could not set wallpaper: %s", exc)


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=[
            "monumental pristine concrete architecture",
            "hypermodern flawless surfaces beside weathered ancient ruins",
            "colossal enigmatic megastructure looming on the horizon",
            "elevated walkways and straight causeways",
            "monumental staircases and terraced platforms",
            "freestanding gates and arched portals",
            "open plazas and courtyards",
            "abstract concrete sculptures",
            "towering highly stylised humanoid statues",
            "precise water channels and reflecting pools",
            "clean geometric forms at impossible scale",
            "distinct scales of objects with some very large for contrast",
            "serene, deserted, contemplative",
            "completely out of place architecture",
        ],
        help="Core aesthetic keywords shared across all biomes.",
    )
    wall.add_argument(
        "--taste-sample-frac",
        type=positive_fraction,
        default=0.8,
        help=(
            "Fraction of taste keywords randomly sampled per run; lower "
            "values give more varied, less cluttered prompts."
        ),
    )
    wall.add_argument(
        "--biomes",
        nargs="+",
        default=[
            "mediterranean: sun-bleached stone terraces, azure sea, rocky coastline, cypress and olive trees, warm limestone cliffs, salt haze over the water",
            "autumn forest: dense woodland in golden and crimson foliage, fallen leaves, thin low mist between trunks, damp moss-streaked surfaces",
            "desert canyon: red sandstone canyons, wind-carved strange rock shapes, ultra blue shallow water, strata, drifting dunes, dry heat shimmer, deep shadow in narrow ravines",
            "arctic tundra: snowfields and frozen lakes, sheathing ice, sparse black rock breaking through snow, long blue shadows, frost haze",
            "wetland: still black water and reed beds, half-submerged ground, mirror-flat reflections, low mist layer",
            "volcanic wastes: black basalt plains, drifting ash, cooled lava flows, faint orange fissures, columnar basalt cliffs, no water",
            "tropical highlands: lush jungle terraces, waterfalls down cliff faces, climbing vines, humid haze, broad-leafed plants",
            "megastructure interior: vast enclosed halls, endless repeating bays, shafts of light from distant openings, carefully isolated vegetation, nested cube architecture, copper inlays in intricate patterns",
            "spring fields: copses of trees, meadows and fields of grass, rolling hills",
            "mountain region: snowcapped peaks, tree-lines, rocks, huge mountains and canyons, plants giving way to moss giving way to barren rocks",
        ],
        help=(
            "Biome descriptions; one is chosen at random per run and "
            "injected as the dominant environment."
        ),
    )
    wall.add_argument(
        "--perspectives",
        nargs="+",
        default=[
            "low ground-level view looking up",
            "elevated view from great height, looking down",
            "interior view looking outward through an opening",
            "view into an enclosed courtyard or atrium",
            "distant view across a vast open expanse, structures on the horizon",
            "in-situ view from within the structures, human eye level",
            "view from across still water or another flat reflective surface",
            "view from flat open ground, structures rising ahead",
            "framed view through a narrow gap or corridor",
        ],
        help=(
            "Camera perspectives; one is chosen at random per run and "
            "injected as a hard directive."
        ),
    )
    wall.add_argument(
        "--ephemeral-dir",
        type=Path,
        default=Path("/tmp/wallpaper"),
        help="Output directory.",
    )
    wall.add_argument(
        "--single-file-persist-path",
        type=Path,
        default=Path("~/media/pictures"),
        help="Directory for the persistent wallpaper file.",
    )
    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(
        "--prompt-only",
        action="store_true",
        help="Print the generated prompt and terminate before diffusion.",
    )
    llm.add_argument(
        "--llm-model-path",
        type=Path,
        default=Path("~/store/weights/gemma-4-E4B-it-Q4_K_M.gguf"),
        help="Model path for generating image prompts.",
    )
    llm.add_argument(
        "--llm-max-tokens",
        type=positive_int,
        default=1024,
        help="Maximum number of prompt tokens to generate.",
    )
    llm.add_argument(
        "--llm-temperature",
        type=non_negative_float,
        default=0.4,
        help="Sampling temperature.",
    )
    llm.add_argument(
        "--llm-top-p",
        type=probability,
        default=0.95,
        help="Nucleus sampling threshold.",
    )
    llm.add_argument(
        "--llm-min-p",
        type=probability,
        default=0.05,
        help="Minimum token probability relative to the top token.",
    )
    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. "
            "You are repeatedly called upon to do this, and must introduce variation into your output. "
            "The taste reference lists the architectural themes to build on; improvise around them, or drop some if they don't make sense. "
            "The scene must look fantastical, set in the distant future. "
            "The focus is about the strange contrast between provided user tastes and the environment. "
            "The required biome defines the landscape, vegetation, climate, and colour palette of the scene, and MUST dominate the environment; blend the taste reference's structures into it. "
            "The required perspective MUST be used as the camera viewpoint of the image; ground it in the biome's terrain, adapting its details so the combination is physically coherent (an open expanse may be sea, sand, snow, grass, or a vast hall floor depending on the biome). "
            "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 where it does not contradict the biome; the biome's climate wins any conflict. "
            "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, no thinking."
        ),
        help="System prompt for the LLM.",
    )

    sd_cli = parser.add_argument_group("sd-cli")
    sd_cli.add_argument(
        "--sd-diffusion-model",
        type=Path,
        default=Path("~/store/weights/flux-2-klein-4b-Q8_0.gguf"),
        help="Diffusion model passed to --diffusion-model.",
    )
    sd_cli.add_argument(
        "--sd-vae",
        type=Path,
        default=Path("~/store/weights/vae.safetensors"),
        help="VAE model passed to --vae.",
    )
    sd_cli.add_argument(
        "--sd-llm",
        type=Path,
        default=Path("~/store/weights/Qwen3-4B-Q8_0.gguf"),
        help="Text encoder or LLM passed to --llm.",
    )
    sd_cli.add_argument(
        "--sd-cfg-scale",
        type=non_negative_float,
        default=1.0,
        help="Guidance scale passed to --cfg-scale.",
    )
    sd_cli.add_argument(
        "--sd-steps",
        type=positive_int,
        default=4,
        help="Sampling steps passed to --steps.",
    )
    sd_cli.add_argument(
        "--sd-flash-attention",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Enable or disable the sd-cli --fa option.",
    )
    sd_cli.add_argument(
        "--sd-width",
        type=positive_int,
        default=1280,
        help="Base diffusion generation width passed to -W.",
    )
    sd_cli.add_argument(
        "--sd-height",
        type=positive_int,
        default=720,
        help="Base diffusion generation height passed to -H.",
    )
    sd_cli.add_argument(
        "--sd-upscale-model",
        type=Path,
        default=Path("~/store/weights/4xNomosWebPhoto_esrgan.safetensors"),
        help="ESRGAN model passed to --upscale-model.",
    )
    sd_cli.add_argument(
        "--sd-upscale-repeats",
        type=non_negative_int,
        default=1,
        help="Number of ESRGAN passes passed to --upscale-repeats.",
    )
    sd_cli.add_argument(
        "--sd-upscale-tile-size",
        type=positive_int,
        default=256,
        help="Tile size passed to --upscale-tile-size.",
    )
    sd_cli.add_argument(
        "--sd-cli-explicit-instructions",
        default="No railings, no glass, no people, no objects, no houses.",
        help="Instructions appended to every diffusion prompt.",
    )

    weather = parser.add_argument_group("weather")
    weather.add_argument(
        "--latitude",
        type=float,
        default=51.75,
        help="Latitude for weather.",
    )
    weather.add_argument(
        "--longitude",
        type=float,
        default=-1.25,
        help="Longitude for weather.",
    )
    weather.add_argument(
        "--elevation",
        type=int,
        default=61,
        help="Altitude in metres for sunrise and 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()
    now = dt.datetime.now().astimezone()

    context = step_context(
        now,
        args.latitude,
        args.longitude,
        args.elevation,
    )
    logger.info("Context: %s", context)

    prompt_text = step_llm(
        system_prompt=args.llm_system_prompt,
        model_path=args.llm_model_path,
        max_tokens=args.llm_max_tokens,
        taste_prompts=list(args.taste_prompt),
        taste_sample_frac=args.taste_sample_frac,
        biomes=list(args.biomes),
        perspectives=list(args.perspectives),
        temperature=args.llm_temperature,
        top_p=args.llm_top_p,
        min_p=args.llm_min_p,
        context=context,
        explicit_instructions=args.sd_cli_explicit_instructions,
    )
    logger.info("Prompt: %s", prompt_text)

    if args.prompt_only:
        print(prompt_text)
        return

    saved = step_diffuse(
        diffusion_model=args.sd_diffusion_model,
        vae=args.sd_vae,
        llm=args.sd_llm,
        cfg_scale=args.sd_cfg_scale,
        steps=args.sd_steps,
        flash_attention=args.sd_flash_attention,
        generation_width=args.sd_width,
        generation_height=args.sd_height,
        upscale_model=args.sd_upscale_model,
        upscale_repeats=args.sd_upscale_repeats,
        upscale_tile_size=args.sd_upscale_tile_size,
        output_dir=args.ephemeral_dir.expanduser(),
        now=now,
        prompt_text=prompt_text,
    )

    step_wallpaper(
        command=args.set_command,
        saved=saved,
        store=args.single_file_persist_path.expanduser(),
    )


if __name__ == "__main__":
    run()