"""
DoodleDreams — ZeroGPU orchestrator
Draw + voice → illustrated bedtime storybook narrated in your cloned voice.
"""
import os, sys, json, time, tempfile, logging, threading
import torch

sys.path.insert(0, os.path.dirname(__file__))

try:
    import spaces
except ModuleNotFoundError:
    class _Shim:
        @staticmethod
        def GPU(*a, **k):
            return a[0] if a and callable(a[0]) else (lambda fn: fn)
    spaces = _Shim()

import gradio as gr
from config import (
    FLUX_MODEL, STORY_MODEL, TTS_MODEL, TRANSLATION_MODEL,
    KANNADA_TTS_MODEL, BASE_SEED, FLUX_STEPS, FLUX_GUIDANCE, FLUX_SIZE,
    COLOR_ART_STYLE, COLOR_PAGE_SUFFIX, STORY_LENGTHS, GENRES, MOODS,
)
from book_builder import build_book_html, export_pdf, magic_loader_html
from ui.layout import create_layout

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))

_FLUX       = None
_STORY_M    = None; _STORY_TOK = None
_TTS_M      = None
_TRANS_M    = None; _TRANS_TOK = None
_KAN_TTS_M  = None
_LOAD_ERRORS = {}


def _load_flux():
    global _FLUX
    if _FLUX is None:
        from diffusers import Flux2KleinPipeline
        _FLUX = Flux2KleinPipeline.from_pretrained(
            FLUX_MODEL.hub_id, torch_dtype=torch.bfloat16).to("cuda")
    return _FLUX


def _load_story():
    global _STORY_M, _STORY_TOK
    if _STORY_M is None:
        from transformers import AutoTokenizer, AutoModelForCausalLM
        _STORY_TOK = AutoTokenizer.from_pretrained(
            STORY_MODEL.hub_id, trust_remote_code=True)
        _STORY_M = AutoModelForCausalLM.from_pretrained(
            STORY_MODEL.hub_id, torch_dtype=torch.float16, trust_remote_code=True,
        ).to("cuda").eval()
    return _STORY_M, _STORY_TOK


def _load_tts():
    global _TTS_M
    if _TTS_M is None:
        from voxcpm import VoxCPM
        _TTS_M = VoxCPM.from_pretrained(TTS_MODEL.hub_id, load_denoiser=False)
    return _TTS_M


def _load_translation():
    global _TRANS_M, _TRANS_TOK
    if _TRANS_M is None:
        from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
        _TRANS_TOK = AutoTokenizer.from_pretrained(
            TRANSLATION_MODEL.hub_id, trust_remote_code=True)
        _TRANS_M = AutoModelForSeq2SeqLM.from_pretrained(
            TRANSLATION_MODEL.hub_id, trust_remote_code=True,
        ).to("cuda").eval()
    return _TRANS_M, _TRANS_TOK


def _load_kannada_tts():
    global _KAN_TTS_M
    if _KAN_TTS_M is None:
        from indic_tts import _get_model
        _KAN_TTS_M = _get_model()
    return _KAN_TTS_M


if ON_ZEROGPU:
    for _n, _fn in [("flux", _load_flux), ("story", _load_story),
                    ("tts", _load_tts), ("translation", _load_translation),
                    ("kannada_tts", _load_kannada_tts)]:
        try:
            _fn()
        except Exception as e:
            _LOAD_ERRORS[_n] = repr(e)
            logger.exception(f"Module-level load failed: {_n}")


# ── Genre/mood → deterministic story templates (fallback) ──────────────

