LLMVideoPipeline/pipeline/script_gen.py

"""
Script Generation Module using AMDLLM (llama.cpp Qwen3.5-7B).
Generates bilingual documentary scripts, A/B test titles, SEO metadata,
dynamic chapter markers, visual generation prompts, and stock search queries.
"""

import json
import logging
import re
from typing import Dict, Any, List, Optional
from .remote import RemotePC
from .research import ViralTopicResearcher

logger = logging.getLogger("YTFactory.Script")

class ScriptGen:
    def __init__(self, pcs: Dict[str, RemotePC], llm_port: int = 8002):
        self.pcs = pcs
        self.llm_port = llm_port
        self.researcher = ViralTopicResearcher(pcs.get("amdllm"), llm_port)

    def _build_master_prompt(self, topic: str, duration: int, research_dossier: Optional[Dict[str, Any]] = None) -> str:
        # Pacing: ~140 words per minute for professional documentary delivery.
        num_segments = max(4, duration // 20)
        target_words = max(70, int(duration * 140 / 60))
        words_per_seg = max(25, target_words // num_segments)

        dossier = research_dossier or {}
        viral_hook = dossier.get("viral_hook_angle", "The unexpected mystery behind this subject.")
        misconceptions = ", ".join(dossier.get("common_misconceptions", [])) or "Standard historical misconceptions"
        revelations = ", ".join(dossier.get("key_revelations", [])) or "Groundbreaking discoveries"
        visual_style = dossier.get("visual_style_anchor", "cinematic documentary style, warm archival tones, 35mm film texture, photorealistic 8k")

        return f"""You are a master YouTube retention strategist and award-winning documentary showrunner (modeling 10M+ view explainer channels like Veritasium, Kurzgesagt, ColdFusion, and Wendover).
Generate an authoritative, high-retention faceless explainer documentary script for the topic: "{topic}".

COMPETITIVE VIRAL INTELLIGENCE & RESEARCH:
- Primary Curiosity Gap Hook: "{viral_hook}"
- Critical Myths to Debunk: {misconceptions}
- Breakthrough Revelations & Step-by-Step Mechanisms: {revelations}
- Unified Visual Art Direction: "{visual_style}"

SPECIFICATIONS:
- Target Duration: Exactly ~{duration} seconds of continuous narration
- Total Word Count: ~{target_words} spoken words across all segments (~{words_per_seg} words per segment)
- Number of Segments: Exactly {num_segments} chronological segments
- Tone: Mysterious, authoritative, engaging, cinematic, insightful
- Narrative Flow (5-Act Explainer Blueprint):
  * Act 1 (Opening): Irresistible hook & psychological curiosity gap
  * Act 2 (The Illusion): The common myth or historical puzzle
  * Act 3 (The Turning Point): The catalyst or critical breakthrough
  * Act 4 (The Mechanism): Step-by-step breakdown of how it actually happened
  * Act 5 (The Paradigm Shift): Profound modern conclusion and lasting impact
- Spoken Depth: Each segment MUST contain a full, rich spoken paragraph ({words_per_seg - 5} to {words_per_seg + 10} words) so the narration fills the entire video duration without silence or rushing.

For EACH segment provide:
1. "chapter_title": Short 2-4 word chapter title (e.g., "The Ancient Mystery", "The Critical Discovery")
2. "text": English spoken narration (rich 2-4 sentence narrative paragraph, natural pacing, compelling storytelling)
3. "visual_prompt": Highly detailed visual scene description for an AI video/image generator (cinematic lighting, camera angle, {visual_style})
4. "stock_query": 2-4 concrete, high-precision visual search keywords describing physical objects/scenery that exist in stock video libraries (e.g., "ancient temple monoliths", "medieval cathedral altar", "archaeological excavation desert", "prehistoric stone tools", "ancient scrolls manuscript"). Never use abstract concepts, full sentences, or obscure person names.
5. "on_screen_text": Optional 2-5 word lower-third label, key statistic, date, or scene title to display on screen (e.g., "Depth: 11,034m", "Year: 1960", "Pacific Ocean Abyss", or "" if none)
6. "is_historical": boolean (true if this segment depicts past events, ancient civilizations, historical figures, archaeological discoveries, or vintage eras; false if modern/contemporary)
7. "person_name": Name of any specific historical person, scientist, inventor, philosopher, leader, or notable figure discussed in this scene (e.g., "Alan Turing", "Steve Wozniak", "Charles Babbage", "Jack the Ripper", or "" if none)
8. "duration_estimate": Estimated duration in seconds ({max(10, duration // num_segments)}s)

METADATA REQUIREMENTS:
- "titles": Exactly 3 distinct viral, high-CTR YouTube titles for A/B testing (max 85 chars each)
- "hook": Compelling 5-second opening question or statement
- "description": Comprehensive 3-paragraph YouTube description with engaging summary, key takeaways, and relevant hashtags (#Documentary #History #Science)
- "tags": 12-18 highly relevant YouTube search tags

STRICT OUTPUT FORMAT:
Output MUST be a single, valid JSON object starting with {{ and ending with }}. Do NOT include markdown code blocks (```), comments, or text before/after.

JSON Schema:
{{
  "titles": ["Title A", "Title B", "Title C"],
  "hook": "Opening hook line...",
  "description": "SEO description...",
  "tags": ["tag1", "tag2", "tag3"],
  "script_segments": [
    {{
      "chapter_title": "Chapter Name",
      "text": "English spoken narration for segment 1...",
      "visual_prompt": "Cinematic visual description...",
      "stock_query": "relevant broll keyword",
      "on_screen_text": "Key Metric or Title",
      "is_historical": true,
      "duration_estimate": 18
    }}
  ]
}}"""

    def _extract_json(self, text: str) -> Dict[str, Any]:
        """Sanitize and parse JSON from LLM response with multiple fallback passes."""
        text = text.strip()
        text = re.sub(r'^\s*```(?:json)?\s*', '', text, flags=re.IGNORECASE)
        text = re.sub(r'\s*```\s*$', '', text)
        text = text.strip()

        start = text.find('{')
        end = text.rfind('}')
        if start == -1 or end <= start:
            raise ValueError("No JSON object found in response")

        json_str = text[start:end + 1]
        # Clean trailing commas
        json_str = re.sub(r',\s*}', '}', json_str)
        json_str = re.sub(r',\s*]', ']', json_str)

        try:
            return json.loads(json_str)
        except json.JSONDecodeError:
            pass

        # Try json5 or partial bracket matching if available
        try:
            import json5
            return json5.loads(json_str)
        except Exception:
            pass

        for end_idx in range(end, start, -1):
            if text[end_idx] == '}':
                candidate = text[start:end_idx + 1]
                candidate = re.sub(r',\s*}', '}', candidate)
                candidate = re.sub(r',\s*]', ']', candidate)
                try:
                    return json.loads(candidate)
                except json.JSONDecodeError:
                    continue

        raise ValueError("Unable to parse JSON from LLM response after all sanitization passes.")

    def _chat(self, pc: RemotePC, prompt: str, max_tokens: int = 8192) -> Optional[str]:
        payload = {
            "model": "llama",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.7,
            "max_tokens": max_tokens
        }
        r = pc.post_json(self.llm_port, "/v1/chat/completions", payload, timeout=600)
        if not r or "choices" not in r:
            return None
        return r["choices"][0]["message"]["content"]

    def _translate_to_spanish(self, pc: RemotePC, script_data: Dict[str, Any]) -> Dict[str, Any]:
        """Translate narration, titles, chapters, description, and on-screen text into Latin American Spanish."""
        logger.info("[SCRIPT] Translating script to Latin-American Spanish...")
        
        segments = script_data.get("script_segments", [])
        titles = script_data.get("titles", [])
        desc = script_data.get("description", "")
        
        translation_payload = {
            "titles": titles,
            "narrations": [s.get("text", "") for s in segments],
            "chapters": [s.get("chapter_title", f"Part {i+1}") for i, s in enumerate(segments)],
            "on_screen_texts": [s.get("on_screen_text", "") for s in segments],
            "description": desc
        }

        prompt = f"""You are a master Spanish documentary voiceover narrator and localization expert (National Geographic / Discovery Channel en Español).
Translate the following English YouTube documentary package into captivating, natural, spoken Latin-American Spanish for voiceover narration and YouTube publishing.

English Data:
{json.dumps(translation_payload, ensure_ascii=False, indent=2)}

CRITICAL SPOKEN VOICE REQUIREMENTS:
1. Spoken Eloquence & Natural Flow:
   - Write as if spoken by a native Latin-American documentary orator with rich, gripping cadence.
   - Use natural spoken Spanish sentence connectors ("Sin embargo,", "A lo largo de los siglos,", "Lo más fascinante es que...", "No obstante,").
   - Avoid rigid word-for-word literal translations from English.
2. Punctuation for Breathing & Timing:
   - Use commas and question marks (¿...?) thoughtfully so the text-to-speech engine pauses with natural human breathing.
3. Proper Nouns & Accuracy:
   - Keep recognized proper names accurate ("Whitechapel", "Londres", "Alejandría", "Göbeklitepe").
4. "titles_es": 3 high-CTR Spanish titles
5. "narrations_es": array of exact same length containing the rich Spanish voiceover text
6. "chapters_es": array of exact same length containing short 2-4 word Spanish chapter titles
7. "on_screen_texts_es": array of exact same length containing Spanish lower-third labels
8. "description_es": full Spanish description with Spanish hashtags

STRICT OUTPUT FORMAT:
Output STRICT JSON only: {{"titles_es": [...], "narrations_es": [...], "chapters_es": [...], "on_screen_texts_es": [...], "description_es": "..."}}"""

        for attempt in (1, 2):
            text = self._chat(pc, prompt, max_tokens=4096)
            if not text:
                continue
            try:
                es_data = self._extract_json(text)
                narr_es = es_data.get("narrations_es", [])
                titles_es = es_data.get("titles_es", [])
                chapters_es = es_data.get("chapters_es", [])
                ost_es = es_data.get("on_screen_texts_es", [])
                desc_es = es_data.get("description_es", "")

                if len(narr_es) == len(segments):
                    for i, seg in enumerate(segments):
                        seg["text_es"] = narr_es[i]
                        seg["chapter_title_es"] = chapters_es[i] if i < len(chapters_es) else seg.get("chapter_title", "")
                        seg["on_screen_text_es"] = ost_es[i] if i < len(ost_es) else seg.get("on_screen_text", "")
                    
                    script_data["titles_es"] = titles_es if len(titles_es) >= len(titles) else titles
                    script_data["description_es"] = desc_es or desc
                    logger.info("[SCRIPT] Spanish translation completed successfully.")
                    return script_data
                else:
                    logger.warning(f"[SCRIPT] Segment length mismatch on ES translation ({len(narr_es)} vs {len(segments)})")
            except Exception as e:
                logger.warning(f"[SCRIPT] ES translation parse attempt {attempt} failed: {e}")

        logger.warning("[SCRIPT] Spanish translation failed after retries. Using English fallbacks.")
        for seg in segments:
            seg.setdefault("text_es", seg.get("text", ""))
            seg.setdefault("chapter_title_es", seg.get("chapter_title", ""))
            seg.setdefault("on_screen_text_es", seg.get("on_screen_text", ""))
        script_data["titles_es"] = script_data.get("titles", [])
        script_data["description_es"] = script_data.get("description", "")
        return script_data

    def _proofread_and_refine_script(self, pc: RemotePC, script_data: Dict[str, Any], topic: str) -> Dict[str, Any]:
        """
        Secondary LLM pass to proofread, fact-check, and fix proper noun/phonetic errors
        (e.g., 'white tecapel' -> 'Whitechapel', 'tutankamen' -> 'Tutankhamun').
        """
        logger.info("[SCRIPT:AUDIT] Running secondary LLM proofreading & fact-checking pass...")
        
        prompt = f"""You are a master documentary script editor, fact-checker, and phonetics auditor.
Your job is to rigorously review and correct this draft documentary script for the topic: "{topic}".

AUDITING OBJECTIVES:
1. Proper Nouns & Historical Places:
   - Identify and correct ANY misspellings, phonetic corruptions, or garbled proper nouns (for example, if you see "white tecapel" or "white te capel", change it to "Whitechapel"; fix any names of people, cities, landmarks, vessels, historical dates, or scientific terms).
2. Phonetic Cleanliness for Text-to-Speech:
   - Ensure words are written cleanly so speech synthesizers pronounce them naturally without awkward syllable breaks.
3. Flow & Grammar:
   - Enhance sentence rhythm, natural documentary pacing, and dramatic intrigue.
4. Structure:
   - Preserve the exact same JSON schema, segment count, keys, visual prompts, and metadata.

DRAFT SCRIPT TO AUDIT:
{json.dumps(script_data, ensure_ascii=False, indent=2)}

STRICT OUTPUT FORMAT:
Return ONLY the corrected script as a single valid JSON object starting with {{ and ending with }}. No markdown, no commentary."""

        try:
            raw_text = self._chat(pc, prompt, max_tokens=8192)
            if raw_text:
                refined = self._extract_json(raw_text)
                if "script_segments" in refined and len(refined["script_segments"]) == len(script_data.get("script_segments", [])):
                    logger.info("[SCRIPT:AUDIT] Secondary LLM proofreading pass passed successfully.")
                    return refined
                else:
                    logger.warning("[SCRIPT:AUDIT] Segment count mismatch in audit pass. Keeping original.")
        except Exception as e:
            logger.warning(f"[SCRIPT:AUDIT] Secondary LLM audit pass encountered error: {e}. Keeping original.")

        return script_data

    def _proofread_spanish_translation(self, pc: RemotePC, script_data: Dict[str, Any], topic: str) -> Dict[str, Any]:
        """
        Secondary LLM pass to proofread Spanish translation for natural Latin-American idioms,
        proper noun spelling (e.g. 'Whitechapel', 'Alejandría'), accent marks, and pronunciation.
        """
        logger.info("[SCRIPT:AUDIT] Running secondary LLM Spanish audit pass...")
        
        segments = script_data.get("script_segments", [])
        payload = {
            "titles_es": script_data.get("titles_es", []),
            "description_es": script_data.get("description_es", ""),
            "segments_es": [
                {
                    "chapter_title_es": s.get("chapter_title_es", ""),
                    "text_es": s.get("text_es", ""),
                    "on_screen_text_es": s.get("on_screen_text_es", "")
                }
                for s in segments
            ]
        }

        prompt = f"""You are a senior Spanish documentary voiceover director and phonetic editor.
Review, polish, and perfect this Spanish documentary narration for topic "{topic}".

OBJECTIVES FOR 100% NATURAL SPOKEN SPANISH:
1. Native Cadence & Phrasing:
   - Eliminate any awkward literal translations or robotic sentence structures.
   - Refine sentence flow so every line sounds natural, authoritative, and compelling when spoken aloud.
2. Proper Punctuation & Breathing Pauses:
   - Ensure natural comma placement to give the voice synthesizer natural breathing pauses.
3. Proper Noun & Pronunciation Accuracy:
   - Ensure historical figures, locations, and proper nouns are written in their correct standard Spanish forms (e.g., "Whitechapel", "Jack el Destripador", "Londres", "Alejandría", "Göbeklitepe").
4. Preserve Structure:
   - Maintain the exact same number of segments.

Spanish Data:
{json.dumps(payload, ensure_ascii=False, indent=2)}

STRICT OUTPUT FORMAT:
Return STRICT JSON only: {{"titles_es": [...], "description_es": "...", "segments_es": [{{"chapter_title_es": "...", "text_es": "...", "on_screen_text_es": "..."}}, ...]}}"""

        try:
            raw_text = self._chat(pc, prompt, max_tokens=8192)
            if raw_text:
                refined = self._extract_json(raw_text)
                segs_es = refined.get("segments_es", [])
                if len(segs_es) == len(segments):
                    for i, s in enumerate(segments):
                        s["text_es"] = segs_es[i].get("text_es", s.get("text_es", ""))
                        s["chapter_title_es"] = segs_es[i].get("chapter_title_es", s.get("chapter_title_es", ""))
                        s["on_screen_text_es"] = segs_es[i].get("on_screen_text_es", s.get("on_screen_text_es", ""))
                    if "titles_es" in refined and refined["titles_es"]:
                        script_data["titles_es"] = refined["titles_es"]
                    if "description_es" in refined and refined["description_es"]:
                        script_data["description_es"] = refined["description_es"]
                    logger.info("[SCRIPT:AUDIT] Spanish audit pass completed successfully.")
        except Exception as e:
            logger.warning(f"[SCRIPT:AUDIT] Spanish audit pass error: {e}")

        return script_data

    def _generate_phonetic_pronunciation_lexicon(self, pc: RemotePC, script_data: Dict[str, Any], topic: str) -> Dict[str, Dict[str, str]]:
        """
        Scan script for foreign, ancient, archaeological, mythological, or diacritic-heavy words
        (e.g., 'Göbeklitepe', 'Derinkuyu', 'Tutankhamun', 'Quetzalcoatl') and look up phonetic respellings
        for English and Spanish voice synthesizers.
        """
        logger.info("[SCRIPT:PHONETICS] Analyzing and looking up proper noun phonetic pronunciations...")
        
        segments = script_data.get("script_segments", [])
        sample_texts = [seg.get("text", "") for seg in segments[:10]] + [topic]
        combined_text = "\n".join(sample_texts)

        prompt = f"""You are an expert linguistic phonetics engine and voice synthesis director.
Analyze the following documentary script for topic "{topic}".

OBJECTIVE:
Identify any foreign, ancient, archaeological, mythological, scientific, non-standard, or diacritic-heavy words (e.g., "Göbeklitepe", "Derinkuyu", "Tutankhamun", "Quetzalcoatl", "Chichen Itza", "Oppenheimer", "Oumuamua", "Tiahuanaco", "Mohenjo-Daro", etc.) that a Text-to-Speech (TTS) voice synthesizer might struggle to pronounce correctly.

For EACH identified word, provide:
1. "en": Phonetic respelling for English TTS (using clear syllables, e.g., "Göbeklitepe" -> "Goh-beck-lee Teh-peh", "Chichen Itza" -> "Chee-chen Eet-zah")
2. "es": Phonetic respelling for Spanish TTS (using Spanish phonetic rules, e.g., "Göbeklitepe" -> "Guebekli Tepe", "Stonehenge" -> "Ston-jench")

If no difficult words exist, return empty objects.

Script text:
{combined_text}

STRICT JSON OUTPUT:
{{
  "phonetic_mappings": {{"en": {{"word": "phonetic_respelling_en"}}, "es": {{"word": "phonetic_respelling_es"}}}}
}}"""

        try:
            raw_text = self._chat(pc, prompt, max_tokens=2048)
            if raw_text:
                parsed = self._extract_json(raw_text)
                mappings = parsed.get("phonetic_mappings", {})
                if mappings and (mappings.get("en") or mappings.get("es")):
                    logger.info(f"[SCRIPT:PHONETICS] Discovered {len(mappings.get('en', {}))} phonetic pronunciation overrides.")
                    return mappings
        except Exception as e:
            logger.debug(f"[SCRIPT:PHONETICS] Phonetic lookup parse note: {e}")

        return {"en": {}, "es": {}}

    def generate(self, topic: str, duration: int) -> Dict[str, Any]:
        """Generate master bilingual script with secondary LLM verification & proofreading passes."""
        logger.info(f"[SCRIPT] Generating script for topic: '{topic}' (~{duration}s)...")
        pc = self.pcs.get("amdllm")
        if not pc or not pc.check(self.llm_port):
            raise RuntimeError(f"AMDLLM llama.cpp not reachable on port {self.llm_port}")

        # ── Step 1.0: Viral Topic Research & High-Retention Explainer Blueprint ──
        research_dossier = self.researcher.research_topic(topic)
        prompt = self._build_master_prompt(topic, duration, research_dossier=research_dossier)
        
        for attempt in range(1, 4):
            logger.info(f"[SCRIPT] Requesting script generation from AMDLLM (attempt {attempt}/3)...")
            raw_text = self._chat(pc, prompt, max_tokens=8192)
            if not raw_text:
                logger.warning(f"[SCRIPT] AMDLLM returned empty response on attempt {attempt}")
                continue

            try:
                data = self._extract_json(raw_text)
                
                # Validation & Normalization
                titles = data.get("titles", [])
                if isinstance(titles, str):
                    titles = [titles]
                while len(titles) < 3:
                    titles.append(f"{topic} - Part {len(titles)+1}")
                data["titles"] = titles[:3]
                data["title"] = titles[0]
                data["research_dossier"] = research_dossier
                data["visual_style_anchor"] = research_dossier.get("visual_style_anchor", "cinematic documentary style, warm archival tones, 35mm film texture, photorealistic 8k")

                segments = data.get("script_segments", [])
                if not segments:
                    raise ValueError("No script_segments generated")

                for i, s in enumerate(segments):
                    s.setdefault("chapter_title", f"Part {i+1}")
                    s.setdefault("duration_estimate", max(15, duration // len(segments)))
                    s.setdefault("stock_query", topic)
                    s.setdefault("visual_style_anchor", data["visual_style_anchor"])
                    if not s.get("visual_prompt"):
                        s["visual_prompt"] = f"Cinematic documentary footage illustrating: {s.get('text', topic)}"

                # ── Step 1.5: Secondary LLM Proofreading & Fact-Checking Pass (English) ──
                data = self._proofread_and_refine_script(pc, data, topic)

                # ── Step 1.6: Spanish Translation ──
                data = self._translate_to_spanish(pc, data)

                # ── Step 1.7: Secondary LLM Spanish Translation Audit ──
                data = self._proofread_spanish_translation(pc, data, topic)

                # ── Step 1.8: Phonetic Pronunciation Lookup for Non-Standard Vocabulary ──
                data["phonetic_mappings"] = self._generate_phonetic_pronunciation_lexicon(pc, data, topic)
                
                logger.info(f"[SCRIPT] Master script successfully created & audited: '{data['title']}' ({len(data.get('script_segments', []))} segments)")
                return data

            except Exception as e:
                logger.warning(f"[SCRIPT] Parse error on attempt {attempt}: {e}")

        raise RuntimeError("Failed to generate and parse script after 3 attempts.")