import os
from openai import OpenAI
import re # Import regex for parsing conversation turns
from typing import Optional, Union # Need Optional for settings

# Ensure the OPENROUTER_API_KEY environment variable is set
api_key = "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"
if not api_key:
    raise ValueError("OPENROUTER_API_KEY environment variable not set.")

# Point the OpenAI client to the OpenRouter API
client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=api_key,
)

# --- Core Generation Functions ---

def generate_synthetic_text(
    prompt: str,
    model: str = "deepseek/deepseek-chat-v3-0324:free",
    system_message: str = "You are a helpful assistant generating synthetic data.",
    temperature: Optional[float] = 0.7, # Default temperature
    top_p: Optional[float] = None,       # Default top_p (let API decide if None)
    max_tokens: Optional[int] = None   # Default max_tokens (let API decide if None)
) -> str:
    """
    Generates synthetic text using an OpenRouter model via Chat Completions,
    including model parameter controls.

    Args:
        prompt: The user's input prompt.
        model: The model ID.
        system_message: The system message context.
        temperature: Controls randomness (0.0 to 2.0). None means API default.
        top_p: Nucleus sampling probability. None means API default.
        max_tokens: Maximum number of tokens to generate. None means API default.

    Returns:
        The generated text string or an error message.
    """
    if not api_key or api_key == "YOUR_API_KEY_HERE_OR_SET_ENV_VAR":
         return "Error: OPENROUTER_API_KEY not configured properly. Please set the environment variable."

    # Prepare parameters, only including them if they are not None
    params = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_message},
                {"role": "user", "content": prompt},
            ],
        "extra_headers": {
             # "HTTP-Referer": "YOUR_SITE_URL",
             "X-Title": "SynthGen",
         }
    }
    if temperature is not None:
        params["temperature"] = temperature
    if top_p is not None:
        params["top_p"] = top_p
    if max_tokens is not None:
        params["max_tokens"] = max_tokens

    try:
        response = client.chat.completions.create(**params) # Use dictionary unpacking

        if response.choices and response.choices[0].message and response.choices[0].message.content:
            return response.choices[0].message.content.strip()
        else:
            print(f"Warning: No content in response for model {model}. Response: {response}")
            return "Error: No content generated by the model."
    except Exception as e:
        print(f"Error during API call to model {model}: {e}")
        return f"Error during API call: {e}"

def generate_prompts(
    num_prompts: int,
    model: str,
    topic_hint: str = "diverse and interesting",
    temperature: Optional[float] = 0.7, # Pass settings through
    top_p: Optional[float] = None,
    max_tokens: Optional[int] = 200 # Set a reasonable default max for prompts
) -> list[str]:
    """
    Generates a list of conversation prompts using an AI model.

    Args:
        num_prompts: The number of prompts to generate.
        model: The model ID to use for generation.
        topic_hint: Optional hint for the kind of topics (e.g., "related to technology").
        temperature: Controls randomness (0.0 to 2.0). None means API default.
        top_p: Nucleus sampling probability. None means API default.
        max_tokens: Maximum number of tokens to generate. None means API default.

    Returns:
        A list of generated prompts.
    """
    instruction = (
        f"Generate exactly {num_prompts} unique, {topic_hint} system prompts or starting topics suitable "
        f"for generating synthetic conversations between a user and an AI assistant. "
        f"Each prompt should be concise (ideally one sentence) and focus on a clear task or subject. "
        f"Present each prompt on a new line, with no other introductory or concluding text."
        f"\n\nExamples:\n"
        f"- Act as a travel agent planning a trip to Japan.\n"
        f"- Explain the concept of black holes to a 5-year-old.\n"
        f"- Write a python function to reverse a string."
    )
    system_msg = "You are an expert prompt generator. Follow the user's instructions precisely."

    # Pass the settings down to generate_synthetic_text
    generated_text = generate_synthetic_text(
        instruction,
        model,
        system_message=system_msg,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens
    )

    if generated_text.startswith("Error:"):
        raise ValueError(generated_text)

    # Split into lines and clean up any extra whitespace or empty lines
    prompts = [p.strip() for p in generated_text.strip().split('\n') if p.strip()]
    prompts = [p.replace("- ", "") for p in prompts]
    if not prompts:
        # Log the raw generated text if parsing failed
        print(f"Warning: Failed to parse prompts from generated text. Raw text:\n{generated_text}")
        raise ValueError("AI failed to generate prompts in the expected format.")

    # Optional: Truncate or pad if the model didn't generate the exact number
    return prompts[:num_prompts]


