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
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 prepare_topic_generation, generate_topics, generate_single_content
import random

TEXT_POSBILITY = 0.3

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):
             print(f"Error: Config file '{config_path}' is missing one or more required keys.")
             print(f"Required keys are: {required_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)


def generate_topics_step(config):
    """Generates topics based on the configuration."""
    print("Step 1: Generating Topics...")
    ai_agent, system_prompt, user_prompt, base_output_dir = prepare_topic_generation(
        config.get("date", datetime.now().strftime("%Y-%m-%d")), # Use current date if not specified
        config["num"], config["topic_system_prompt"], config["topic_user_prompt"],
        config["api_url"], config["model"], config["api_key"], config["prompts_dir"],
        config["resource_dir"], config["output_dir"]
    )

    run_id, tweet_topic_record = generate_topics(
        ai_agent, system_prompt, user_prompt, config["output_dir"],
        config.get("topic_temperature", 0.2), config.get("topic_top_p", 0.5), config.get("topic_max_tokens", 1.5) # Added defaults for safety
    )

    if not run_id or not tweet_topic_record:
        print("Topic generation failed. Exiting.")
        ai_agent.close()
        return None, None

    output_dir = os.path.join(config["output_dir"], run_id)
    os.makedirs(output_dir, exist_ok=True)
    tweet_topic_record.save_topics(os.path.join(output_dir, "tweet_topic.json"))
    tweet_topic_record.save_prompt(os.path.join(output_dir, "tweet_prompt.txt"))
    ai_agent.close()
    print(f"Topics generated successfully. Run ID: {run_id}")
    return run_id, tweet_topic_record


def generate_content_and_posters_step(config, run_id, tweet_topic_record):
    """Generates content and posters based on generated topics."""
    if not run_id or not tweet_topic_record:
        print("Missing run_id or topics data. Skipping content and poster generation.")
        return

    print("Step 2: Generating Content and Posters...")
    base_output_dir = config["output_dir"]
    output_dir = os.path.join(base_output_dir, run_id) # Directory for this specific run

    # Load content generation system prompt
    content_system_prompt = ResourceLoader.load_system_prompt(config["content_system_prompt"])
    if not content_system_prompt:
        print("Warning: Content generation system prompt is empty. Using default logic if available or might fail.")
        # Potentially load topic system prompt as fallback if needed, or handle error
        # content_system_prompt = ResourceLoader.load_system_prompt(config["topic_system_prompt"])


    image_base_dir = config.get("image_base_dir", None)
    if not image_base_dir:
         print("Error: 'image_base_dir' not specified in config. Cannot locate images.")
         return
    camera_image_subdir = config.get("camera_image_subdir", "相机") # Default '相机'
    modify_image_subdir = config.get("modify_image_subdir", "modify") # Default 'modify'


    for i, topic in enumerate(tweet_topic_record.topics_list):
        topic_index = i + 1
        print(f"Processing Topic {topic_index}/{len(tweet_topic_record.topics_list)}: {topic.get('title', 'N/A')}")
        tweet_content_list = []

        # --- Content Generation Loop ---
        for j in range(config["variants"]):
            variant_index = j + 1
            print(f"  Generating Variant {variant_index}/{config['variants']}...")
            time.sleep(random.random()) # Keep the random delay? Okay for now.
            ai_agent = AI_Agent(config["api_url"], config["model"], config["api_key"])
            try:
                tweet_content, gen_result = generate_single_content(
                    ai_agent, content_system_prompt, topic,
                    config["prompts_dir"], config["resource_dir"],
                    output_dir, run_id, topic_index, variant_index, config.get("content_temperature", 0.3) # Added default
                )
                if tweet_content:
                    tweet_content_list.append(tweet_content.get_json_file()) # Assuming this returns the structured data needed later
                else:
                     print(f"  Failed to generate content for Topic {topic_index}, Variant {variant_index}. Skipping.")
            except Exception as e:
                print(f"  Error during content generation for Topic {topic_index}, Variant {variant_index}: {e}")
                traceback.print_exc()
            finally:
                ai_agent.close() # Ensure agent is closed

        if not tweet_content_list:
            print(f"  No content generated for Topic {topic_index}. Skipping poster generation.")
            continue

        # --- Poster Generation Setup ---
        object_name = topic.get("object", "")
        if not object_name:
            print("  Warning: Topic object name is missing. Cannot determine image path.")
            continue

        # Clean object name (consider making this a utility function)
        try:
            object_name = object_name.split(".")[0]
            if "景点信息-" in object_name:
                 object_name = object_name.split("景点信息-")[1]
            # Handle cases like "景点A+景点B"? Needs clearer logic if required.
        except Exception as e:
            print(f"  Warning: Could not fully clean object name '{object_name}': {e}")

        # Construct and check image paths using config base dir
        # Path for collage/poster input images (e.g., from 'modify' dir)
        input_img_dir_path = os.path.join(image_base_dir, modify_image_subdir, object_name)
        # Path for potential description file (e.g., from '相机' dir)
        camera_img_dir_path = os.path.join(image_base_dir, camera_image_subdir, object_name)
        description_file_path = os.path.join(camera_img_dir_path, "description.txt")


        if not os.path.exists(input_img_dir_path) or not os.path.isdir(input_img_dir_path):
            print(f"  Image directory not found or not a directory: '{input_img_dir_path}'. Skipping poster generation for this topic.")
            continue

        info_directory = []
        if os.path.exists(description_file_path):
             info_directory = [description_file_path]
             print(f"  Using description file: {description_file_path}")
        else:
             print(f"  Description file not found: '{description_file_path}'. Using generated content for poster text.")

        # --- Generate Text Configurations for Posters ---
        content_gen = contentGen.ContentGenerator()
        try:
            # Assuming tweet_content_list contains the JSON data needed by content_gen.run
            poster_text_configs_raw = content_gen.run(info_directory, config["variants"], tweet_content_list)
            print(f"  Raw poster text configs: {poster_text_configs_raw}") # For debugging
            poster_config_summary = posterGen.PosterConfig(poster_text_configs_raw)
        except Exception as e:
            print(f"  Error running ContentGenerator or parsing poster configs: {e}")
            traceback.print_exc()
            continue # Skip poster generation for this topic


        # --- Poster Generation Loop ---
        poster_num = config["variants"] # Same as content variants
        target_size = tuple(config.get("poster_target_size", [900, 1200])) # Add default size

        for j_index in range(poster_num):
            variant_index = j_index + 1
            print(f"    Generating Poster {variant_index}/{poster_num}...")
            try:
                poster_config = poster_config_summary.get_config_by_index(j_index)
                if not poster_config:
                    print(f"    Warning: Could not get poster config for index {j_index}. Skipping.")
                    continue

                variant_output_dir = os.path.join(output_dir, f"{topic_index}_{variant_index}")
                collage_output_dir = os.path.join(variant_output_dir, "collage_img")
                poster_output_dir = os.path.join(variant_output_dir, "poster")
                os.makedirs(collage_output_dir, exist_ok=True)
                os.makedirs(poster_output_dir, exist_ok=True)

                # --- Image Collage ---
                img_list = simple_collage.process_directory(
                    input_img_dir_path,
                    target_size=target_size,
                    output_count=1, # Assuming 1 collage image per poster variant
                    output_dir=collage_output_dir
                )
                # print(f"    Collage image list: {img_list}") # Debugging

                if not img_list or len(img_list) == 0 or not img_list[0].get('path'):
                    print(f"    Failed to generate collage image for Variant {variant_index}. Skipping poster.")
                    continue
                collage_img_path = img_list[0]['path']
                print(f"    Using collage image: {collage_img_path}")

                # --- Create Poster ---
                poster_gen_instance = posterGen.PosterGenerator() # Renamed to avoid conflict
                
                # Prepare text data (Simplified logic, adjust TEXT_POSBILITY if needed)
                # Consider moving text data preparation into posterGen or a dedicated function
                text_data = {
                    "title": poster_config.get('main_title', 'Default Title'),
                    "subtitle": "", # Subtitle seems unused?
                    "additional_texts": []
                }
                texts = poster_config.get('texts', [])
                if texts:
                     text_data["additional_texts"].append({"text": texts[0], "position": "bottom", "size_factor": 0.5})
                     if len(texts) > 1 and random.random() < TEXT_POSBILITY: # Apply possibility check
                          text_data["additional_texts"].append({"text": texts[1], "position": "bottom", "size_factor": 0.5})

                # print(f"    Text data for poster: {text_data}") # Debugging
                
                final_poster_path = os.path.join(poster_output_dir, "poster.jpg")
                result_path = poster_gen_instance.create_poster(collage_img_path, text_data, final_poster_path)
                if result_path:
                    print(f"    Successfully generated poster: {result_path}")
                else:
                    print(f"    Poster generation function did not return a valid path.")

            except Exception as e:
                print(f"    Error during poster generation for Variant {variant_index}: {e}")
                traceback.print_exc()
                continue # Continue to next variant

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

    # Execute steps sequentially
    run_id, tweet_topic_record = generate_topics_step(config)

    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()