script_generator.py

"""
script_generator.py

LLM  : openai/gpt-oss-20b via HuggingFace InferenceClient
Output: deeply structured scripts with compelling intro, rich body, insightful conclusion
"""

from __future__ import annotations
import os
import re
import time
from typing import Optional

HF_TOKEN = os.environ.get("HF_TOKEN", "")
MODEL    = "openai/gpt-oss-20b"


# ---------------------------------------------------------------------------
# LLM call with retry
# ---------------------------------------------------------------------------

def _call_llm(system_prompt: str, user_prompt: str) -> Optional[str]:
    if not HF_TOKEN:
        print("[ScriptGen] No HF_TOKEN — using template.")
        return None
    from huggingface_hub import InferenceClient
    client = InferenceClient(token=HF_TOKEN, model=MODEL)
    for attempt in range(1, 4):
        try:
            print(f"[ScriptGen] {MODEL} attempt {attempt}/3...")
            response = client.chat_completion(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user",   "content": user_prompt},
                ],
                max_tokens=1400,
                temperature=0.72,
                top_p=0.92,
            )
            text = response.choices[0].message.content.strip()
            print(f"[ScriptGen] Returned {len(text)} chars")
            if len(text) > 200:
                return text
            print(f"[ScriptGen] Too short, waiting 8s...")
            time.sleep(8)
        except Exception as e:
            print(f"[ScriptGen] Attempt {attempt} failed: {e}")
            if attempt < 3:
                time.sleep(5)
    print("[ScriptGen] All retries failed — using template.")
    return None


# ---------------------------------------------------------------------------
# System prompt  (the character the LLM inhabits)
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """\
You are a world-class audio content producer — think NPR meets TED Talks.
Your scripts are praised for three things:
1. A HOOK that immediately pulls the listener in with a surprising fact, provocative question, or vivid scenario.
2. A BODY that builds understanding through concrete examples, contrasts, and 'aha' moments — never just lists of facts.
3. A CONCLUSION that doesn't just summarise — it gives the listener something to CARRY with them: a reframing, a call to action, or an insight they'll repeat to others.

You write ONLY the spoken script — no titles, headers, markdown, stage directions, or meta-commentary.
You never copy source text verbatim. You always rewrite in your own warm, intelligent, spoken voice.
Every sentence must sound natural when read aloud."""


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

FORMAT_MAP = {
    "Podcast (Single Host)":
        "a single-host podcast episode. One curious, intelligent narrator takes the listener on a journey.",
    "Podcast (Dual Host Debate)":
        "a two-host podcast. Label every line HOST_A: or HOST_B:. "
        "HOST_A introduces ideas; HOST_B challenges, probes, and adds real-world examples. "
        "They build on each other — it feels like a real conversation, not a script.",
    "Educational Narration":
        "a clear educational narration. Explain the topic from first principles. "
        "Use analogies and concrete examples. One calm, authoritative voice.",
    "Storytelling Episode":
        "an immersive narrative episode. Open with a scene or character. "
        "Weave the key ideas into the story arc. End with a revelation.",
    "News Bulletin":
        "a professional news bulletin. Lead with the most important fact. "
        "Crisp sentences. Anchor tone — authoritative and clear.",
    "Rap / Song Mode":
        "a rap. Structure: [Verse 1], [Chorus], [Verse 2], [Bridge], [Outro]. "
        "Punchy rhyming lines. Keep the wordplay tight. dont tag the structure, just the rap song",
    "Executive Summary":
        "a 90-second executive audio briefing. Every word earns its place. "
        "State the situation, the insight, and the implication.",
}