_TEMPLATES = {
    "Animals": [
        ("{hero} loved exploring the meadow every evening.", "{hero} walking through a golden meadow at dusk"),
        ("One night, {hero} heard a tiny sound in the tall grass.", "{hero} listening carefully near the rustling grass"),
        ("A little firefly needed help finding its family.", "{hero} meeting a tiny glowing firefly"),
        ("{hero} gently carried the firefly through the dark forest.", "{hero} walking carefully through a moonlit forest"),
        ("Together they found the firefly's home, glowing warm and bright.", "{hero} and the firefly reunited with the glowing firefly family"),
        ("Tired and happy, {hero} curled up under the stars.", "{hero} sleeping peacefully under a starry sky"),
    ],
    "Kingdom": [
        ("In a cozy kingdom, {hero} was the kindest helper of all.", "{hero} standing cheerfully in a small fairy-tale kingdom"),
        ("One sleepy evening, the king's golden crown went missing.", "{hero} seeing the worried king without his crown"),
        ("{hero} searched the royal garden by moonlight.", "{hero} searching carefully through a moonlit garden"),
        ("A sleepy mouse had borrowed it for a bed!", "{hero} discovering a tiny mouse asleep inside the crown"),
        ("{hero} found the mouse a proper bed made of petals.", "{hero} tucking the tiny mouse into a flower-petal bed"),
        ("The king smiled, and the whole kingdom slept in peace.", "{hero} and the king smiling together under the night sky"),
    ],
    "Space": [
        ("{hero} loved watching the stars from the garden.", "{hero} lying in the grass gazing at the starry sky"),
        ("One night, a small star blinked and fell from the sky.", "{hero} seeing a little star tumbling down"),
        ("{hero} caught the star in a jar of moonlight.", "{hero} gently catching a glowing star in a jar"),
        ("The star was lost and didn't know how to get home.", "{hero} listening to the sad little star"),
        ("{hero} climbed the tallest hill and let the star go free.", "{hero} releasing the star from the hilltop into the sky"),
        ("The star zoomed home, and {hero} fell fast asleep smiling.", "{hero} smiling and drifting off to sleep under the stars"),
    ],
    "Dragons": [
        ("{hero} lived near a mountain where a shy dragon slept.", "{hero} looking up at a misty mountain at night"),
        ("One evening, the dragon sneezed and lost its flame.", "{hero} watching the dragon sneeze sadly"),
        ("{hero} brought warm soup and a soft blanket to the dragon.", "{hero} carrying a steaming bowl of soup to the dragon"),
        ("The dragon felt better and puffed a tiny grateful flame.", "{hero} and the dragon sharing a warm cozy moment"),
        ("Together they lit the lanterns along the sleepy village path.", "{hero} and the dragon lighting lanterns in the quiet village"),
        ("The dragon curled up, and {hero} tucked it in with a smile.", "{hero} tucking the dragon in for the night"),
    ],
    "Ocean": [
        ("{hero} sat by the shore watching the moonlight on the waves.", "{hero} sitting peacefully by the ocean at night"),
        ("A little fish splashed up and looked worried.", "{hero} seeing a small worried fish near the surface"),
        ("The fish had lost its shell-home in a big wave.", "{hero} listening to the little fish explain its problem"),
        ("{hero} dove gently and found the shell on the sandy floor.", "{hero} swimming carefully along the moonlit ocean floor"),
        ("The fish swam home, and the sea became calm and quiet.", "{hero} watching the happy fish return to its shell"),
        ("{hero} fell asleep to the soft sound of the waves.", "{hero} sleeping peacefully beside the calm, moonlit sea"),
    ],
    "Forest": [
        ("{hero} walked into the whispering forest as the moon rose.", "{hero} stepping into a moonlit forest path"),
        ("The trees were worried — an owl had lost its song.", "{hero} hearing the trees whisper about the silent owl"),
        ("{hero} climbed a mossy rock and hummed a gentle tune.", "{hero} humming softly on a mossy rock under the moon"),
        ("The owl listened and slowly remembered its melody.", "{hero} watching the owl open its eyes and begin to sing"),
        ("The whole forest filled with soft nighttime music.", "{hero} smiling as the forest glows with peaceful sound"),
        ("{hero} yawned and drifted off to sleep among the roots.", "{hero} sleeping curled up peacefully at the base of a great tree"),
    ],
}

FEW_SHOT = """
Write a 6-page children's bedtime story for age 5 about Luna the cat. Genre: Animals. Mood: Calming.

Return ONLY valid JSON:
{
  "title": "Luna and the Sleepy Firefly",
  "character_description": "A small grey cat named Luna with soft fur, big green eyes, and a white tip on her tail",
  "pages": [
    {"page": 1, "text": "Luna loved sitting in the garden when the moon came out.", "scene": "Luna sitting in a moonlit garden"},
    {"page": 2, "text": "One night, she heard a tiny buzzing sound in the flowers.", "scene": "Luna listening near a flower patch"},
    {"page": 3, "text": "A little firefly was lost and couldn't find its family.", "scene": "Luna meeting a tiny glowing firefly"},
    {"page": 4, "text": "Luna walked gently through the dark, lighting the way.", "scene": "Luna walking with the firefly glowing beside her"},
    {"page": 5, "text": "They found the firefly's home, glowing warm and bright.", "scene": "Luna and firefly arriving at a cluster of glowing lights"},
    {"page": 6, "text": "Luna purred softly and curled up under the stars.", "scene": "Luna sleeping peacefully under a starry sky"}
  ]
}
"""


