#!/usr/bin/env python3
"""
Test script for the Visual Decider on the Tobacco Company premise.
"""

import json
import sys
from pathlib import Path

# Add project to path
sys.path.insert(0, str(Path(__file__).parent))

from video_generation import layer4_visual_decider, run_layer4_v4_pipeline

def load_tobacco_opportunity():
    """Load the Tobacco Company opportunity with research."""
    output_file = Path(__file__).parent / "output" / "layer4_output_one_minute_history.json"

    with open(output_file) as f:
        data = json.load(f)

    # The file has the full opportunity data
    return data

def run_visual_decider_only(opportunity):
    """Run just the Visual Decider on an existing script."""
    script = opportunity.get("script")
    if not script:
        print("ERROR: No script found in opportunity")
        return None

    print("=" * 60)
    print("RUNNING VISUAL DECIDER TEST")
    print("=" * 60)
    print(f"Premise: {script.get('metadata', {}).get('premise', 'Unknown')}")
    print(f"Beats: {len(script.get('beats', []))}")

    total_segments = sum(len(b.get('segments', [])) for b in script.get('beats', []))
    print(f"Total segments: {total_segments}")
    print()

    # Run the visual decider
    enhanced_script = layer4_visual_decider(script, opportunity)

    return enhanced_script

def print_segment_analysis(script):
    """Print detailed analysis of each segment's visual decision."""
    print("\n" + "=" * 60)
    print("SEGMENT-BY-SEGMENT VISUAL ANALYSIS")
    print("=" * 60)

    for beat in script.get("beats", []):
        beat_name = beat.get("beat_name", "Unknown")
        print(f"\n### BEAT {beat.get('beat_number')}: {beat_name}")
        print("-" * 40)

        for i, segment in enumerate(beat.get("segments", [])):
            seg_id = f"b{beat.get('beat_number')}_s{i+1}"

            print(f"\n[{seg_id}] {segment.get('time_start', 0):.1f}s - {segment.get('time_end', 0):.1f}s")
            vo = segment.get('voiceover', '')
            print(f"VOICEOVER: \"{vo[:100]}...\"" if len(vo) > 100 else f"VOICEOVER: \"{vo}\"")

            # Check if we have structured visual
            vs = segment.get("visual_structured", {})
            decision = vs.get("decision", "UNKNOWN")
            print(f"DECISION: {decision}")

            if decision == "IMAGE":
                img = vs.get("image", {})
                print(f"SEARCH: \"{img.get('search_query', 'N/A')}\"")
                if img.get("era"):
                    print(f"ERA: {img.get('era')}")
            elif decision == "VIDEO":
                vid = vs.get("video", {})
                print(f"SEARCH: \"{vid.get('search_query', 'N/A')}\"")
            elif decision == "DATA_VISUALIZATION":
                dv = vs.get("data_visualization", {})
                print(f"CHART: {dv.get('chart_type', 'N/A')}")
                print(f"TITLE: {dv.get('title', 'N/A')}")
                data = dv.get("data", [])
                print(f"DATA POINTS: {len(data)}")
                for dp in data[:3]:
                    print(f"  - {dp.get('label')}: {dp.get('value')} {dp.get('unit', '')}")
                if dv.get("source"):
                    print(f"SOURCE: {dv.get('source')[:80]}...")

            # Show thinking
            thinking = vs.get("thinking", segment.get("visual_thinking", ""))
            if thinking:
                print(f"THINKING: {thinking[:120]}...")

            # Show data gap if present
            data_gap = vs.get("if_data_not_available") or segment.get("visual_data_gap")
            if data_gap:
                print(f"DATA GAP: {data_gap.get('why_not_possible', 'N/A')}")
                print(f"FALLBACK: {data_gap.get('fallback_image_query', 'N/A')}")

            print()

def save_results(script, filename="test_visual_decider_output.json"):
    """Save the enhanced script for review."""
    output_path = Path(__file__).parent / "output" / filename
    with open(output_path, "w") as f:
        json.dump(script, f, indent=2)
    print(f"\nResults saved to: {output_path}")
    return output_path

if __name__ == "__main__":
    # Load opportunity
    print("Loading Tobacco Company opportunity...")
    opportunity = load_tobacco_opportunity()

    # Run visual decider
    enhanced_script = run_visual_decider_only(opportunity)

    if enhanced_script:
        # Print analysis
        print_segment_analysis(enhanced_script)

        # Save results
        output_path = save_results(enhanced_script)

        # Print summary
        print("\n" + "=" * 60)
        print("SUMMARY")
        print("=" * 60)

        metadata = enhanced_script.get("visual_decider_metadata", {})
        print(f"Processed: {metadata.get('processed_segments', 0)}/{metadata.get('total_segments', 0)} segments")
        print(f"Duration: {metadata.get('duration_seconds', 0):.1f}s")

        # Count types
        type_counts = {"IMAGE": 0, "VIDEO": 0, "DATA_VISUALIZATION": 0, "UNKNOWN": 0}
        for beat in enhanced_script.get("beats", []):
            for segment in beat.get("segments", []):
                vs = segment.get("visual_structured", {})
                decision = vs.get("decision", "UNKNOWN")
                type_counts[decision] = type_counts.get(decision, 0) + 1

        print(f"\nType distribution:")
        for vtype, count in type_counts.items():
            if count > 0:
                print(f"  {vtype}: {count}")
    else:
        print("ERROR: Visual decider failed")