# script_gen.py — AI Script Generation
#
# Primary:  Groq (Llama 4 Maverick) — 1,000 req/day free, no credit card
# Fallback: Google Gemini 2.0 Flash  — ~20 req/day free, no credit card

from groq import Groq
import google.genai as genai
from google.genai import types as genai_types


# ── Tone modifiers injected directly into system prompts ─────────────────────
# Tone must live in the SYSTEM prompt — not the user message — because the LLM
# treats system prompt instructions as its core identity. Tone in the user
# message gets overridden by the system prompt's established style.

TONE_SYSTEM_MODIFIER = {
    "Professional": (
        "TONE: Formal and authoritative. Use precise, clear language. "
        "No slang, no casual expressions. Every sentence should feel considered and expert."
    ),
    "Casual": (
        "TONE: Relaxed and conversational. Write like you're chatting with a friend. "
        "Use everyday expressions, contractions, and a warm, approachable vibe. "
        "Avoid stiff or formal phrasing."
    ),
    "Dramatic": (
        "TONE: Expressive and emotive. Build tension. Emphasise stakes. "
        "Use vivid, punchy language. Pause for effect. Make the audience feel something. "
        "Every sentence should have weight and urgency."
    ),
    "Humorous": (
        "TONE: Witty and playful. Weave in jokes, puns, and light-hearted observations "
        "based on the content. Keep it fun without being silly. "
        "Surprise the listener with unexpected but relevant humour. "
        "A well-placed joke should land every 3-4 exchanges."
    ),
    "Educational": (
        "TONE: Patient and instructive. Write as if explaining to someone encountering "
        "this topic for the first time. Define terms immediately. Use simple analogies. "
        "Build understanding step by step. Never assume prior knowledge."
    ),
}

TONE_USER_REMINDER = {
    "Professional": "Maintain formal, authoritative language throughout. No exceptions.",
    "Casual":       "Keep it conversational and friendly all the way through — contractions, warmth, zero stiffness.",
    "Dramatic":     "Stay dramatic — every line should carry tension, weight, or urgency. Don't let it go flat.",
    "Humorous":     "Stay playful and witty throughout. Include jokes or puns naturally woven into the content.",
    "Educational":  "Stay instructive and clear throughout — always explain, never assume.",
}


# ── Base system prompts (tone injected dynamically) ───────────────────────────

BASE_SYSTEM_PROMPTS = {
    "podcast": """{tone_modifier}

You are a professional podcast scriptwriter.
Write a natural, engaging 2-host podcast script using ONLY the provided source material.

Format EVERY spoken line strictly as one of:
ALEX: <what Alex says>
JAMIE: <what Jamie says>

Structure:
- ALEX opens with a hook and introduces the topic (2-3 lines)
- Both hosts discuss 3-4 key points from the material, alternating naturally
- Include moments of reaction, curiosity, and follow-up questions between hosts
- JAMIE delivers a clear summary of key takeaways (2-3 lines)
- ALEX closes with a sign-off (1-2 lines)

Rules:
- Natural spoken English — no bullet points in the dialogue
- No markdown, no headers, no asterisks in the output
- Every single line must start with ALEX: or JAMIE:
- Ground all content strictly in the source material provided
- Apply the tone above consistently to BOTH hosts' voices""",

    "debate": """{tone_modifier}

You are a professional debate scriptwriter.
Write a structured, passionate debate using ONLY the provided source material.

Format EVERY spoken line strictly as one of:
PRO: <argument>
CON: <counterargument>

Structure:
- PRO opening statement (2-3 lines)
- CON opening statement (2-3 lines)
- Round 1: PRO argument → CON rebuttal
- Round 2: PRO argument → CON rebuttal
- Round 3: PRO argument → CON rebuttal
- PRO closing statement
- CON closing statement

Rules:
- Each line should be 2-4 sentences long
- Arguments must be grounded in the source material
- Every line must start with PRO: or CON:
- Apply the tone above consistently — it should colour how each side argues""",

    "storytelling": """{tone_modifier}

You are a master storyteller and narrator.
Transform the source material into a compelling audio narrative.

Format EVERY line strictly as:
NARRATOR: <narration text>

Structure:
- Open with a vivid hook drawn directly from the source material
- Use analogies and relatable comparisons to explain concepts in the source
- Build a story arc using ONLY what is in the source: setup → complexity → insight
- Close with a reflection based strictly on what the source concludes

Rules:
- Vivid spoken English — no bullet points, no lists, no markdown
- Every line must start with NARRATOR:
- Do NOT invent background, history, or origin stories not in the source
- Apply the tone above to the narrator's voice and style throughout""",

    "lecture": """{tone_modifier}

You are an engaging university professor.
Deliver a clear, structured lecture using ONLY the provided source material.

Format EVERY spoken line strictly as:
PROFESSOR: <lecture content>

Structure:
- Open with a compelling hook or real-world question
- Introduce core concepts one by one with explanations and examples
- Use analogies to make abstract ideas concrete
- Summarise key learnings at the end
- Close with a thought-provoking question or next step

Rules:
- Explain jargon immediately when used
- No bullet points or lists in the speech itself
- Every line must start with PROFESSOR:
- Apply the tone above to how the professor speaks throughout""",

    "poem": """{tone_modifier}

You are a gifted spoken word poet.
Transform the source material into a beautiful, evocative poem meant to be read aloud.
Every image, idea, and emotion must come directly from the source — no invented content.

Format EVERY line strictly as:
NARRATOR: <poem line>

Structure:
- Opening stanza (4-6 lines): A vivid, arresting image or idea drawn from the source
- Body stanzas (3-4 stanzas of 4-6 lines each): Explore the key ideas poetically —
  use metaphor, rhythm, and imagery to illuminate what the source material means
- Closing stanza (4-6 lines): A resonant, memorable conclusion grounded in the source

Rules:
- Write in free verse or with a natural rhyme scheme — never forced rhymes
- Every line must start with NARRATOR:
- Use line breaks deliberately — short lines for emphasis, longer lines for flow
- Favour concrete images over abstract statements
- Do NOT invent facts, history, or context not present in the source
- The poem should feel like it was written by a human poet, not generated
- Apply the tone above to the voice and emotional register of the poem""",
}


