TravelContentCreator/main.py

import os
import time
from datetime import datetime
import argparse
import sys
import traceback
import json

from core.ai_agent import AI_Agent
# from core.topic_parser import TopicParser # No longer needed directly in main?
import core.contentGen as contentGen
import core.posterGen as posterGen
import core.simple_collage as simple_collage
from utils.resource_loader import ResourceLoader
from utils.tweet_generator import ( # Import the moved functions
    run_topic_generation_pipeline, 
    generate_content_for_topic,
    generate_posters_for_topic
)
from utils.prompt_manager import PromptManager # Import PromptManager
import random

def load_config(config_path="poster_gen_config.json"):
    """Loads configuration from a JSON file."""
    if not os.path.exists(config_path):
        print(f"Error: Configuration file '{config_path}' not found.")
        print("Please copy 'example_config.json' to 'poster_gen_config.json' and customize it.")
        sys.exit(1)
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        # Basic validation (can be expanded)
        required_keys = ["api_url", "model", "api_key", "resource_dir", "prompts_dir", "output_dir", "num", "variants", "topic_system_prompt", "topic_user_prompt", "content_system_prompt", "image_base_dir"]
        if not all(key in config for key in required_keys):
             missing_keys = [key for key in required_keys if key not in config]
             print(f"Error: Config file '{config_path}' is missing required keys: {missing_keys}")
             sys.exit(1)
        # Resolve relative paths based on config location or a defined base path if necessary
        # For simplicity, assuming paths in config are relative to project root or absolute
        return config
    except json.JSONDecodeError:
        print(f"Error: Could not decode JSON from '{config_path}'. Check the file format.")
        sys.exit(1)
    except Exception as e:
        print(f"Error loading configuration from '{config_path}': {e}")
        sys.exit(1)

# Removed generate_topics_step function definition from here
# Its logic is now in utils.tweet_generator.run_topic_generation_pipeline

# --- Main Orchestration Step (Remains in main.py) --- 

def generate_content_and_posters_step(config, run_id, tweet_topic_record):
    """Generates content and posters by calling decoupled functions for each topic."""
    if not run_id or not tweet_topic_record or not tweet_topic_record.topics_list:
        print("Missing run_id or valid topics data. Skipping content and poster generation.")
        return

    print("\nStep 2: Processing Topics for Content and Posters...")
    base_output_dir = config["output_dir"]
                                                        
    # --- Initialize shared resources once --- 
    prompt_manager = None
    ai_agent = None
    # We might not need to initialize these here if the moved functions handle it
    # content_gen = None 
    # poster_gen_instance = None
    initialized_ok = False
    try:
        prompt_manager = PromptManager(config)
        print(f"Initializing AI Agent ({config['model']})...")
        ai_agent = AI_Agent(config["api_url"], config["model"], config["api_key"])
        # content_gen = core_contentGen.ContentGenerator() # Instantiation moved to generate_posters_for_topic
        # poster_gen_instance = core_posterGen.PosterGenerator() # Instantiation moved to generate_posters_for_topic
        
        # Check image base directory validity once
        image_base_dir = config.get("image_base_dir", None)
        if not image_base_dir or not os.path.isdir(image_base_dir):
            raise ValueError(f"'image_base_dir' ({image_base_dir}) not specified or not a valid directory in config.")
        initialized_ok = True
        
    except Exception as e:
        print(f"Error initializing resources for content/poster generation: {e}")
        traceback.print_exc()
        if ai_agent: ai_agent.close()
        return
        
    # --- Process each topic using decoupled functions --- 
    total_topics = len(tweet_topic_record.topics_list)
    for i, topic in enumerate(tweet_topic_record.topics_list):
        topic_index = i + 1
        print(f"\nProcessing Topic {topic_index}/{total_topics}: {topic.get('title', 'N/A')}")
        
        # 1. Generate Content for this topic (Call function from utils.tweet_generator)
        tweet_content_list = generate_content_for_topic(
            ai_agent, prompt_manager, config, topic, 
            base_output_dir, run_id, topic_index
        )
        
        # 2. Generate Posters for this topic (Call function from utils.tweet_generator)
        if tweet_content_list:
            # Pass only necessary parts of config? Or the whole config?
            # Pass config for now, the function extracts what it needs.
            generate_posters_for_topic(
                config, # Pass config
                topic, tweet_content_list, 
                base_output_dir, run_id, topic_index
                # Instances are now created inside generate_posters_for_topic
                # poster_gen_instance, content_gen, 
            )
        else:
            print(f"  Skipping poster generation for Topic {topic_index} due to content generation failure.")

    # --- Cleanup --- 
    if ai_agent:
        print("\nClosing shared AI Agent...")
        ai_agent.close()
    print("\nFinished processing all topics.")

def main():
    # No argparse for now, directly load default config
    config = load_config() # Load from poster_gen_config.json

    # Execute steps sequentially
    # Step 1: Generate Topics (using the function from utils.tweet_generator)
    run_id, tweet_topic_record = run_topic_generation_pipeline(config)

    # Step 2: Generate Content and Posters (if Step 1 was successful)
    if run_id and tweet_topic_record:
        generate_content_and_posters_step(config, run_id, tweet_topic_record)
    else:
        print("Exiting due to issues in topic generation.")

if __name__ == "__main__":
    main()