def generate_synthetic_conversation(
    system_prompt: str,
    model: str,
    num_turns: int,
    temperature: Optional[float] = 0.7, # Pass settings through
    top_p: Optional[float] = None,
    max_tokens: Optional[int] = 1000 # Set a reasonable default max for conversations
) -> str:
    """
    Generates a synthetic conversation with a specified number of turns.

    Args:
        system_prompt: The initial system prompt defining the context or AI persona.
        model: The model ID to use for generation.
        num_turns: The desired number of conversational turns (1 turn = 1 User + 1 Assistant).
        temperature: Controls randomness (0.0 to 2.0). None means API default.
        top_p: Nucleus sampling probability. None means API default.
        max_tokens: Maximum number of tokens to generate. None means API default.

    Returns:
        A string containing the formatted conversation.
    """
    # We'll ask the model to generate the whole conversation in one go for simplicity.
    # More complex approaches could involve iterative calls.
    instruction = (
        f"Generate a realistic conversation between a 'User' and an 'Assistant'. "
        f"The conversation should start based on the following system prompt/topic: '{system_prompt}'.\n"
        f"The conversation should have approximately {num_turns} pairs of User/Assistant turns.\n"
        f"Format the output clearly, starting each line with 'User:' or 'Assistant:'.\n\n"
        f"Example Format:\n"
        f"User: Hello!\n"
        f"Assistant: Hi there! How can I help you today?\n"
        f"User: Can you explain photosynthesis?\n"
        f"Assistant: Certainly! Photosynthesis is the process..."
    )

    # Use the user-provided system prompt for the *conversation's* context,
    # but a generic one for the generation *task* itself.
    system_msg_for_generation = f"You are an AI assistant simulating a conversation. The context for the conversation you generate is: {system_prompt}"

    # Pass the settings down to generate_synthetic_text
    conversation_text = generate_synthetic_text(
        prompt=instruction,
        model=model,
        system_message=system_msg_for_generation,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens
    )

    if conversation_text.startswith("Error:"):
        # Propagate the error message
        return f"Error generating conversation for prompt '{system_prompt}':\n{conversation_text}"

    # Basic validation/cleanup (optional)
    if not re.search(r"User:|Assistant:", conversation_text, re.IGNORECASE):
         print(f"Warning: Generated text for conversation '{system_prompt}' might not be in the expected format. Raw text:\n{conversation_text}")
         # Return the raw text anyway, maybe the model format is slightly different
         return f"Generated conversation for prompt '{system_prompt}':\n(Format might vary)\n\n{conversation_text}"

    return f"Generated conversation for prompt '{system_prompt}':\n\n{conversation_text}"

