script_generator.py

"""
script_generator.py

LLM: openai/gpt-oss-20b via HuggingFace InferenceClient (chat completions)
Fallback: high-quality template engine
"""

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

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


# ---------------------------------------------------------------------------
# InferenceClient call  (chat completions format)
# ---------------------------------------------------------------------------

def _call_llm(system_prompt: str, user_prompt: str) -> Optional[str]:
    if not HF_TOKEN:
        print("[ScriptGen] No HF_TOKEN — skipping LLM, using template.")
        return None

    import time
    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=1200,
                temperature=0.72,
                top_p=0.92,
            )
            text = response.choices[0].message.content.strip()
            print(f"[ScriptGen] {MODEL} returned {len(text)} chars")

            if len(text) > 200:
                return text

            print(f"[ScriptGen] Too short — model cold starting, 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 exhausted — using template.")
    return None


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

SYSTEM_PROMPT = """You are a professional podcast scriptwriter with years of experience.
Your scripts sound completely natural when read aloud — like real human conversation, not AI text.
You never copy source material verbatim. You always rewrite ideas in a warm, engaging, spoken style.
You never use bullet points, headers, markdown, or stage directions in your output."""


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

    words  = int(duration_target * 130)
    focus  = f"Focus specifically on: {custom_focus.strip()}." if custom_focus.strip() else ""

    format_map = {
        "Podcast (Single Host)":
            "a single-host podcast episode. One narrator speaks throughout — warm, curious, engaging.",
        "Podcast (Dual Host Debate)":
            "a two-host podcast. Label EVERY line as either HOST_A: or HOST_B: — no exceptions. "
            "The hosts have a real conversation: they ask each other questions, build on each other's points, "
            "sometimes disagree, and bring in examples.",
        "Educational Narration":
            "a clear educational narration. One voice explains the topic step by step, simply and clearly.",
        "Storytelling Episode":
            "a narrative storytelling episode. Vivid, immersive, draws the listener in with a story arc.",
        "News Bulletin":
            "a professional news bulletin. Crisp, factual, authoritative anchor tone.",
        "Rap / Song Mode":
            "a rap song. Use Verse 1, Chorus, Verse 2, Outro structure. Rhythmic, punchy lines.",
        "Executive Summary":
            "a 90-second executive audio briefing. Dense, direct, every word earns its place.",
    }
    fmt = format_map.get(content_type, "a spoken audio piece")

    cleaned_context = _pre_clean_context(context)

    return f"""Write {fmt}

Topic: {doc_name}
Tone: {tone}
Target length: approximately {words} words
{focus}

STRICT RULES:
- Output ONLY the spoken script. No titles, headers, labels, stage directions, or markdown.
- Every sentence must be fluent, complete, and natural to speak aloud.
- Use contractions, rhetorical questions, and natural spoken transitions.
- Structure: compelling hook → substantive body (3 main points) → memorable closing.
- Base ALL facts only on the source material below. Do not invent anything not in the source.
- Do NOT copy sentences from the source — rewrite everything in your own spoken voice.

SOURCE MATERIAL:
\"\"\"
{cleaned_context[:4000]}
\"\"\""""


def _pre_clean_context(context: str) -> str:
    """Strip leftover PDF junk before sending to LLM."""
    lines = context.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
# ---------------------------------------------------------------------------

def _extract_good_sentences(context: str, n: int = 15) -> list:
    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|"
                r"^\d+\.|ibid|op\. cit|et al\.",
                s, re.IGNORECASE
            )
            and len(re.findall(r"[A-Za-z]", s)) > 30
        ):
            good.append(s)
    seen, unique = set(), []
    for s in good:
        key = s[:50]
        if key not in seen:
            seen.add(key)
            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 topic"

    sents = _extract_good_sentences(context, n=15)
    if len(sents) < 4:
        raise ValueError(
            "Not enough clean sentences to build a script. "
            "The PDF may be scanned or heavily formatted. Try a different document."
        )

    intro     = sents[:2]
    body      = sents[2:11]
    outro     = sents[11:13]

    openers = {
        "Engaging & Conversational":
            f"Have you ever really thought about {topic}? It's one of those subjects that seems straightforward at first — but the deeper you go, the more there is to unpack. Let's get into it.",
        "Professional & Formal":
            f"Today we examine {topic} — a subject with significant practical implications for professionals in this field.",
        "Playful & Fun":
            f"Okay, buckle up — we're diving into {topic} today, and I promise it's way more interesting than it sounds!",
        "Dramatic & Intense":
            f"What if everything you assumed about {topic} was only half the story? Today, we fill in the gaps.",
        "Calm & Reflective":
            f"Let's take a moment — slow down, breathe — and really explore {topic}. It deserves more attention than we usually give it.",
    }
    opener  = openers.get(tone, openers["Engaging & Conversational"])
    closing = (
        f"So here's what it all comes down to when we talk about {topic}: "
        f"{'  '.join(outro) if outro else 'the ideas we explored today are genuinely worth applying.'} "
        f"Thank you for listening. Until next time."
    )

    transitions = [
        "Let's start with the foundation.",
        "Now here's where it gets really interesting.",
        "And this is the piece most people overlook.",
    ]
    body_paras = []
    chunk = max(1, len(body) // 3)
    for i in range(3):
        group = body[i * chunk: (i + 1) * chunk]
        if group:
            body_paras.append(f"{transitions[i]} {' '.join(group)}")

    if content_type == "Rap / Song Mode":
        v1 = " / ".join(body[:4])
        v2 = " / ".join(body[4:8])
        return (
            f"Verse one.\n{v1}\n\n"
            f"Chorus.\nKnowledge drops, truth unlocks — {topic} around the clock.\n\n"
            f"Verse two.\n{v2}\n\n"
            f"Outro.\n{closing}"
        )
    if content_type == "News Bulletin":
        return (
            f"Good evening. Tonight's lead story: {topic}. "
            f"{' '.join(intro)} "
            f"{' '.join(body_paras)} "
            f"Reporting complete."
        )
    if num_hosts == 2:
        return "\n\n".join([
            f"HOST_A: {opener}",
            f"HOST_B: Exactly. And what really struck me is that {intro[0] if intro else ''}",
            f"HOST_A: {body_paras[0] if body_paras else ''}",
            f"HOST_B: That's such a key point. {body_paras[1] if len(body_paras) > 1 else ''}",
            f"HOST_A: And building on that — {body_paras[2] if len(body_paras) > 2 else ''}",
            f"HOST_B: Right. The real-world takeaway here is about applying this consistently.",
            f"HOST_A: {closing}",
            f"HOST_B: Couldn't agree more. See you next time.",
        ])

    return f"{opener}\n\n{chr(10).join(body_paras)}\n\n{closing}"


# ---------------------------------------------------------------------------
# Post-process
# ---------------------------------------------------------------------------

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"\[(?!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.IGNORECASE
        )
    ]
    script = "\n".join(lines)
    script = re.sub(r"\n{3,}", "\n\n", 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:

        # Try LLM
        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-generated script.")
            return _clean_script(script)

        # Template fallback
        print("[ScriptGen] Using template engine.")
        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)