STRUCTURE_GUIDE = {
    "Podcast (Single Host)": """\
STRUCTURE:
- INTRO (15-20% of length): Open with a hook — a surprising fact, bold question, or vivid scenario drawn from the material. Don't say "Welcome to VoiceVerse" — just dive in.
- BODY (65-70%): Cover 3 distinct angles or insights. Each one should: (a) state the idea clearly, (b) ground it with a specific example or detail from the source, (c) explain WHY it matters. Use transitions like "But here's what makes this fascinating..." or "And this is where most people miss the point..."
- CONCLUSION (15%): Don't summarise — reframe. Give the listener one sharp insight they'll carry with them. End with energy, not a fade.""",

    "Podcast (Dual Host Debate)": """\
STRUCTURE:
- INTRO: HOST_A opens with the hook. HOST_B immediately reacts with curiosity or mild challenge.
- BODY: 3 rounds of exchange. Each round = HOST_A raises a point from the material → HOST_B probes with "But doesn't that assume...?" or adds a counterexample → HOST_A deepens or concedes → natural pivot to next point.
- CONCLUSION: HOST_B asks "So what does this actually mean for someone listening right now?" HOST_A delivers the key insight. HOST_B adds one final thought.""",

    "Educational Narration": """\
STRUCTURE:
- INTRO: Frame the problem or question this topic answers. Why should anyone care?
- BODY: Teach 3 concepts in logical order. Each concept: define it simply → show it in action with an example → connect it to the next concept.
- CONCLUSION: The 'so what' — how does understanding this change how the listener thinks or acts?""",

    "Storytelling Episode": """\
STRUCTURE:
- OPENING SCENE: Drop the listener into a moment. A person, a place, a decision.
- RISING TENSION: Introduce the central tension or question that the document addresses.
- INSIGHT REVEAL: The key ideas from the material emerge through the story.
- RESOLUTION: How the story resolves — and what it means for the listener.""",

    "News Bulletin": """\
STRUCTURE:
- HEADLINE: The single most important fact, stated as a crisp declarative sentence.
- CONTEXT: Why this matters and what led here (2-3 sentences).
- DETAILS: The key specifics — who, what, how.
- SIGNIFICANCE: What this means going forward.""",

    "Rap / Song Mode": """\
STRUCTURE:
- [Verse 1]: Introduce the topic and core tension. 8 lines.
- [Chorus]: The memorable hook — 4 lines that capture the essence. Repeat twice.
- [Verse 2]: Go deeper — add nuance, examples, a twist. 8 lines.
- [Bridge]: A moment of reflection or contrast. 4 lines.
- [Outro]: The call to action or final insight. 4 lines.""",

    "Executive Summary": """\
STRUCTURE:
- SITUATION (1 sentence): What is the context?
- COMPLICATION (1-2 sentences): What is the tension, challenge, or opportunity?
- INSIGHT (2-3 sentences): What does the material reveal that's non-obvious?
- IMPLICATION (1-2 sentences): What should the listener think, decide, or do differently?""",
}


def _build_user_prompt(context: str, content_type: str, tone: str,
                       num_hosts: int, duration_target: float,
                       doc_name: str, custom_focus: str) -> str:

    # At natural speech: ~130 wpm. Speed slider affects TTS, not word count.
    # We scale word count directly with duration so script IS the right length.
    words  = int(duration_target * 130)
    focus  = f"\nFocus specifically on: {custom_focus.strip()}." if custom_focus.strip() else ""
    fmt    = FORMAT_MAP.get(content_type, "a spoken audio piece")
    struct = STRUCTURE_GUIDE.get(content_type, STRUCTURE_GUIDE["Podcast (Single Host)"])
    ctx    = _pre_clean_context(context)

    return f"""Write {fmt}

Topic: {doc_name}
Tone: {tone}
TARGET LENGTH: {words} words (= {duration_target:.1f} minutes of audio at natural pace).
  - If you are running short: add more examples, analogies, elaboration, and deeper explanation.
  - If you are running long: cut ruthlessly — tighten sentences, remove repetition.
  - Count your output. {words} words is the target. Do not stop until you reach it.{focus}

{struct}

QUALITY RULES:
- The intro must NOT start with "Welcome" or "Today we're going to talk about". Start mid-thought.
- Every body point must include a SPECIFIC detail or example from the source material.
- The conclusion must deliver a genuine insight — not a summary. Something the listener will remember.
- Use natural spoken rhythms: short punchy sentences mixed with longer flowing ones.
- All facts must come ONLY from the source below. Do not invent anything.

SOURCE MATERIAL:
\"\"\"
{ctx[:4500]}
\"\"\""""