def _build_story_locally(hero_name: str, genre: str) -> dict:
    hero = (hero_name or "Little Hero").strip() or "Little Hero"
    beats = _TEMPLATES.get(genre, _TEMPLATES["Animals"])
    pages = [
        {"page": i+1, "text": t.format(hero=hero), "scene": s.format(hero=hero)}
        for i, (t, s) in enumerate(beats)
    ]
    return {
        "title": f"{hero}'s Bedtime Dream",
        "character_description": (
            f"{hero}, a friendly children's storybook hero with bright colors, "
            "bold outlines, and a cheerful expressive face"
        ),
        "pages": pages,
    }


def _parse_story_json(raw: str) -> dict | None:
    import re
    m = re.search(r'\{[\s\S]*\}', raw or "")
    if not m:
        return None
    try:
        d = json.loads(m.group(0))
        if "pages" in d and "title" in d:
            return d
    except Exception:
        pass
    return None


# ── ZeroGPU inference ──────────────────────────────────────────────────

@spaces.GPU(duration=60)
def _gen_story_gpu(hero_name: str, genre: str, mood: str) -> dict:
    try:
        model, tok = _load_story()
        prompt = (
            f"{FEW_SHOT}\n\n"
            f"Write a 6-page children's bedtime story for age 5 about {hero_name}. "
            f"Genre: {genre}. Mood: {mood}. Keep it gentle and sleepy.\n\n"
            f"Return ONLY valid JSON:\n"
        )
        inputs = tok.apply_chat_template(
            [{"role": "user", "content": prompt}],
            add_generation_prompt=True, enable_thinking=False,
            return_dict=True, return_tensors="pt",
        ).to("cuda")
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=800, do_sample=False)
        response = tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        parsed = _parse_story_json(response)
        if parsed:
            return parsed
    except Exception as e:
        logger.warning(f"Story GPU failed: {e}")
    return _build_story_locally(hero_name, genre)


@spaces.GPU(duration=150)
def _gen_images_gpu(char_desc: str, scenes: list,
                    doodle_bytes: bytes | None, seed: int) -> list:
    import io
    from PIL import Image
    pipe = _load_flux()
    canonical = None
    if doodle_bytes:
        try:
            ref = Image.open(io.BytesIO(doodle_bytes)).convert("RGB")
            canonical = pipe(
                prompt=(
                    "Turn this child's drawing into a clean, full-body cartoon character "
                    "for a children's storybook. Keep the EXACT same creature. "
                    f"{COLOR_ART_STYLE}, plain white background, full character visible, centered."
                ),
                image=ref, height=FLUX_SIZE, width=FLUX_SIZE,
                guidance_scale=FLUX_GUIDANCE, num_inference_steps=FLUX_STEPS,
                generator=torch.Generator("cuda").manual_seed(seed),
            ).images[0]
        except Exception as e:
            logger.warning(f"Canonical pass failed ({e}); text2img fallback")
    images = []
    for i, scene in enumerate(scenes):
        if canonical is not None:
            kw = dict(image=canonical,
                      prompt=f"The same character. {scene}. {COLOR_ART_STYLE}, {COLOR_PAGE_SUFFIX}")
        else:
            kw = dict(prompt=f"{char_desc}. Scene: {scene}. {COLOR_ART_STYLE}, centered.")
        kw.update(height=FLUX_SIZE, width=FLUX_SIZE, guidance_scale=FLUX_GUIDANCE,
                  num_inference_steps=FLUX_STEPS,
                  generator=torch.Generator("cuda").manual_seed(seed + i + 1))
        images.append(pipe(**kw).images[0])
        logger.info(f"Page {i+1}/{len(scenes)} illustrated")
    return images


@spaces.GPU(duration=120)
def _gen_tts_gpu(text: str, ref_wav: str | None,
                 mood: str, energy: float, language: str) -> str:
    if language == "Kannada":
        from indic_text import translate_to_kannada
        from indic_tts import narrate_kannada
        kannada_text = translate_to_kannada(text)
        ref_txt = "ಇದು ನನ್ನ ಧ್ವನಿ"
        return narrate_kannada(ref_wav or "", ref_txt, kannada_text, mood, energy)
    else:
        from tts import clone_and_speak
        return clone_and_speak(
            ref_wav=ref_wav,
            text=text,
            speed=0.9,
            mood=mood.lower(),
            energy=energy,
        )


# ── heartbeat helper ───────────────────────────────────────────────────

def _with_heartbeat(blocking_fn, frame_fn, poll=4.0):
    box = {}
    def _run():
        try:   box["val"] = blocking_fn()
        except BaseException as e: box["err"] = e
    th = threading.Thread(target=_run, daemon=True)
    th.start()
    t0 = time.time()
    while th.is_alive():
        th.join(timeout=poll)
        if th.is_alive():
            yield ("hb", frame_fn(int(time.time() - t0)))
    if "err" in box:
        raise box["err"]
    yield ("done", box["val"])


