"""Documentary mode, following the 'Ultimate AI Documentary Master Prompt' rules, automated:

  research (research.py) -> blueprint + research status + chunked script -> fact check against sources
  -> one 4K cinematic photoreal image prompt per chunk, with continuity and shot variety.

The PDF's interactive intake (questions 1-6) maps to UI controls (topic, duration, language, audience,
tone); its NEXT/STOP phases run here one after another. Every function takes `ask(system, user,
temperature, max_tokens) -> dict` so it can be tested without a real model.
"""
from __future__ import annotations

import re

TONES = ["Cinematic & Dramatic", "Serious & Investigative", "Educational & Informative",
         "Emotional & Human-Centered", "Suspenseful & Tense", "Balanced Professional Documentary"]
AUDIENCES = ["General Audience", "History Enthusiasts", "Students / Educational",
             "YouTube General Audience", "Documentary / Film Audience"]
FUNCTIONS = ["Hook", "Setup", "Development", "Tension", "Turning Point", "Climax", "Aftermath", "Closing"]
SHOTS = ["extreme wide", "wide", "establishing", "medium", "close-up", "extreme close-up", "low angle",
         "high angle", "over-the-shoulder", "environmental portrait", "detail", "tracking"]
FOUR_K = "4K cinematic, photorealistic live-action historical documentary"
MAX_WORDS = 22

GENERIC_PHRASES = ["Nobody could have imagined what happened next", "Everything was about to change",
                   "A disaster was waiting to happen", "The world would never be the same",
                   "The situation became increasingly dangerous", "It was a moment in history"]


def chunk_count(seconds: int, words_per_second: float, max_chunks: int) -> int:
    """PDF: one 16-21 word chunk per visual unit (1 minute = 6 chunks at a slow English pace).
    Scaled by the narrator's real speaking rate so the video fills the chosen length."""
    return max(4, min(max_chunks, round(seconds * words_per_second / 19)))


def words(text: str) -> int:
    return len(re.sub(r"[*\"'“”]", "", text).split())


# ------------------------------------------------------------------------------------ 1. script
def write_script(ask, topic: str, research: dict, language_rules: str, seconds: int, n_chunks: int,
                 tone: str, audience: str) -> dict:
    system = f"""You are a professional historical documentary researcher and scriptwriter.
RESEARCH FIRST, STORY SECOND. You are given SOURCE TEXT. Use ONLY facts that appear in it.
Never invent: dates, times, numbers, names, places, quotes, statistics, dialogue, private thoughts, emotions,
motivations, secret meetings, heroic or villainous behaviour. If the sources are vague or disagree, say so honestly
("approximately", "records indicate", "historians generally agree", "the exact figure remains disputed").
Never create false precision. Never present a reconstruction as confirmed fact.

PROCESS: from the sources establish WHEN -> WHERE -> WHO -> WHAT -> HOW -> WHY -> CONSEQUENCE, then write.

NARRATION RULES
- Exact facts drive the story and are embedded naturally: DATE -> LOCATION -> PEOPLE -> SITUATION -> FACTUAL
  DEVELOPMENT -> CONSEQUENCE -> NEXT EVENT. Not a list of statistics.
- Prefer specific verified statements over vague ones ("left Southampton on 10 April 1912 with more than 2,000
  people aboard", not "carried many people"). Every number must serve the story.
- Strict chronology. Move the viewer through the timeline naturally ("On ...", "Four days later ...", "By the evening of ...").
- Let the facts create the drama, where the evidence supports it: NORMALITY -> CONTEXT -> WARNING -> DOCUMENTED
  DEVELOPMENT -> ESCALATION -> CONSEQUENCE -> CRISIS -> CLIMAX -> AFTERMATH. Do not force a mystery, twist,
  villain or hero the evidence does not support.
- Every chunk must advance the story AND the viewer's knowledge (a fact, a time, a place, a person, a cause, a
  consequence). Never spend a chunk only on atmosphere.
- Avoid generic filler such as: {'; '.join(GENERIC_PHRASES)}.
- One professional narrator, continuous story: each chunk follows naturally from the previous one.
- This is a SHORT documentary: essential chronology, major verified facts, key numbers, the turning point, the outcome.
- Chunk 1 is the hook: the most striking VERIFIED fact or question, no clickbait.
- People appear only through documented behaviour.

FORMAT
- Exactly {n_chunks} chunks. Each chunk: one or two complete sentences, maximum {MAX_WORDS} words, preferred 16-21.
- Story functions in order, chosen from: {', '.join(FUNCTIONS)}.
- Tone: {tone}. Audience: {audience}.
- Wrap at most 3 key facts (a date, a number) in *asterisks* for caption highlighting.
{language_rules}
- "title" and "title_options" in Roman-script Hinglish or English, honest, max 60 characters.

Return ONLY JSON:
{{"blueprint": {{"title": "", "core_question": "", "premise": "one line", "period": "", "locations": [], "story_structure": ""}},
 "research_status": {{"verified_facts": ["max 8, each with its figure/date"], "requiring_verification": [],
                     "reconstructions": [], "uncertainties": []}},
 "chunks": [{{"chunk": 1, "story_function": "Hook", "narration": "", "visual_intent": "what the viewer should see"}}],
 "title_options": ["5 titles"], "description": "2-3 lines for YouTube", "hashtags": ["#shorts", "..."]}}"""
    user = f"TOPIC: {topic}\nTARGET RUNTIME: about {seconds} seconds\n\nSOURCE TEXT:\n{research['text']}"
    doc = ask(system, user, 0.4, 5000)
    chunks = [c for c in (doc.get("chunks") or []) if str(c.get("narration", "")).strip()]
    if len(chunks) < 3:
        raise ValueError(f"documentary script has only {len(chunks)} usable chunks")
    for i, c in enumerate(chunks, start=1):
        c["chunk"] = i
        c["narration"] = str(c["narration"]).strip()
        if c.get("story_function") not in FUNCTIONS:
            c["story_function"] = FUNCTIONS[min(i - 1, len(FUNCTIONS) - 1)]
    doc["chunks"] = chunks[:n_chunks + 2]          # tolerate +2, never pad
    bp = doc.get("blueprint") or {}
    doc["title"] = str(bp.get("title") or doc.get("title") or topic).strip()
    return doc