def _clean_for_speech(text: str) -> str:
    """
    Fix common PDF-extraction artefacts so text reads cleanly aloud:
    - Rejoin soft-hyphenated line-breaks  (e.g. "con-\ntent" → "content")
    - Collapse runs of whitespace / stray newlines inside a paragraph
    - Fix missing space after punctuation
    - Remove stray Unicode ligatures and control chars
    - Normalise quotes/dashes
    """
    # 1. Rejoin hyphenated line-breaks (word-\nfragment → wordfragment)
    text = re.sub(r"-\s*\n\s*([a-z])", r"\1", text)
    # 2. Collapse mid-sentence newlines (not paragraph breaks)
    text = re.sub(r"(?<!\n)\n(?!\n)", " ", text)
    # 3. Collapse multiple spaces
    text = re.sub(r"[ \t]{2,}", " ", text)
    # 4. Fix missing space after sentence-ending punctuation when followed by a capital
    text = re.sub(r"([.!?])([A-Z])", r"\1 \2", text)
    # 5. Fix missing space after comma/semicolon
    text = re.sub(r"([,;:])([^\s\d])", r"\1 \2", text)
    # 6. Remove stray bullet/list chars that read poorly
    text = re.sub(r"^\s*[-•·]\s*", "", text, flags=re.MULTILINE)
    # 7. Normalise dashes → comma-space for smoother TTS
    text = re.sub(r"\s*—\s*", ", ", text)
    text = re.sub(r"\s*–\s*", ", ", text)
    # 8. Normalise quotes
    text = re.sub(r"[""„]", '"', text)
    text = re.sub(r"[''‛]", "'", text)
    # 9. Drop page-number artifacts, e.g. lines that are purely digits or "Page N"
    text = re.sub(r"(?m)^\s*(Page\s+)?\d+\s*$", "", text)
    # 10. Final whitespace tidy
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def _pre_clean_context(ctx: str) -> str:
    lines = ctx.splitlines()
    good  = []
    for line in lines:
        line = line.strip()
        if not line:
            good.append("")
            continue
        if re.search(
            r"authorized for use|no part of|hbs no\.|s p jain|spjimr|"
            r"^\d+\s*$|^page\s+\d+|from dec \d{4}|to jun \d{4}|"
            r"all rights reserved|transmitted in any form",
            line, re.IGNORECASE
        ):
            continue
        if len(line) < 20 and not line.endswith((".", "!", "?")):
            continue
        good.append(line)
    return "\n".join(good).strip()


# ---------------------------------------------------------------------------
# Template fallback  (used only when LLM fails)
# ---------------------------------------------------------------------------

def _extract_good_sentences(context: str, n: int = 18) -> list:
    context = _clean_for_speech(context)
    raw  = re.split(r"(?<=[.!?])\s+", context)
    good = []
    for s in raw:
        s = s.strip()
        if (len(s) > 55 and s[0].isupper() and s[-1] in ".!?"
                and not re.search(r"hbs|authorized|spjimr|©|fig\.|table \d|^\d+\.|ibid|et al\.", s, re.I)
                and len(re.findall(r"[A-Za-z]", s)) > 30):
            good.append(s)
    seen, unique = set(), []
    for s in good:
        if s[:50] not in seen:
            seen.add(s[:50])
            unique.append(s)
    return unique[:n]