def build_system_prompt(mode: str, tone: str) -> str:
    """Build a system prompt with tone baked directly into it."""
    base = BASE_SYSTEM_PROMPTS.get(mode, BASE_SYSTEM_PROMPTS["podcast"])
    modifier = TONE_SYSTEM_MODIFIER.get(tone, TONE_SYSTEM_MODIFIER["Professional"])
    return base.format(tone_modifier=modifier)


USER_PROMPT_TEMPLATE = """Create a complete, production-ready {mode} script based STRICTLY on the source material below.

Duration target: Approximately {duration} minute(s) of spoken audio.
Word count target: Aim for approximately {target_words} words total.
(Spoken audio averages ~130 words per minute)

TONE REMINDER: {tone_reminder}

Critical rules:
- Do NOT add any facts, claims, or details not explicitly present in the source
- Keep all content grounded and accurate to the source
- Write in natural spoken English — this will be converted to audio
- Hit the word count target as closely as possible
- Follow the format exactly as instructed (every line must have the correct speaker label)

SOURCE MATERIAL:
{context}

Write the complete script now."""


# ── Provider functions ────────────────────────────────────────────────────────

def generate_with_groq(context: str, mode: str, api_key: str,
                       tone: str = "Professional", duration: float = 2.0) -> str:
    client = Groq(api_key=api_key)

    target_words  = int(duration * 130)
    max_tokens    = min(4096, max(1024, int(target_words * 1.6)))
    tone_reminder = TONE_USER_REMINDER.get(tone, TONE_USER_REMINDER["Professional"])
    system_prompt = build_system_prompt(mode, tone)

    response = client.chat.completions.create(
        model="meta-llama/llama-4-maverick-17b-128e-instruct",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": USER_PROMPT_TEMPLATE.format(
                mode=mode, context=context,
                duration=duration, target_words=target_words,
                tone_reminder=tone_reminder,
            )},
        ],
        max_tokens=max_tokens,
        temperature=0.75,  # slightly higher for tone variety
    )

    result = response.choices[0].message.content
    if not result or not result.strip():
        raise ValueError("Groq returned an empty response.")
    return result


def generate_with_gemini(context: str, mode: str, api_key: str,
                         tone: str = "Professional", duration: float = 2.0) -> str:
    client = genai.Client(api_key=api_key)

    target_words  = int(duration * 130)
    max_tokens    = min(4096, max(1024, int(target_words * 1.6)))
    tone_reminder = TONE_USER_REMINDER.get(tone, TONE_USER_REMINDER["Professional"])
    system_prompt = build_system_prompt(mode, tone)

    full_prompt = f"{system_prompt}\n\n{USER_PROMPT_TEMPLATE.format(mode=mode, context=context, duration=duration, target_words=target_words, tone_reminder=tone_reminder)}"

    response = client.models.generate_content(
        model="gemini-2.0-flash",
        contents=full_prompt,
        config=genai_types.GenerateContentConfig(
            max_output_tokens=max_tokens,
            temperature=0.75,
        ),
    )

    result = response.text
    if not result or not result.strip():
        raise ValueError("Gemini returned an empty response.")
    return result


# ── Main entry point ──────────────────────────────────────────────────────────

def generate_script(
    context: str,
    mode: str,
    groq_key: str = "",
    gemini_key: str = "",
    tone: str = "Professional",
    duration: float = 2.0,
) -> str:
    errors = []

    if groq_key.strip():
        try:
            return generate_with_groq(context, mode, groq_key.strip(), tone, duration)
        except Exception as e:
            errors.append(f"Groq error: {str(e)}")
            print(f"[VoiceVerse] Groq failed, trying Gemini. Reason: {e}")

    if gemini_key.strip():
        try:
            return generate_with_gemini(context, mode, gemini_key.strip(), tone, duration)
        except Exception as e:
            errors.append(f"Gemini error: {str(e)}")
            print(f"[VoiceVerse] Gemini also failed. Reason: {e}")

    raise ValueError(
        "Script generation failed.\n\n"
        + "\n".join(errors)
        + "\n\nPlease check that your API keys are correctly set in HF Space Secrets."
    )