# ── main generator ─────────────────────────────────────────────────────

def create_book(doodle_image, ref_audio, hero_name, genre, mood, language, length_label):
    t0 = time.perf_counter()
    hero_name = (hero_name or "").strip() or "Little Hero"
    energy    = 0.45

    trace = {
        "backend": "zerogpu", "hero": hero_name,
        "genre": genre, "mood": mood, "language": language,
        "seed": BASE_SEED, "ts": time.strftime("%Y-%m-%d %H:%M:%S"),
    }
    if _LOAD_ERRORS:
        trace["load_errors"] = _LOAD_ERRORS

    _no   = gr.update(visible=False)
    _keep = gr.update()

    yield (magic_loader_html("story", hero_name),
           "Writing the bedtime story…", None, _no, {}, "")

    try:
        story = _gen_story_gpu(hero_name, genre, mood)
    except Exception as e:
        yield (f"<div class='page-loading'>Error: {e}</div>",
               f"Error: {e}", None, _no, {}, "")
        return

    title      = story.get("title", "A Bedtime Story")
    pages      = story.get("pages", [])
    char_desc  = story.get("character_description", "")
    scenes     = [p.get("scene", "") for p in pages]
    page_texts = [p.get("text", "")  for p in pages]
    full_text  = f"{title}. {' '.join(page_texts)}"
    trace.update(title=title, char_desc=char_desc)

    yield (magic_loader_html("images", hero_name),
           f"{title} — illustrating…", None, _no, story, json.dumps(trace, indent=2))

    doodle_bytes = None
    if doodle_image is not None:
        import io
        from PIL import Image
        buf = io.BytesIO()
        Image.fromarray(doodle_image).save(buf, format="PNG")
        doodle_bytes = buf.getvalue()

    # narration starts in parallel with illustration
    voice_box = {}
    def _do_voice():
        try:
            voice_box["path"] = _gen_tts_gpu(full_text, ref_audio, mood, energy, language)
        except Exception as e:
            voice_box["err"] = e

    voice_th = threading.Thread(target=_do_voice, daemon=True)
    voice_th.start()

    def _audio_now():
        return voice_box.get("path")

    img_bytes, engine = None, "sketch"
    try:
        for kind, payload in _with_heartbeat(
            lambda: _gen_images_gpu(char_desc, scenes, doodle_bytes, BASE_SEED),
            lambda s: (
                magic_loader_html("images", hero_name),
                f"{title} — illustrating… {s}s"
                + ("  · narration ready ▶" if _audio_now() else "  · recording…"),
                _audio_now(), _no, story, json.dumps(trace, indent=2),
            ),
        ):
            if kind == "hb":
                yield payload
            else:
                import io
                img_bytes = []
                for img in payload:
                    buf = io.BytesIO(); img.save(buf, format="PNG")
                    img_bytes.append(buf.getvalue())
                engine = "flux"
    except Exception as e:
        logger.exception("Image generation failed")
        trace["image_error"] = repr(e)
        from services.images import generate_placeholder_images
        img_bytes = generate_placeholder_images(char_desc, scenes, doodle_bytes)

    book_html = build_book_html(img_bytes, page_texts, title, engine)

    while voice_th.is_alive():
        voice_th.join(timeout=4)
        if voice_th.is_alive():
            yield (book_html, f"{title} — finishing narration…",
                   _audio_now(), _no, story, json.dumps(trace, indent=2))
    audio_path = _audio_now()
    if voice_box.get("err"):
        trace["tts_error"] = repr(voice_box["err"])

    pdf_path = None
    try:
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
            pdf_path = export_pdf(img_bytes, page_texts, title, f.name)
    except Exception as e:
        logger.warning(f"PDF failed: {e}")

    trace["total_sec"] = round(time.perf_counter() - t0, 2)
    trace["engine"] = engine

    pdf_update = gr.update(value=pdf_path, visible=True) if pdf_path else _keep

    yield (
        book_html,
        f"Done: {title} · {len(img_bytes)} pages · {language} · {trace['total_sec']}s",
        audio_path, pdf_update, story, json.dumps(trace, indent=2),
    )


if __name__ == "__main__":
    demo = create_layout(create_book_fn=create_book)
    demo.queue(default_concurrency_limit=2, max_size=8)
    demo.launch(share=False, allowed_paths=[tempfile.gettempdir()],
                **demo.design_kwargs)