#1 GENERATE AND SAMPLE BEST_TOPIC
from datetime import datetime
import json
import os
from typing import List, Dict, Any
from utils import get_llm_response, extract_json_from_string
from utils import read_json_to_string

scoring_topics_template = read_json_to_string("./prompt_template/scoring_topics.json")
initializing_topics_template = read_json_to_string("./prompt_template/initializing_topics.json")

def initialize_topics(model_name: str, temperature: float) -> str:
    """
    Call the large model to generate 5 original novel worldview themes.
    Args:
        model_name: LLM model name
        temperature: Temperature parameter
    Returns:
        Theme string
    """
    prompt = f"""
    Please generate 5 novel worldview themes, including the book name and theme. The theme should be novel, attractive, and meet the following requirements:
    1. The theme must be original and unrelated to existing and published literary works/film works/game works/animation works/network novels, etc.
    2. The theme can be a completely fictional worldview, or a real worldview, or a combination of real worldview and fictional worldview, and can be imagined.
    3. The theme needs to show a certain depth and breadth of thinking.
    4. The theme needs to have a certain degree of innovation.
    Please strictly follow the following JSON format, and do not use double quotes in the value of JSON:
    ```json
    {initializing_topics_template}
    ```
    """
    try:
        response = get_llm_response(prompt, model_name, temperature)
        return extract_json_from_string(response)
    except Exception as e:
        print(f"Failed to initialize themes: {e}")
        return json.dumps([])

def scoring_topics(topics: str, model_name: str) -> str:
    """
    Score and reason the theme list.
    Args:
        topics: Theme string
        model_name: LLM model name
    Returns:
        JSON string with score and reason
    """
    criteria = f"""
    1. The theme must be original and unrelated to existing and published literary works/film works/game works/animation works/network novels, etc.
    2. The theme can be a completely fictional worldview, or a real worldview, or a combination of real worldview and fictional worldview, and can be imagined.
    3. The theme needs to show a certain depth and breadth of thinking.
    4. The theme needs to have a certain degree of innovation.
    """
    prompt = f"""
    The scoring range is 0-10, please score each theme in topics according to the following criteria, and give the scoring reason:
    {criteria}
    Here is the theme list:
    {topics}
    Please strictly follow the following JSON format, and do not use double quotes in the value of JSON:
    ```json
    {scoring_topics_template}
    ```
    """
    try:
        response = get_llm_response(prompt, model_name)
        return extract_json_from_string(response)
    except Exception as e:
        print(f"Failed to score themes: {e}")
        return json.dumps([])

def extract_json_from_response(response: str) -> List[Dict[str, Any]]:
    """
    Extract the json part from the string returned by LLM and parse it.
    Args:
        response: String returned by LLM
    Returns:
        List of themes and scores
    """
    try:
        json_str = response.split("```json")[1].split("```", 1)[0]
        data = json.loads(json_str)
        # compatible with score being a string
        for item in data:
            if isinstance(item["score"], str):
                item["score"] = int(item["score"])
        return data
    except Exception as e:
        print(f"Failed to parse the content returned by LLM: {e}")
        return []

def get_topic_with_score(topics_with_score: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Select the theme with the highest score.
    Args:
        topics_with_score: List of themes and scores
    Returns:
        Dictionary of the theme with the highest score
    """
    if not topics_with_score:
        return {}
    return max(topics_with_score, key=lambda x: x["score"])

def save_to_file(data: Any, filename: str) -> None:
    """
    Save data to file.
    Args:
        data: Data to save
        filename: File name
    Returns:
        None
    """
    with open(filename, 'w', encoding='utf-8') as f:
        if isinstance(data, (dict, list)):
            json.dump(data, f, ensure_ascii=False, indent=2)
        else:
            f.write(str(data))

def extract_topics(topics_json: str) -> List[str]:
    """
    Extract all topic fields from the theme JSON string and return a list.
    Args:
        topics_json: JSON string containing theme information
    Returns:
        List of all topics
    """
    try:
        if isinstance(topics_json, str):
            topics_data = json.loads(topics_json)
        else:
            topics_data = topics_json  # if it is already a list or dictionary, use it directly
        
        topics_list = []
        for item in topics_data:
            if "topic" in item:
                topics_list.append(item["topic"])
        
        return topics_list
    except Exception as e:
        print(f"Failed to extract topics: {e}")
        return []

def generate_topic(model: str, temperature: float) -> str:
    """
    Generate theme
    Args:
        model: Model name
        temperature: Temperature coefficient
    Returns:
        best_topic
    """
    print(f"Using model: {model}")
    print(f"Temperature coefficient for generating theme: {temperature}")
    # Initialize timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    try:
        topics = initialize_topics(model, temperature)
        print("\nGenerated themes:")
        print(topics)
        # Create folders
        os.makedirs("raw_topics", exist_ok=True)
        os.makedirs("scored_topics", exist_ok=True)
        os.makedirs("best_topic", exist_ok=True)
        
        save_to_file(topics, f"raw_topics/topics_raw_{timestamp}.json")

        # Extract topic list
        topics_list = extract_topics(topics)
        topics_str = json.dumps(topics_list, ensure_ascii=False)
        
        scored_topics = scoring_topics(topics_str, model)
        print("\nScoring results:")
        print(scored_topics)
        save_to_file(scored_topics, f"scored_topics/topics_scored_{timestamp}.json")

        best_topic = get_topic_with_score(scored_topics)
        print("\nThe theme with the highest score:")
        print(json.dumps(best_topic, ensure_ascii=False, indent=4))
        save_to_file(best_topic, f"best_topic/best_topic_{timestamp}.json")
        # Extract topic from json
        topic = best_topic["topic"]
        return topic
    except Exception as e:
        print(f"An error occurred during execution: {e}")
        return None



if __name__ == "__main__":
    best_topic = generate_topic(model="claude-3-7-sonnet-20250219", temperature=1.0)
    print(best_topic)