# ------------------------------------------------------------------------------------ 2. fact check
def fact_check(ask, research: dict, doc: dict, language_rules: str) -> tuple[dict, list[dict]]:
    system = f"""You are a strict documentary fact-checker. Compare every narration chunk with the SOURCE TEXT.
Flag: any date, time, number, name, place, quote, cause, sequence or claim that the source does not support,
anything stated with more precision than the source gives, invented dialogue/thoughts/emotions, wrong chronology.
Rewrite ONLY the flawed chunks so they are supported by the source (hedge if needed). Keep the same language and
style, maximum {MAX_WORDS} words, keep *asterisks* style. Keep the chunk count and order.
{language_rules}
Return ONLY JSON: {{"issues": [{{"chunk": 1, "problem": "", "fix": ""}}],
 "chunks": [{{"chunk": 1, "narration": "all chunks, unchanged ones copied exactly"}}]}}"""
    script = "\n".join(f"CHUNK {c['chunk']}: {c['narration']}" for c in doc["chunks"])
    res = ask(system, f"SOURCE TEXT:\n{research['text']}\n\nSCRIPT:\n{script}", 0.1, 4000)
    issues = [i for i in (res.get("issues") or []) if isinstance(i, dict)]
    fixed = res.get("chunks") or []
    if len(fixed) == len(doc["chunks"]):
        for c, f in zip(doc["chunks"], fixed):
            new = str(f.get("narration", "")).strip()
            if new and new != c["narration"]:
                c["narration_before_check"] = c["narration"]
                c["narration"] = new
    return doc, issues


