#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import json
import time
from datetime import datetime
import logging

# --- Setup Project Root Path ---
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
    sys.path.append(PROJECT_ROOT)
# --- End Setup ---

# --- Imports from the project ---
try:
    from utils.tweet_generator import run_topic_generation_pipeline, generate_content_for_topic, generate_posters_for_topic
    from core.topic_parser import TopicParser
    from utils.prompt_manager import PromptManager # Needed for content generation
    from core.ai_agent import AI_Agent # Needed for content generation
    # from utils.tweet_generator import tweetTopicRecord # No longer needed directly
    from utils.output_handler import FileSystemOutputHandler, OutputHandler # Import handlers
except ImportError as e:
    logging.critical(f"ImportError: {e}. Ensure all core/utils modules are available and '{PROJECT_ROOT}' is in sys.path.")
    sys.exit(1)
# --- End Imports ---

def load_config(config_path):
    """Loads configuration from a JSON file."""
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        logging.info(f"Config loaded successfully from {config_path}")
        return config
    except FileNotFoundError:
        logging.error(f"Error: Configuration file not found at {config_path}")
        return None
    except json.JSONDecodeError:
        logging.error(f"Error: Could not decode JSON from {config_path}")
        return None
    except Exception as e:
        logging.exception(f"An unexpected error occurred loading config {config_path}:")
        return None

