# SYNTHESIZING CHARACTER SCENE POOL PIPELINE
import os
import sys
import json
import argparse
from datetime import datetime
from typing import List, Dict, Optional, Tuple
import traceback

# Import modules for each step
from topic_generation import generate_topic
from worldview_generation import generate_worldview_from_scratch, generate_worldview, save_worldview
from character_generation import (
    generate_character_profile_from_saved_worldview, 
    generate_character_profile,
    save_character_profile,
    save_separated_character_profile,
    update_tracking
)
from scene_generation import generate_scene, save_scene, load_character_profiles
from initial_dialogue_generation import generate_dialogue, save_dialogue, load_character_profiles as load_chars_for_dialogue
from character_selector import integrate_with_dialogue_generation
from utils import read_yaml_to_string, get_llm_response


class SynthesizingPipeline:
    """Complete synthesizing pipeline class"""
    
    def __init__(self, model_name: str = "claude-3-7-sonnet-20250219", temperature: float = 1.0):
        """
        Initialize pipeline
        
        Args:
            model_name: Used language model name
            temperature: Temperature parameter when generating
        """
        self.model_name = model_name
        self.temperature = temperature
        self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # Store result paths for each step
        self.results = {
            "topic": None,
            "worldview_path": None,
            "character_profile_path": None,
            "separated_characters_path": None,
            "scene_path": None,
            "dialogue_path": None
        }
        
        # Create necessary directories
        self._create_directories()
    
    def _create_directories(self):
        """Create necessary directory structure"""
        directories = [
            "raw_topics", "scored_topics", "best_topic",
            "worldview", "character_profile", 
            f"character_pool/separated_{self.timestamp}",
            "scene", "dialogue", "pipeline_results"
        ]
        
        for directory in directories:
            os.makedirs(directory, exist_ok=True)
    
    def step1_generate_topic(self) -> str:
        """
        Step 1: Generate topic
        
        Returns:
            str: Generated best topic
        """
        print("=" * 60)
        print("Step 1/5: Generate topic (Topic Generation)")
        print("=" * 60)
        
        try:
            topic = generate_topic(model=self.model_name, temperature=self.temperature)
            self.results["topic"] = topic
            print(f"✓ Topic generation completed: {topic}")
            return topic
        except Exception as e:
            print(f"✗ Topic generation failed: {e}")
            raise
    
    def step2_generate_worldview(self, topic: str) -> str:
        """
        Step 2: Generate worldview
        
        Args:
            topic: Topic
            
        Returns:
            str: Worldview file path
        """
        print("\n" + "=" * 60)
        print("Step 2/5: Generate worldview (Worldview Generation)")
        print("=" * 60)
        
        try:
            worldview_content = generate_worldview(topic, self.model_name)
            worldview_path = f"worldview/worldview_{self.timestamp}.yaml"
            save_worldview(worldview_content, worldview_path)
            
            self.results["worldview_path"] = worldview_path
            print(f"✓ Worldview generation completed, saved to: {worldview_path}")
            return worldview_path
        except Exception as e:
            print(f"✗ Worldview generation failed: {e}")
            raise
    
    def step3_generate_characters(self, worldview_path: str, num_characters: int = 3) -> Tuple[str, str]:
        """
        Step 3: Generate characters
        
        Args:
            worldview_path: Worldview file path
            num_characters: Number of characters to generate
            
        Returns:
            Tuple[str, str]: (Character profile file path, Separated characters folder path)
        """
        print("\n" + "=" * 60)
        print("Step 3/5: Generate characters (Character Generation)")
        print("=" * 60)
        
        try:
            # Read worldview
            worldview_content = read_yaml_to_string(worldview_path)
            
            # Generate character profile
            character_profile = generate_character_profile(
                worldview_content, 
                self.model_name, 
                number_of_characters=num_characters,
                temperature=self.temperature
            )
            
            # Save character profile
            character_profile_path = f"character_profile/character_profile_{self.timestamp}.yaml"
            separated_path = f"character_pool/separated_{self.timestamp}"
            
            save_character_profile(character_profile, character_profile_path)
            save_separated_character_profile(character_profile, separated_path)
            
            # Update tracking information
            update_tracking(worldview_path, character_profile_path, num_characters, separated_path)
            
            self.results["character_profile_path"] = character_profile_path
            self.results["separated_characters_path"] = separated_path
            
            print(f"✓ Character generation completed, saved to: {character_profile_path}")
            print(f"✓ Separated characters saved to: {separated_path}")
            
            return character_profile_path, separated_path
        except Exception as e:
            print(f"✗ Character generation failed: {e}")
            raise
    
    def step4_generate_scene(self, worldview_path: str, separated_characters_path: str) -> str:
        """
        Step 4: Generate scene
        
        Args:
            worldview_path: Worldview file path
            separated_characters_path: Separated characters folder path
            
        Returns:
            str: Scene file path
        """
        print("\n" + "=" * 60)
        print("Step 4/5: Generate scene (Scene Generation)")
        print("=" * 60)
        
        try:
            # Read worldview
            worldview_content = read_yaml_to_string(worldview_path)
            
            # Load character profile
            character_paths = [f"{separated_characters_path}/*.yaml"]
            character_profiles = load_character_profiles(character_paths)
            
            if not character_profiles:
                raise ValueError("No valid character files found")
            
            print(f"✓ Loaded {len(character_profiles)} characters")
            
            # Generate scene
            scene_content = generate_scene(
                worldview_content, 
                character_profiles, 
                self.model_name,
                temperature=self.temperature
            )
            
            # Save scene
            scene_path = f"scene/scene_{self.timestamp}.yaml"
            save_scene(scene_content, scene_path)
            
            # Record used character files
            record_file = f"scene/scene_{self.timestamp}_character_files.txt"
            with open(record_file, "w", encoding="utf-8") as f:
                f.write(f"Worldview: {worldview_path}\n")
                f.write(f"Characters: {separated_characters_path}/*.yaml\n")
            
            self.results["scene_path"] = scene_path
            print(f"✓ Scene generation completed, saved to: {scene_path}")
            
            return scene_path
        except Exception as e:
            print(f"✗ Scene generation failed: {e}")
            raise
    
    def step5_generate_dialogue(self, scene_path: str, worldview_path: str, separated_characters_path: str) -> str:
        """
        Step 5: Generate dialogue
        
        Args:
            scene_path: Scene file path
            worldview_path: Worldview file path
            separated_characters_path: Separated characters folder path
            
        Returns:
            str: Dialogue file path
        """
        print("\n" + "=" * 60)
        print("Step 5/5: Generate dialogue (Initial Dialogue Generation)")
        print("=" * 60)
        
        try:
            # Read scene and worldview
            scene_content = read_yaml_to_string(scene_path)
            worldview_content = read_yaml_to_string(worldview_path)
            
            # Load character profile
            character_paths = [f"{separated_characters_path}/*.yaml"]
            character_profiles, character_files = load_chars_for_dialogue(character_paths)
            
            if not character_profiles:
                raise ValueError("No valid character files found")
            
            print(f"✓ Loaded {len(character_profiles)} characters")
            
            # Extract character names
            character_names = []
            for file_path in character_files:
                base_name = os.path.basename(file_path)
                name_without_ext = os.path.splitext(base_name)[0]
                character_name = name_without_ext.replace('_', ' ')
                character_names.append(character_name)
            
            # Generate dialogue
            dialogue_content = generate_dialogue(
                scene_content,
                worldview_content,
                character_profiles,
                self.model_name,
                temperature=self.temperature,
                character_names=character_names
            )
            
            # Save dialogue
            dialogue_path = f"dialogue/dialogue_{self.timestamp}.json"
            save_dialogue(dialogue_content, dialogue_path)
            
            # Record used files
            record_file = f"dialogue/dialogue_{self.timestamp}_files.txt"
            with open(record_file, "w", encoding="utf-8") as f:
                f.write(f"Scene: {scene_path}\n")
                f.write(f"Worldview: {worldview_path}\n")
                f.write(f"Characters: {separated_characters_path}/*.yaml\n")
            
            self.results["dialogue_path"] = dialogue_path
            print(f"✓ Dialogue generation completed, saved to: {dialogue_path}")
            
            return dialogue_path
        except Exception as e:
            print(f"✗ Dialogue generation failed: {e}")
            raise
    
    def run_full_pipeline(self, num_characters: int = 3, use_selector: bool = False) -> Dict[str, str]:
        """
        Run complete pipeline
        
        Args:
            num_characters: Number of characters to generate
            use_selector: Whether to use character selector
            
        Returns:
            Dict[str, str]: Result file paths for each step
        """
        print("🚀 Start running complete character scene pool synthesizing pipeline")
        print(f"📋 Configuration: model={self.model_name}, temperature={self.temperature}, number of characters={num_characters}")
        print(f"🕐 Timestamp: {self.timestamp}")
        
        try:
            # Step 1: Generate topic
            topic = self.step1_generate_topic()
            
            # Step 2: Generate worldview
            worldview_path = self.step2_generate_worldview(topic)
            
            # Step 3: Generate characters
            character_profile_path, separated_characters_path = self.step3_generate_characters(
                worldview_path, num_characters
            )
            
            # If character selector is enabled, let user select characters
            if use_selector:
                print("\n" + "=" * 60)
                print("Character Selector (Character Selector)")
                print("=" * 60)
                try:
                    selected_worldview, selected_characters = integrate_with_dialogue_generation()
                    if selected_worldview and selected_characters:
                        worldview_path = selected_worldview
                        # Here we need to handle the selected character files
                        print(f"✓ User selected worldview: {selected_worldview}")
                        print(f"✓ User selected {len(selected_characters)} characters")
                    else:
                        print("⚠ User did not select characters, using default generated characters")
                except Exception as e:
                    print(f"⚠ Character selector failed: {e}")
                    print("Using default generated characters continue")
            
            # Step 4: Generate scene
            scene_path = self.step4_generate_scene(worldview_path, separated_characters_path)
            
            # Step 5: Generate dialogue
            dialogue_path = self.step5_generate_dialogue(scene_path, worldview_path, separated_characters_path)
            
            # Save complete pipeline result summary
            self._save_pipeline_summary()
            
            print("\n" + "🎉" * 20)
            print("✅ Complete pipeline execution successful!")
            print("🎉" * 20)
            print("\n📁 Generated files:")
            for step, path in self.results.items():
                if path:
                    print(f"  {step}: {path}")
            
            return self.results
            
        except Exception as e:
            print(f"\n❌ Pipeline execution failed: {e}")
            print("\n🔍 Error details:")
            traceback.print_exc()
            raise
    
    def _save_pipeline_summary(self):
        """Save pipeline execution summary"""
        summary = {
            "timestamp": self.timestamp,
            "model_name": self.model_name,
            "temperature": self.temperature,
            "results": self.results,
            "topic": self.results.get("topic"),
            "execution_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        
        summary_path = f"pipeline_results/pipeline_summary_{self.timestamp}.json"
        with open(summary_path, "w", encoding="utf-8") as f:
            json.dump(summary, f, ensure_ascii=False, indent=4)
        
        print(f"📋 Pipeline summary saved to: {summary_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Complete character scene pool synthesizing pipeline",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
        Usage examples:
        # Basic usage
        python synthesizing_character_scene_pool_pipeline.py
        
        # Specify model and parameters
        python synthesizing_character_scene_pool_pipeline.py --model claude-3-7-sonnet-20250219 --temperature 1.0 --num-characters 3
        
        # Enable character selector
        python synthesizing_character_scene_pool_pipeline.py --use-selector
        
        # Use different model
        python synthesizing_character_scene_pool_pipeline.py --model gemini-2.5-pro-preview-06-05 --temperature 0.8
        """
    )
    
    parser.add_argument(
        "--model", 
        type=str, 
        default="claude-3-7-sonnet-20250219",
        help="Used language model name (default: claude-3-7-sonnet-20250219)"
    )
    
    parser.add_argument(
        "--temperature", 
        type=float, 
        default=1.0,
        help="Temperature parameter when generating (default: 1.0)"
    )
    
    parser.add_argument(
        "--num-characters", 
        type=int, 
        default=3,
        help="Number of characters to generate (default: 3)"
    )
    
    parser.add_argument(
        "--use-selector", 
        action="store_true",
        help="Enable character selector, allow user to manually select characters"
    )
    
    parser.add_argument(
        "--output-dir", 
        type=str, 
        default=".",
        help="Output directory (default: current directory)"
    )
    
    args = parser.parse_args()
    
    # Switch to output directory
    if args.output_dir != ".":
        os.makedirs(args.output_dir, exist_ok=True)
        os.chdir(args.output_dir)
    
    # Create and run pipeline
    pipeline = SynthesizingPipeline(
        model_name=args.model,
        temperature=args.temperature
    )
    
    try:
        results = pipeline.run_full_pipeline(
            num_characters=args.num_characters,
            use_selector=args.use_selector
        )
        
        print(f"\n🎯 All files have been generated!")
        print(f"📂 View pipeline_results/ directory for complete summary")
        
        return 0
        
    except KeyboardInterrupt:
        print("\n⚠ User interrupted execution")
        return 1
    except Exception as e:
        print(f"\n❌ Execution failed: {e}")
        return 1


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