# Function to generate different types of content based on a topic
def generate_corpus_content(
    topic: str,
    content_type: str, # e.g., "Corpus Snippets", "Short Story", "Article"
    length_param: int, # Meaning depends on type (e.g., num snippets, approx words)
    model: str,
    system_message_base: str = "You are a helpful assistant generating synthetic content.",
    temperature: Optional[float] = 0.7,
    top_p: Optional[float] = None,
    max_tokens: Optional[int] = None # Use a larger default if None
) -> str:
    """
    Generates different types of synthetic content based on a topic.

    Args:
        topic: The central topic for the content.
        content_type: The type of content to generate.
        length_param: A parameter controlling length/quantity (meaning depends on type).
        model: The model ID.
        system_message_base: Base system message (will be specialized).
        temperature: Model temperature.
        top_p: Model top_p.
        max_tokens: Model max_tokens.

    Returns:
        The generated content string or an error message.
    """

    prompt = ""
    system_message = system_message_base # Start with base

    # --- Construct Prompt based on Content Type ---
    if content_type == "Corpus Snippets":
        if length_param <= 0: length_param = 5 # Default number of snippets
        prompt = (
            f"Generate exactly {length_param} distinct text snippets related to the topic: '{topic}'. "
            f"Each snippet should be a few sentences long and focus on a different aspect if possible. "
            f"Present each snippet clearly, perhaps separated by a blank line or a marker like '---'."
        )
        system_message = "You are an AI generating diverse text snippets for a data corpus."
        # Adjust max_tokens based on expected number of snippets if not set
        if max_tokens is None: max_tokens = length_param * 150 # Estimate

    elif content_type == "Short Story":
        if length_param <= 0: length_param = 300 # Default approx words
        prompt = (
            f"Write a short story (approximately {length_param} words) centered around the topic: '{topic}'. "
            f"The story should have a clear beginning, middle, and end."
        )
        system_message = "You are a creative AI writing a short story."
        if max_tokens is None: max_tokens = int(length_param * 2.5) # Estimate

    elif content_type == "Article":
        if length_param <= 0: length_param = 500 # Default approx words
        prompt = (
            f"Write an informative article (approximately {length_param} words) about the topic: '{topic}'. "
            f"The article should be well-structured, factual (to the best of your ability), and engaging."
        )
        system_message = "You are an AI assistant writing an informative article."
        if max_tokens is None: max_tokens = int(length_param * 2.5) # Estimate

    else:
        return f"Error: Unknown content type '{content_type}'."

    if not prompt:
         return "Error: Could not construct a valid prompt."

    # --- Call the core generation function ---
    generated_text = generate_synthetic_text(
        prompt=prompt,
        model=model,
        system_message=system_message,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens
    )

    # Return the result (includes potential errors from generate_synthetic_text)
    # Add a title for clarity
    if not generated_text.startswith("Error:"):
         return f"Generated {content_type} for topic '{topic}':\n\n{generated_text}"
    else:
         return generated_text # Propagate the error

# --- Main Execution (Example Usage) ---
if __name__ == "__main__":
    print("--- Testing Basic Text Generation ---")
    test_prompt = "Describe the benefits of using synthetic data."
    text_result = generate_synthetic_text(test_prompt, temperature=0.5, max_tokens=100) # Example with settings
    print(f"Prompt: {test_prompt}\nResult:\n{text_result}\n")

    print("\n--- Testing Prompt Generation ---")
    try:
        num_prompts_to_gen = 3
        prompts_result = generate_prompts(num_prompts_to_gen, "deepseek/deepseek-chat-v3-0324:free")
        print(f"Generated {len(prompts_result)} prompts:")
        for i, p in enumerate(prompts_result):
            print(f"{i+1}. {p}")
    except ValueError as e:
        print(f"Error generating prompts: {e}")

    print("\n--- Testing Conversation Generation ---")
    conv_prompt = "Act as a helpful expert explaining the difference between nuclear fission and fusion."
    num_conv_turns = 3
    conv_result = generate_synthetic_conversation(conv_prompt, "deepseek/deepseek-chat-v3-0324:free", num_conv_turns)
    print(f"{conv_result}\n")

    print("\n--- Testing with Invalid API Key (if applicable) ---")
    # Temporarily use an invalid key for testing error handling
    original_key = client.api_key
    client.api_key = "invalid-key"
    error_text_result = generate_synthetic_text("Test prompt")
    print(f"Result with invalid key: {error_text_result}")
    client.api_key = original_key # Restore original key

    print("\nGeneration tests complete.")