def main_test():
    # --- Basic Logging Setup --- 
    logging.basicConfig(
        level=logging.INFO, 
        format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    # --- End Logging Setup --- 

    logging.info("--- Starting Pipeline Steps Test ---")

    # 1. Load Configuration
    config_path = os.path.join(PROJECT_ROOT, "poster_gen_config.json") # Use project root path
    config = load_config(config_path)
    if config is None:
        logging.critical("Failed to load configuration. Exiting test.")
        sys.exit(1)

    # --- Initialize Output Handler --- 
    output_handler: OutputHandler = FileSystemOutputHandler(config.get("output_dir", "result"))
    logging.info(f"Using Output Handler: {output_handler.__class__.__name__}")
    # --- End Output Handler Init --- 

    # 2. Define a unique Run ID for this test run
    test_run_id = f"test_pipeline_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    logging.info(f"Using Run ID: {test_run_id}")

    # --- Step 1: Topic Generation --- 
    logging.info("\n--- Testing Step 1: Topic Generation ---")
    step1_start = time.time()
    # run_topic_generation_pipeline still takes config internally
    run_id_step1, topics_list, system_prompt, user_prompt = run_topic_generation_pipeline(config, run_id=test_run_id)
    step1_end = time.time()

    if run_id_step1 is None or topics_list is None:
        logging.error("Topic generation (Step 1) failed. Cannot proceed.")
        sys.exit(1)
    if run_id_step1 != test_run_id:
        logging.warning(f"run_id returned from step 1 ({run_id_step1}) differs from the one provided ({test_run_id}). Using the returned one.")
        test_run_id = run_id_step1 # Use the ID actually used by the pipeline

    logging.info(f"Step 1 finished in {step1_end - step1_start:.2f} seconds.")
    logging.info(f"Generated {len(topics_list)} topics for Run ID: {test_run_id}")
    # Use output handler to save topic results (mimics main.py)
    output_handler.handle_topic_results(test_run_id, topics_list, system_prompt, user_prompt)

    # --- Step 2: Content and Poster Generation --- 
    logging.info("\n--- Testing Step 2: Content and Poster Generation ---")
    step2_start = time.time()

    # Initialize resources needed for step 2, extracting params from config
    prompt_manager = None
    ai_agent = None
    step2_success_flag = False
    try:
        # --- Create PromptManager --- 
        prompt_manager = PromptManager(
            topic_system_prompt_path=config.get("topic_system_prompt"),
            topic_user_prompt_path=config.get("topic_user_prompt"),
            content_system_prompt_path=config.get("content_system_prompt"),
            prompts_dir=config.get("prompts_dir"),
            resource_dir_config=config.get("resource_dir", []),
            topic_gen_num=config.get("num", 1),
            topic_gen_date=config.get("date", "")
        )
        logging.info("PromptManager instance created for Step 2 test.")
        
        # --- Create AI Agent --- 
        ai_api_url = config.get("api_url")
        ai_model = config.get("model")
        ai_api_key = config.get("api_key")
        request_timeout = config.get("request_timeout", 30)
        max_retries = config.get("max_retries", 3)
        stream_chunk_timeout = config.get("stream_chunk_timeout", 60) # Read stream timeout
        if not all([ai_api_url, ai_model, ai_api_key]):
            raise ValueError("Missing required AI configuration (api_url, model, api_key)")
        logging.info("Initializing AI Agent for content generation test...")
        ai_agent = AI_Agent(
            base_url=ai_api_url,
            model_name=ai_model,
            api=ai_api_key,
            timeout=request_timeout,
            max_retries=max_retries,
            stream_chunk_timeout=stream_chunk_timeout # Pass stream timeout
        )

        total_topics = len(topics_list)
        logging.info(f"Processing {total_topics} topics for content/posters...")
        for i, topic_item in enumerate(topics_list):
            topic_index = topic_item.get('index', i + 1)
            logging.info(f"--- Processing Topic {topic_index}/{total_topics} ---")

            # --- Generate Content --- 
            content_variants = config.get("variants", 1)
            content_temp = config.get("content_temperature", 0.3)
            content_top_p = config.get("content_top_p", 0.4)
            content_presence_penalty = config.get("content_presence_penalty", 1.5)
            
            content_success = generate_content_for_topic(
                ai_agent=ai_agent, 
                prompt_manager=prompt_manager,
                topic_item=topic_item, 
                run_id=test_run_id, 
                topic_index=topic_index, 
                output_handler=output_handler,
                variants=content_variants,
                temperature=content_temp,
                top_p=content_top_p,
                presence_penalty=content_presence_penalty
            )

            if content_success:
                 logging.info(f"Content generation successful for Topic {topic_index}.")
                 # --- Generate Posters --- 
                 poster_variants = config.get("variants", 1)
                 poster_assets_dir = config.get("poster_assets_base_dir")
                 img_base_dir = config.get("image_base_dir")
                 mod_img_subdir = config.get("modify_image_subdir", "modify")
                 res_dir_config = config.get("resource_dir", [])
                 poster_size = tuple(config.get("poster_target_size", [900, 1200]))
                 txt_possibility = config.get("text_possibility", 0.3)
                 collage_subdir = config.get("output_collage_subdir", "collage_img")
                 poster_subdir = config.get("output_poster_subdir", "poster")
                 poster_filename = config.get("output_poster_filename", "poster.jpg")
                 cam_img_subdir = config.get("camera_image_subdir", "相机")
                 
                 if not poster_assets_dir or not img_base_dir:
                     logging.error(f"Missing critical paths for poster generation. Skipping posters for topic {topic_index}.")
                     continue
                     
                 posters_attempted = generate_posters_for_topic(
                     topic_item=topic_item, 
                     output_dir=config.get("output_dir", "result"),
                     run_id=test_run_id, 
                     topic_index=topic_index,
                     output_handler=output_handler,
                     variants=poster_variants,
                     poster_assets_base_dir=poster_assets_dir,
                     image_base_dir=img_base_dir,
                     modify_image_subdir=mod_img_subdir,
                     resource_dir_config=res_dir_config,
                     poster_target_size=poster_size,
                     text_possibility=txt_possibility,
                     output_collage_subdir=collage_subdir,
                     output_poster_subdir=poster_subdir,
                     output_poster_filename=poster_filename,
                     camera_image_subdir=cam_img_subdir
                 )
                 if posters_attempted:
                    logging.info(f"Poster generation process completed for Topic {topic_index}.")
                    step2_success_flag = True
                 else:
                    logging.warning(f"Poster generation skipped/failed for Topic {topic_index}.")
            else:
                logging.warning(f"Content generation failed for Topic {topic_index}. Skipping posters.")
            logging.info(f"--- Finished Topic {topic_index} ---")

    except ValueError as e:
        logging.error(f"Configuration error during Step 2 setup: {e}")
        step2_success_flag = False
    except Exception as e:
        logging.exception("An error occurred during Step 2 processing:")
        step2_success_flag = False # Ensure flag is false on error
    finally:
        if ai_agent:
            logging.info("Closing AI Agent for content generation test...")
            ai_agent.close()
    # --- End Simulated Step 2 Logic ---

    step2_end = time.time()
    if step2_success_flag:
        logging.info(f"Step 2 finished in {step2_end - step2_start:.2f} seconds.")
    else:
        logging.warning(f"Step 2 finished in {step2_end - step2_start:.2f} seconds, but encountered errors or generated no output.")

    # --- Finalize Output --- 
    if test_run_id:
        output_handler.finalize(test_run_id)
    # --- End Finalize --- 

    # --- Test Summary ---
    logging.info("\n--- Pipeline Steps Test Summary ---")
    logging.info(f"Run ID: {test_run_id}")
    output_location = os.path.join(config["output_dir"], test_run_id)
    logging.info(f"Check output files in: {output_location}")
    if os.path.exists(output_location):
        logging.info("Output directory exists.")
    else:
        logging.warning("Output directory NOT found.")

if __name__ == "__main__":
    main_test()