import random
import os
import torch

# Patch for torch < 2.4 which lacks torch.xpu (required by diffusers >= 0.30)
if not hasattr(torch, "xpu"):
    class _MockXPU:
        is_available    = staticmethod(lambda: False)
        device_count    = staticmethod(lambda: 0)
        empty_cache     = staticmethod(lambda: None)
        manual_seed     = staticmethod(lambda seed: None)
        reset_peak_memory_stats = staticmethod(lambda: None)
        max_memory_allocated    = staticmethod(lambda: 0)
        synchronize     = staticmethod(lambda: None)
    torch.xpu = _MockXPU()

# ---------------------------------------------------------------------------
# Detect HuggingFace Spaces ZeroGPU
# ---------------------------------------------------------------------------

IS_HF_SPACE = os.environ.get("SPACE_ID") is not None
if IS_HF_SPACE:
    import spaces

# ---------------------------------------------------------------------------
# Emoji → descriptive text maps (fed into FLUX prompt)
# ---------------------------------------------------------------------------

ANIMAL_MAP: dict[str, str] = {
    "🐶": "puppy",        "🐱": "kitten",       "🐰": "bunny",
    "🦊": "baby fox",     "🐼": "baby panda",   "🐨": "baby koala",
    "🦁": "baby lion",    "🐯": "baby tiger",   "🐸": "baby frog",
    "🐧": "baby penguin", "🦋": "butterfly",    "🦄": "unicorn",
}

PLACE_MAP: dict[str, str] = {
    "🌊": "on a sunny beach with ocean waves",
    "🏔️": "on a snowy mountain top",
    "🌸": "in a cherry blossom garden",
    "🌈": "under a rainbow",
    "🌙": "on a glowing crescent moon",
    "⭐": "surrounded by sparkling stars",
    "🌴": "on a tropical island",
    "🏡": "in a cosy cottage garden",
    "🌺": "in a field of tropical flowers",
    "🍄": "in an enchanted mushroom forest",
}

STYLES = [
    "kawaii digital art",
    "soft watercolor illustration",
    "cute cartoon style",
    "pastel chibi illustration",
]

# ---------------------------------------------------------------------------
# Prompt builder
# ---------------------------------------------------------------------------

def build_prompt(streak: int, animals: list[str], places: list[str]) -> str:
    # Pick one animal and one place from user selection
    animal_emoji = random.choice(animals) if animals else random.choice(list(ANIMAL_MAP))
    place_emoji  = random.choice(places)  if places  else random.choice(list(PLACE_MAP))

    animal_text = ANIMAL_MAP.get(animal_emoji, "bunny")
    place_text  = PLACE_MAP.get(place_emoji,   "in a magical garden")
    style       = random.choice(STYLES)

    if streak <= 2:
        mood   = "cute"
        extras = "with big sparkling eyes, soft pastel colors"
    elif streak <= 5:
        mood   = "super cute and happy"
        extras = "with big sparkling eyes, glitter, pastel rainbow colors, smiling"
    else:
        mood   = "magically adorable, ultra fluffy"
        extras = (
            "with big sparkling eyes, magical sparkles, rainbow aura, "
            "tiny crown, pastel colors, looking amazed"
        )

    prompt = (
        f"A {mood} {animal_text} {place_text}, {style}, {extras}, "
        "white background, children's book illustration style, high quality, detailed"
    )
    return prompt

# ---------------------------------------------------------------------------
# Pipeline loader (cached)
# ---------------------------------------------------------------------------

_pipe = None

def get_pipeline():
    global _pipe
    if _pipe is not None:
        return _pipe

    from diffusers import FluxPipeline

    print("Loading FLUX.1-schnell pipeline…")
    _pipe = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-schnell",
        torch_dtype=torch.bfloat16,
    )

    if torch.cuda.is_available():
        _pipe = _pipe.to("cuda")
    elif torch.backends.mps.is_available():
        _pipe = _pipe.to("mps")
    else:
        _pipe.enable_model_cpu_offload()

    print("Pipeline ready.")
    return _pipe

# ---------------------------------------------------------------------------
# Core generation (always wrapped in try/except)
# ---------------------------------------------------------------------------

def _generate(streak: int, animals: list[str], places: list[str]):
    try:
        pipe   = get_pipeline()
        prompt = build_prompt(streak, animals, places)
        print(f"Generating | streak={streak} | prompt: {prompt}")

        result = pipe(
            prompt=prompt,
            num_inference_steps=4,
            guidance_scale=0.0,
            height=512,
            width=512,
        )
        return result.images[0], prompt

    except Exception as e:
        import traceback
        print(f"[image_generator] ❌ generation failed: {e}")
        print(traceback.format_exc())
        return None, str(e)


# ---------------------------------------------------------------------------
# Public API — two versions depending on environment
# ---------------------------------------------------------------------------

if IS_HF_SPACE:
    @spaces.GPU(duration=60)
    def generate_reward_image(streak: int, animals: list[str], places: list[str]):
        """Generate reward image on HF Spaces ZeroGPU."""
        return _generate(streak, animals, places)
else:
    def generate_reward_image(streak: int, animals: list[str], places: list[str]):
        """Generate reward image locally."""
        return _generate(streak, animals, places)