# ------------------------------------------------------------------------------------ 3. image prompts
def image_prompts(ask, doc: dict, research: dict, aspect: str = "vertical 9:16") -> dict:
    system = f"""You are the director of photography of a premium historical documentary.
Write exactly ONE image prompt per chunk. Every prompt MUST begin with: "{FOUR_K}".
Each prompt establishes: exact subject, exact action, historical period, location, characters and clothing,
foreground / midground / background, camera angle and shot size, lens feeling, depth of field, lighting,
atmosphere, color grading, materials. 60-110 words, English.
RULES
- Historical period lock: no modern buildings, vehicles, technology, clothing, weapons, signage or lighting
  unless they belong to the scene. Geographic lock: architecture, terrain, vegetation and climate match the place.
- No text, letters, numbers, signs, logos, headlines or typography in the image.
- Global visual tone: every image belongs to the same film - one consistent grade, contrast, lens and
  lighting philosophy. No painterly, cartoon, CGI, vintage or oversaturated looks.
- Realism: realistic anatomy, natural skin texture, realistic fabrics, shadows, reflections and scale, natural light.
- Continuity: define recurring characters, locations and props once in "visual_world" and repeat their
  appearance word-for-word in every prompt where they appear.
- Cinematography variety: choose a shot from {', '.join(SHOTS)}; never the same shot size twice in a row;
  open with an establishing or extreme wide shot.
- Human cost is shown with restraint (aftermath, empty objects, distance, scale) - no gore or graphic injury.
- Composition: {aspect}, main subject readable in the centre third.
Return ONLY JSON: {{"visual_world": {{"period": "", "location": "", "look": "grade, lens, light",
 "recurring_characters": [{{"name": "", "appearance": "", "clothing": ""}}], "recurring_props": []}},
 "images": [{{"chunk": 1, "shot": "establishing", "prompt": "{FOUR_K}, ..."}}]}}"""
    blueprint = doc.get("blueprint") or {}
    plan = "\n".join(f"CHUNK {c['chunk']} [{c['story_function']}]: {c['narration']}\n  VISUAL INTENT: "
                     f"{c.get('visual_intent', '')}" for c in doc["chunks"])
    user = (f"DOCUMENTARY: {doc['title']}\nPERIOD: {blueprint.get('period', '')}\n"
            f"LOCATIONS: {', '.join(blueprint.get('locations') or [])}\n\n{plan}\n\n"
            f"SOURCE EXCERPT (for period-accurate detail):\n{research['text'][:5000]}")
    res = ask(system, user, 0.5, 5000)
    images = {int(i.get("chunk", 0)): i for i in (res.get("images") or []) if isinstance(i, dict)}
    out = []
    for c in doc["chunks"]:
        img = images.get(c["chunk"]) or {}
        prompt = str(img.get("prompt") or c.get("visual_intent") or c["narration"]).strip()
        if not prompt.lower().startswith("4k"):
            prompt = f"{FOUR_K}, {prompt}"
        shot = str(img.get("shot", "")).lower()
        out.append({"chunk": c["chunk"], "shot": shot if shot in SHOTS else "", "prompt": prompt})
    return {"visual_world": res.get("visual_world") or {}, "images": out}


def camera_motion(shot: str, i: int) -> str:
    """PDF motion principles: slow push-in, gentle pan, controlled tracking - no excessive zoom."""
    if shot in ("extreme wide", "establishing", "wide", "tracking"):
        return "pan_right" if i % 2 else "pan_left"
    if shot == "low angle":
        return "pan_up"
    if shot == "high angle":
        return "pan_down"
    if shot in ("close-up", "extreme close-up", "detail", "medium", "environmental portrait", "over-the-shoulder"):
        return "zoom_in"
    return ["zoom_in", "pan_right", "zoom_out", "pan_left"][i % 4]


def to_board(doc: dict, visuals: dict, clean_tts) -> dict:
    scenes = []
    for i, (c, img) in enumerate(zip(doc["chunks"], visuals["images"]), start=1):
        scenes.append({"scene_number": i, "narration": c["narration"], "tts_text": clean_tts(c["narration"]),
                       "story_function": c["story_function"], "visual_prompt": c.get("visual_intent", ""),
                       "final_prompt": img["prompt"] + ", vertical 9:16 frame, no text",
                       "shot": img["shot"], "camera_motion": camera_motion(img["shot"], i),
                       "mood": "mysterious", "sfx": "none"})
    return {"character": (visuals.get("visual_world") or {}).get("recurring_characters") or {},
            "visual_style": (visuals.get("visual_world") or {}).get("look") or FOUR_K,
            "visual_world": visuals.get("visual_world") or {}, "music_mood": "mysterious", "scenes": scenes}


def story_view(doc: dict, research: dict, issues: list[dict]) -> dict:
    """Shape the documentary like a story dict so the rest of the pipeline and UI stay the same."""
    bp = doc.get("blueprint") or {}
    return {"mode": "documentary", "title": doc["title"], "hook": doc["chunks"][0]["narration"],
            "story": " ".join(c["narration"] for c in doc["chunks"]), "twist": bp.get("premise", ""),
            "ending": doc["chunks"][-1]["narration"], "concept": bp.get("premise", ""),
            "blueprint": bp, "research_status": doc.get("research_status") or {},
            "fact_check_issues": issues, "sources": research["sources"],
            "title_options": doc.get("title_options") or [doc["title"]],
            "description": doc.get("description", ""), "hashtags": doc.get("hashtags") or ["#shorts", "#history"],
            "chunks": doc["chunks"]}