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
    import sys

    # Suppress Python 3.10 asyncio GC bug: BaseEventLoop.__del__ crashes with
    # "Invalid file descriptor: -1" when a loop is collected by the GC.
    _orig_unraisable = sys.unraisablehook
    def _unraisable_hook(args):
        if args.exc_type is ValueError and "Invalid file descriptor" in str(args.exc_value):
            return  # silently drop — harmless Python 3.10 GC bug
        _orig_unraisable(args)
    sys.unraisablehook = _unraisable_hook

    print("HF Space detected — model will load on first GPU call and stay cached in memory.")

# ---------------------------------------------------------------------------
# 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:
    # Include ALL selected animals (up to 3) — pick one place randomly
    animal_list  = animals[:3] if animals else [random.choice(list(ANIMAL_MAP))]
    place_emoji  = random.choice(places) if places else random.choice(list(PLACE_MAP))

    animal_names = [ANIMAL_MAP.get(a, "bunny") for a in animal_list]
    if len(animal_names) == 1:
        animal_text = animal_names[0]
    elif len(animal_names) == 2:
        animal_text = f"{animal_names[0]} and {animal_names[1]}"
    else:
        animal_text = f"{animal_names[0]}, {animal_names[1]} and {animal_names[2]}"

    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 globally — survives between ZeroGPU calls)
# ---------------------------------------------------------------------------

_pipe = None

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

    from diffusers import FluxPipeline

    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    print(f"Loading FLUX.1-schnell pipeline… (token={'set' if hf_token else 'NOT SET'})")
    _pipe = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-schnell",
        torch_dtype=torch.bfloat16,
        token=hf_token,
    )

    if IS_HF_SPACE:
        # ZeroGPU always has CUDA — don't rely on torch.cuda.is_available()
        # which can return False outside the @spaces.GPU context
        _pipe = _pipe.to("cuda")
    elif torch.cuda.is_available():
        _pipe = _pipe.to("cuda")
    elif torch.backends.mps.is_available():
        _pipe = _pipe.to("mps")
    else:
        _pipe = _pipe.to("cpu")

    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:
    # duration=120: covers first-run model download (~30s) + load (~10s) + generate (~5s).
    # On subsequent calls _pipe is already loaded so only ~5s of GPU time is used.
    @spaces.GPU(duration=120)
    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)