"""Nidra — personalised bedtime story, narrated in your voice.

A two-step agent (NVIDIA Nemotron-Nano-9B: outline → full story) narrates via
VoxCPM2. All GPU work on Modal; this Space is CPU-only.

Hackathon: Thousand Token Wood · Agentic · NVIDIA Nemotron · OpenBMB.
"""
from __future__ import annotations
import io, re, random
import gradio as gr
import modal, numpy as np, soundfile as sf
from theme import build_css

MODAL_APP = "praxy-voice"
DEFAULT_REF = "ashwin_10s.wav"
_STORY = modal.Cls.from_name(MODAL_APP, "NemotronStoryteller")
_VOX   = modal.Cls.from_name(MODAL_APP, "VoxCPM2Cloner")

def _sentences(t): return [s.strip() for s in re.split(r"(?<=[.!?])\s+", t) if s.strip()]

def make_story(name, age, theme, ref_audio_path):
    if not (name and name.strip()) or not (theme and theme.strip()):
        return None, _empty_book(), _empty_outline(), None

    ref_path = ref_audio_path or DEFAULT_REF
    with open(ref_path, "rb") as f: ref_bytes = f.read()

    res = _STORY().write_story.remote(name=name, age=int(age), theme=theme)
    outline, story = res["outline"], res["story"]

    parts, sr = [], 48000
    for s in _sentences(story):
        wb, sr = _VOX().clone.remote(text=s, ref_audio_bytes=ref_bytes)
        arr, _ = sf.read(io.BytesIO(wb), dtype="float32")
        parts.append(arr); parts.append(np.zeros(int(sr*.3), dtype="float32"))
    full = np.concatenate(parts) if parts else np.zeros(1, dtype="float32")
    out_path = "/tmp/nidra_story.wav"
    sf.write(out_path, full, sr)

    trace_path = "/tmp/nidra_trace.md"
    with open(trace_path, "w") as f:
        f.write(f"# Nidra trace\n\n**For:** {name} (age {age})\n**Theme:** {theme}\n\n"
                f"## Agent step 1 — outline\n{outline}\n\n## Step 2 — story\n{story}\n")

    return out_path, _render_book(name, age, theme, story), _render_outline(outline), trace_path


def _empty_book():
    return ('<div style="padding:48px;text-align:center;color:#8086AE;'
            "font-family:'Fraunces',serif;font-size:18px;font-style:italic;line-height:1.8;\">"
            '✦ Your story will appear here, once the words are woven…</div>')

def _empty_outline():
    return ('<div style="padding:30px;text-align:center;color:#8086AE;'
            "font-family:'Plus Jakarta Sans',sans-serif;font-size:13px;\">"
            "The agent's outline will appear here.</div>")

def _render_book(name, age, theme, story):
    paras = [p.strip() for p in story.split("\n\n") if p.strip()]
    styled = []
    for i, p in enumerate(paras):
        if i == 0 and p:
            styled.append(
                f'<p style="margin:0 0 1.3em;">'
                f'<span style="font-family:\'Fraunces\',serif;font-size:4em;float:left;line-height:.78;'
                f'margin:6px 10px -6px 0;color:#F5C842;font-weight:700;">{p[0]}</span>{p[1:]}</p>')
        else:
            styled.append(f'<p style="margin:0 0 1.3em;">{p}</p>')
    body = "".join(styled)
    theme_short = theme[:64] + ("…" if len(theme) > 64 else "")
    return f"""
<div style="background:linear-gradient(165deg,#161C3F,#0B1026);border:1px solid rgba(245,200,66,.22);
  border-radius:20px;overflow:hidden;box-shadow:0 18px 50px rgba(0,0,0,.35);">
  <div style="background:rgba(245,200,66,.07);border-bottom:1px solid rgba(245,200,66,.16);
    padding:26px 34px;text-align:center;">
    <div style="font-size:11px;color:#C9A23A;letter-spacing:.16em;text-transform:uppercase;
      font-family:'Plus Jakarta Sans',sans-serif;font-weight:700;">A bedtime story for</div>
    <div style="font-family:'Fraunces',serif;font-size:32px;font-weight:700;color:#F5C842;margin:6px 0 2px;">{name}</div>
    <div style="font-size:12px;color:#AEB4D6;font-family:'Plus Jakarta Sans',sans-serif;">
      age {age} · <em style="color:#9AA0C6;">{theme_short}</em></div>
  </div>
  <div style="padding:38px 46px;font-family:'Fraunces',Georgia,serif;font-size:18px;
    line-height:2.0;color:#F4EFDE;">{body}</div>
  <div style="text-align:center;padding:20px;color:#C9A23A;font-size:13px;letter-spacing:.34em;
    border-top:1px solid rgba(245,200,66,.10);">✦&nbsp;&nbsp;THE END&nbsp;&nbsp;✦</div>
</div>"""

def _render_outline(outline):
    lines = [l.strip() for l in outline.strip().split("\n") if l.strip()]
    items = "".join(
        f'<div style="font-size:14px;color:#C7CBE6;line-height:1.65;padding:7px 0 7px 16px;'
        f'border-left:2px solid rgba(245,200,66,.35);margin-bottom:8px;">{l}</div>'
        for l in lines)
    return (
        '<div style="background:#141A3A;border:1px solid rgba(158,134,255,.18);border-radius:16px;padding:22px;">'
        '<div style="font-size:11px;font-weight:800;color:#9E86FF;letter-spacing:.14em;'
        "text-transform:uppercase;margin-bottom:14px;font-family:'Plus Jakarta Sans',sans-serif;\">"
        '◆ Agent outline · step 1</div>'
        f'{items}</div>')


