import gradio as gr
from email_generator.main import loop_email_workflow, EMAIL_EVALUATOR_PROMPT, EMAIL_GENERATOR_PROMPT
import json

# Function to generate the email
def generate_email_workflow(persona_json: str, campaign_json: str, sender_json: str, max_attempts: int, openai_api_key: str):
    """
    Generate a complete email with persona, campaign, and sender details.

    Args:
        persona_json (str): A JSON string representing the persona.
        campaign_json (str): A JSON string representing the campaign details.
        sender_json (str): A JSON string representing the sender details.
        max_attempts (int): Maximum number of attempts for generating a valid email.
        openai_api_key (str): The API key for OpenAI, if applicable.

    Returns:
        str: The complete generated email or an error message.
    """
    try:
        # Parse JSON strings to dictionaries
        persona = json.loads(persona_json)
        campaign = json.loads(campaign_json)
        sender = json.loads(sender_json)

        # Determine the model to use based on the API key
        use_huggingface = not bool(openai_api_key)
        model_used = "HuggingFace (Zephyr-7B)" if use_huggingface else "OpenAI (gpt-3.5-turbo)"

        # Run the email generation workflow
        result = loop_email_workflow(
            persona=persona,
            campaign=campaign,
            sender_data=sender,
            evaluator_prompt=EMAIL_EVALUATOR_PROMPT,
            generator_prompt=EMAIL_GENERATOR_PROMPT,
            max_tries=max_attempts,
            use_huggingface=use_huggingface,
            openai_api_key=openai_api_key if not use_huggingface else None,
        )

        if not result["final_email"]:
            return f"Failed to generate a valid email after {max_attempts} attempts. Feedback: {result.get('message', 'No additional information.')}\n\nModel Used: {model_used}"

        # Add sender information to the email content
        generated_email = result["final_email"]

        return generated_email

    except json.JSONDecodeError:
        return "Invalid JSON format. Please ensure all inputs are valid JSON."
    except Exception as e:
        return f"Error: {e}"


# Create Gradio interface
persona_input = gr.Textbox(
    label="Enter Persona (JSON format)",
    lines=10,
    value='{"name": "John", "city": "New York", "hobbies": "Reading"}',
    placeholder='{"name": "John", "city": "New York", "hobbies": "Reading"}'
)
campaign_input = gr.Textbox(
    label="Enter Campaign Details (JSON format)",
    lines=10,
    value='{"subject_line": "Discover Our New Product!", "product": "Backpacks", "discount": "20%", "validity": "Until January 31, 2025"}',
    placeholder='{"subject_line": "Discover Our New Product!", "product": "Backpacks", "discount": "20%", "validity": "Until January 31, 2025"}'
)
sender_input = gr.Textbox(
    label="Enter Sender Details (JSON format)",
    lines=5,
    value='{"name": "Jane Doe", "company": "Outdoor Gear Co."}',
    placeholder='{"name": "Jane Doe", "company": "Outdoor Gear Co.", "cta_text": "Shop Now", "cta_link": "https://example.com"}'
)
max_attempts_input = gr.Slider(
    label="Max Attempts",
    minimum=1,
    maximum=10,
    step=1,
    value=3,
    interactive=True
)
openai_api_key_input = gr.Textbox(
    label="Enter OpenAI API Key (Leave blank to use HuggingFace Zephyr-7B Beta)",
    type="password",
    placeholder="sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
)
email_output = gr.Textbox(
    label="Generated Email",
    lines=15,
    interactive=False,
)

# Interface layout
with gr.Blocks() as interface:
    gr.Markdown(
        """
        # Personalized Email Generator
        Generate a personalized email based on user persona and campaign details. 
        Provide the inputs in JSON format and specify the maximum number of attempts for generation.

        ### Available Models:
        - **OpenAI (gpt-3.5-turbo)**: A highly advanced language model known for its accuracy and contextual understanding, ideal for generating professional and creative emails.
        - **HuggingFace Zephyr-7B Beta**: An open-source model optimized for text generation tasks, offering a cost-effective alternative to proprietary APIs.
        """
    )
    with gr.Row():
        with gr.Column():
            persona_input.render()
            campaign_input.render()
            sender_input.render()
            max_attempts_input.render()
            openai_api_key_input.render()
        with gr.Column():
            email_output.render()

    generate_button = gr.Button("Generate Email")
    generate_button.click(
        fn=generate_email_workflow,
        inputs=[persona_input, campaign_input, sender_input, max_attempts_input, openai_api_key_input],
        outputs=email_output,
    )

# Launch the app
if __name__ == "__main__":
    interface.launch()