def _template_script(context: str, content_type: str, tone: str,
                     num_hosts: int, doc_name: str,
                     custom_focus: str, duration_target: float) -> str:

    topic = (custom_focus.strip() if custom_focus.strip()
             else re.sub(r"^\d+[\.\d]*\s*", "", doc_name).replace("_", " ").replace("-", " ").strip())
    topic = topic or "this subject"

    sents = _extract_good_sentences(context, n=18)
    if len(sents) < 4:
        raise ValueError(
            "Not enough clean text to generate a script. "
            "The PDF may be scanned or heavily formatted."
        )

    hook    = sents[0]
    body    = sents[1:13]
    insight = sents[13:16]
    outro   = sents[16:18]

    # Hook openers by tone
    hooks = {
        "Engaging & Conversational":
            f"Here's something most people get wrong about {topic}: they think they already understand it. {hook}",
        "Professional & Formal":
            f"The evidence on {topic} points to a conclusion that has significant implications. {hook}",
        "Playful & Fun":
            f"Okay, real talk — {topic} is far more interesting than anyone gave it credit for. {hook}",
        "Dramatic & Intense":
            f"What if the conventional wisdom on {topic} is backwards? {hook}",
        "Calm & Reflective":
            f"There's something worth sitting with when we examine {topic} carefully. {hook}",
    }
    opener = hooks.get(tone, hooks["Engaging & Conversational"])

    # Body: 3 structured points
    chunk = max(1, len(body) // 3)
    point_intros = [
        "The first thing to understand is this.",
        "But here's where it gets genuinely interesting.",
        "And this third dimension is what most people overlook entirely.",
    ]
    body_parts = []
    for i in range(3):
        group = body[i * chunk:(i + 1) * chunk]
        if group:
            body_parts.append(f"{point_intros[i]} {' '.join(group)}")

    # Conclusion: reframing insight
    insight_text = " ".join(insight) if insight else ""
    outro_text   = " ".join(outro) if outro else ""
    conclusion = (
        f"So here's the reframe: when it comes to {topic}, the real question isn't just what it is — "
        f"it's what it changes about how you act. {insight_text} "
        f"{'That is the part worth carrying forward.' if not outro_text else outro_text}"
    )

    if content_type == "Rap / Song Mode":
        v1 = " / ".join(body[:4])
        v2 = " / ".join(body[4:8])
        return (f"[Verse 1]\n{v1}\n\n"
                f"[Chorus]\n{topic}, open your eyes / the truth is right here, no disguise\n\n"
                f"[Verse 2]\n{v2}\n\n"
                f"[Bridge]\n{insight_text}\n\n"
                f"[Outro]\n{conclusion}")

    if content_type == "News Bulletin":
        return (f"{body_parts[0] if body_parts else hook} "
                f"{body_parts[1] if len(body_parts) > 1 else ''} "
                f"The significance: {conclusion}")

    if num_hosts == 2:
        return "\n\n".join([
            f"HOST_A: {opener}",
            f"HOST_B: That framing is really striking. Because what it implies is — {body_parts[0] if body_parts else ''}",
            f"HOST_A: Exactly. And if you follow that logic through — {body_parts[1] if len(body_parts) > 1 else ''}",
            f"HOST_B: But doesn't that raise the question: what does this mean practically? Like, for someone trying to apply this?",
            f"HOST_A: {body_parts[2] if len(body_parts) > 2 else ''}",
            f"HOST_B: So the real takeaway here isn't just the concept — it's the implication.",
            f"HOST_A: {conclusion}",
            f"HOST_B: That's a genuinely useful way to think about it. Thanks for listening, everyone.",
        ])

    return f"{opener}\n\n{chr(10).join(body_parts)}\n\n{conclusion}"


# ---------------------------------------------------------------------------
# Clean LLM output
# ---------------------------------------------------------------------------

def _clean_script(script: str) -> str:
    script = re.sub(r"\*{1,3}(.*?)\*{1,3}", r"\1", script)
    script = re.sub(r"#{1,4}\s*", "", script)
    script = re.sub(r"\[(?!Verse|Chorus|Bridge|Outro|INST)[^\]]{0,40}\]", "", script)
    script = re.sub(r"https?://\S+", "", script)
    lines  = [l for l in script.splitlines()
              if not re.search(r"authorized for use|hbs no\.|spjimr|s p jain|©|all rights", l, re.I)]
    script = "\n".join(lines)
    script = re.sub(r"\n{3,}", "\n\n", script)
    script = _clean_for_speech(script)
    return script.strip()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

class ScriptGenerator:

    def generate(self, context: str, content_type: str, tone: str,
                 num_hosts: int, duration_target: float,
                 doc_name: str, custom_focus: str) -> str:

        script = _call_llm(
            system_prompt=SYSTEM_PROMPT,
            user_prompt=_build_user_prompt(
                context, content_type, tone, num_hosts,
                duration_target, doc_name, custom_focus
            ),
        )

        if script and len(script.strip()) > 200:
            print("[ScriptGen] Using LLM script.")
            return _clean_script(script)

        print("[ScriptGen] Using template.")
        script = _template_script(
            context=context, content_type=content_type, tone=tone,
            num_hosts=num_hosts, doc_name=doc_name,
            custom_focus=custom_focus, duration_target=duration_target,
        )
        return _clean_script(script)