# ── night-sky hero (80 seeded CSS stars + mountains + moon) ──────────────────

random.seed(42)
_STARS = "".join(
    f'<div style="position:absolute;top:{random.randint(2,92)}%;left:{random.randint(2,98)}%;'
    f'width:{(s:=random.choice([1,1,1,2,2,3]))}px;height:{s}px;border-radius:50%;'
    f'background:rgba(255,255,255,{round(random.uniform(.3,.95),2)});'
    f'animation:tw {round(random.uniform(1.4,3.5),1)}s {round(random.uniform(0,3),1)}s ease-in-out infinite alternate;"></div>'
    for _ in range(80))

_MOUNTAINS = ('<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1400 180" preserveAspectRatio="none"'
    ' style="position:absolute;bottom:0;left:0;width:100%;height:110px;pointer-events:none;">'
    '<polygon points="0,180 180,70 320,140 520,20 700,110 870,45 1060,130 1200,50 1400,120 1400,180"'
    ' fill="#070A1C" opacity="0.85"/>'
    '<polygon points="0,180 100,110 260,160 440,80 600,145 780,90 950,160 1100,75 1280,140 1400,100 1400,180"'
    ' fill="#070A1C"/></svg>')

HERO = f"""
<div style="background:linear-gradient(185deg,#04040F,#0B1026 45%,#120A30 72%,#070A1C);
  border-radius:26px;padding:72px 28px 120px;text-align:center;position:relative;overflow:hidden;
  border:1px solid rgba(245,200,66,.12);min-height:330px;">
  {_STARS}
  <div style="position:absolute;top:16%;left:9%;width:90px;height:1px;
    background:linear-gradient(90deg,transparent,#fff,transparent);transform:rotate(-20deg);
    animation:shoot 7s 2s ease-out infinite;"></div>
  {_MOUNTAINS}
  <div style="position:relative;z-index:10;max-width:640px;margin:0 auto;">
    <div style="font-size:72px;margin-bottom:6px;display:inline-block;
      animation:mf 5s ease-in-out infinite;filter:drop-shadow(0 0 30px rgba(255,210,80,.75));">🌙</div>
    <h1 style="font-family:'Fraunces',serif;font-size:clamp(2.6rem,8vw,4.6rem);font-weight:900;
      color:#F5C842;margin:0 0 12px;letter-spacing:-.02em;line-height:1;
      text-shadow:0 0 60px rgba(245,200,66,.45);">Nidra</h1>
    <p style="font-family:'Plus Jakarta Sans',sans-serif;font-size:clamp(1rem,2.6vw,1.22rem);
      color:#D7D9EC;margin:0 auto 18px;max-width:520px;line-height:1.6;">
      A bedtime story woven for your child — written by an AI agent,
      <strong style="color:#fff;">narrated in your own voice.</strong></p>
    <div style="display:flex;justify-content:center;gap:8px;flex-wrap:wrap;">
      <span class="pf-pill" style="background:rgba(158,134,255,.22);color:#C7B6FF;border:1px solid rgba(158,134,255,.4);">NVIDIA NEMOTRON</span>
      <span class="pf-pill" style="background:rgba(245,200,66,.16);color:#F5C842;border:1px solid rgba(245,200,66,.36);">VoxCPM2 · OPENBMB</span>
      <span class="pf-pill" style="background:rgba(96,216,144,.16);color:#7FE3A6;border:1px solid rgba(96,216,144,.34);">AGENTIC · 2-STEP</span>
    </div>
    <p style="font-size:12.5px;color:#8086AE;font-family:'Plus Jakarta Sans',sans-serif;margin:16px 0 0;">
      ⏳ First run warms the models on Modal (~3–5 min). After that it's quick.</p>
  </div>
</div>
"""

EXTRA = """
@keyframes tw {0%{opacity:.15;transform:scale(1)}100%{opacity:1;transform:scale(2.1)}}
@keyframes mf {0%,100%{transform:translateY(0) rotate(-5deg)}50%{transform:translateY(-10px) rotate(5deg)}}
@keyframes shoot {0%{opacity:0;transform:translate(0,0)}6%{opacity:1}80%{opacity:.5}100%{opacity:0;transform:translate(220px,108px)}}
"""

with gr.Blocks(title="Nidra — Bedtime Story Maker") as demo:
    gr.HTML(f"<style>{build_css('nidra', EXTRA)}</style>")
    gr.HTML(HERO)

    with gr.Row():
        name_in = gr.Textbox(label="Child's name", placeholder="Arjun", scale=2)
        age_in  = gr.Slider(2, 10, value=5, step=1, label="Age", scale=1)

    theme_in = gr.Textbox(
        label="Story theme",
        placeholder="a curious mongoose who finds a magical library under a banyan tree",
        lines=2)
    ref_in = gr.Audio(sources=["upload", "microphone"], type="filepath",
        label="Narrator's voice  (optional — leave empty for the default warm voice)")

    gen_btn = gr.Button("✨  Weave the story", variant="primary", elem_id="cta")
    audio_out = gr.Audio(label="Listen", type="filepath")
    gr.HTML('<hr class="pf-divider"/>')

    with gr.Row():
        story_out   = gr.HTML(value=_empty_book(), elem_id="story-col")
        outline_out = gr.HTML(value=_empty_outline(), elem_id="outline-col")
    trace_out = gr.File(label="Download the agent trace  (Markdown)")

    gen_btn.click(make_story, inputs=[name_in, age_in, theme_in, ref_in],
                  outputs=[audio_out, story_out, outline_out, trace_out])

if __name__ == "__main__":
    demo.launch()