#!/usr/bin/env python3
"""
Backfill narration_script field for existing Layer 4 output files.

Reads each file, extracts voiceover from all segments in all beats,
assembles the narration_script field, and writes it back.
"""

import json
import sys
from pathlib import Path


def extract_narration_script(script: dict) -> str:
    """
    Extract all voiceover text from a script into a single narration string.

    Pulls voiceover from every segment in every beat, in order.
    Separates segments with blank lines, beats with --- dividers.
    Preserves all [beat], [pause], [thinking], [drops voice] markers.
    """
    beats = script.get("beats", [])
    if not beats:
        return ""

    beat_texts = []

    for beat in beats:
        segments = beat.get("segments", [])
        segment_texts = []

        for segment in segments:
            voiceover = segment.get("voiceover", "")
            if voiceover and voiceover.strip():
                segment_texts.append(voiceover.strip())

        if segment_texts:
            # Join segments within a beat with double newlines
            beat_text = "\n\n".join(segment_texts)
            beat_texts.append(beat_text)

    # Join beats with --- dividers
    narration = "\n\n---\n\n".join(beat_texts)

    return narration


def find_script_in_data(data: dict) -> dict:
    """
    Find the script object in various possible locations.

    Handles both old format (script at top level) and new format (script in output).
    """
    # New format: script inside output
    if "output" in data and isinstance(data["output"], dict):
        script = data["output"].get("script")
        if script and isinstance(script, dict) and "beats" in script:
            return script

    # Old format: script at top level
    if "script" in data and isinstance(data["script"], dict):
        script = data["script"]
        if "beats" in script:
            return script

    return None


def backfill_file(filepath: Path) -> bool:
    """
    Backfill a single file with narration_script.

    Returns True if file was modified, False otherwise.
    """
    try:
        with open(filepath, 'r') as f:
            data = json.load(f)
    except (json.JSONDecodeError, IOError) as e:
        print(f"  [SKIP] Cannot read {filepath.name}: {e}")
        return False

    # Check if already has narration_script
    if "narration_script" in data and data["narration_script"]:
        print(f"  [SKIP] {filepath.name} - already has narration_script")
        return False

    # Find the script
    script = find_script_in_data(data)
    if not script:
        print(f"  [SKIP] {filepath.name} - no script with beats found")
        return False

    # Extract narration
    narration = extract_narration_script(script)
    if not narration:
        print(f"  [SKIP] {filepath.name} - no voiceover content found")
        return False

    # Add narration_script at the top level
    # For new format files, insert after meta
    if "meta" in data:
        # Rebuild dict with narration_script in right position
        new_data = {}
        for key, value in data.items():
            new_data[key] = value
            if key == "meta":
                new_data["narration_script"] = narration
        data = new_data
    else:
        # Old format: just add at top level
        data["narration_script"] = narration

    # Write back
    with open(filepath, 'w') as f:
        json.dump(data, f, indent=2, default=str)

    word_count = len(narration.split())
    print(f"  [OK] {filepath.name} - added narration_script ({word_count} words)")
    return True


def main():
    output_dir = Path(__file__).parent / "output"

    if not output_dir.exists():
        print("No output directory found")
        return 1

    print("=" * 60)
    print("BACKFILLING NARRATION_SCRIPT FIELD")
    print("=" * 60)

    # Find all layer4 files (old format)
    old_format_files = list(output_dir.glob("layer4_*.json"))
    # Exclude input files and summary
    old_format_files = [f for f in old_format_files
                        if "input" not in f.name and "summary" not in f.name]

    # Find all files in channel subfolders (new format)
    channel_files = []
    for channel_dir in output_dir.iterdir():
        if channel_dir.is_dir() and channel_dir.name != "__pycache__":
            channel_files.extend(channel_dir.glob("*.json"))

    all_files = old_format_files + channel_files

    if not all_files:
        print("No Layer 4 output files found to backfill")
        return 0

    print(f"\nFound {len(all_files)} files to process:")
    print(f"  - Old format (output/layer4_*.json): {len(old_format_files)}")
    print(f"  - New format (channel folders): {len(channel_files)}")
    print()

    modified = 0
    for filepath in sorted(all_files):
        if backfill_file(filepath):
            modified += 1

    print()
    print("=" * 60)
    print(f"COMPLETE: Modified {modified}/{len(all_files)} files")
    print("=" * 60)

    return 0


if __name__ == "__main__":
    sys.exit(main())