import gradio as gr
import os
from groq import Groq
import markdown

# Retrieve system prompts from environment variables
GENERATE_PROMPT = os.environ.get("GENERATE_CONTENT_2")
HUMANIZE_PROMPT = os.environ.get("HUMANIZE_CONTENT_2")

def generate_mba_content(topic, api_key):
    if not GENERATE_PROMPT:
        return "Error: Generate prompt not found. Ensure 'GENERATE_PROMPT' is set in your environment."
    
    try:
        client = Groq(api_key=api_key)
    except Exception as e:
        return f"Error: Failed to initialize Groq client. Check your API key. Details: {str(e)}"

    prompt = GENERATE_PROMPT.replace("[TOPIC]", topic)
    try:
        response = client.chat.completions.create(
            model="llama3-70b-8192",
            messages=[
                {"role": "system", "content": prompt},
                {"role": "user", "content": f"Generate content for the topic: {topic}"}
            ],
            temperature=0.7,
            max_tokens=4000,
        )
        content = response.choices[0].message.content
        content = content.replace("—", ",")  # Replace em dashes with commas
        if not content.startswith("#"):
            content = markdown.markdown(content)
        return content
    except Exception as e:
        return f"Error: Failed to generate content. Details: {str(e)}"
    finally:
        api_key = None
        if 'client' in locals():
            del client

# def humanize_text(text, api_key):
#     if not HUMANIZE_PROMPT:
#         return "Error: Humanize prompt not found. Ensure 'HUMANIZE_PROMPT' is set in your environment."
    
#     try:
#         client = Groq(api_key=api_key)
#     except Exception as e:
#         return f"Error: Failed to initialize Groq client. Check your API key. Details: {str(e)}"

#     prompt = HUMANIZE_PROMPT.replace("[TEXT]", text)
#     try:
#         response = client.chat.completions.create(
#             model="llama3-70b-8192",
#             messages=[
#                 {"role": "system", "content": prompt},
#                 {"role": "user", "content": f"Rewrite the following text: {text}"}
#             ],
#             temperature=0.7,
#             max_tokens=4000,
#         )
#         content = response.choices[0].message.content
#         content = content.replace("—", ",")  # Replace em dashes with commas
#         if not content.startswith("#"):
#             content = markdown.markdown(content)
#         return content
#     except Exception as e:
#         return f"Error: Failed to humanize text. Details: {str(e)}"
#     finally:
#         api_key = None
#         if 'client' in locals():
#             del client

def humanize_text(text, api_key):
    SYSTEM_PROMPT = "You are a MBA teacher as well as a content writer. Your task is to rewrite provided content to sound more natural and human while following instructions."
    # HUMANIZE_PROMPT = os.environ.get("HUMANIZE_CONTENT_2")  # Ex: "Rewrite the following text to be more human: [TEXT]"

    if not HUMANIZE_PROMPT:
        return "Error: Humanize prompt not found. Ensure 'HUMANIZE_PROMPT' is set in your environment."

    try:
        client = Groq(api_key=api_key)
    except Exception as e:
        return f"Error: Failed to initialize Groq client. Check your API key. Details: {str(e)}"

    # Place the input text into the prompt template for the user message
    user_prompt = HUMANIZE_PROMPT.replace("[TEXT]", text)
    try:
        response = client.chat.completions.create(
            model="llama3-70b-8192",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt}
            ],
            temperature=0.7,
            max_tokens=4000,
        )
        content = response.choices[0].message.content
        content = content.replace("—", ",")
        if not content.startswith("#"):
            content = markdown.markdown(content)
        return content
    except Exception as e:
        return f"Error: Failed to humanize text. Details: {str(e)}"
    finally:
        api_key = None
        if 'client' in locals():
            del client

# Define Gradio interface
def gradio_app():
    with gr.Blocks(title="MBA Content Tools") as app:
        gr.Markdown("# MBA Content Tools")
        gr.Markdown("Generate or humanize content in a professional MBA-style format.")
        
        api_key_input = gr.Textbox(label="Groq API Key", type="password", placeholder="Enter your Groq API key here")
        
        with gr.Tab("Generate Content from Topic"):
            topic_input = gr.Textbox(label="Topic", placeholder="e.g., Strategic Management")
            generate_btn = gr.Button("Generate Content")
            generate_output = gr.Markdown(label="Generated Content")
            generate_btn.click(generate_mba_content, inputs=[topic_input, api_key_input], outputs=generate_output)
        
        with gr.Tab("Humanize Existing Text"):
            text_input = gr.TextArea(label="Text to Humanize", placeholder="Paste your article, report, or paragraph here")
            humanize_btn = gr.Button("Humanize Text")
            humanize_output = gr.Markdown(label="Humanized Content")
            humanize_btn.click(humanize_text, inputs=[text_input, api_key_input], outputs=humanize_output)
        
        gr.Markdown("""
        **Note:** Your API key is used securely for this session and cleared from memory afterward. 
        Ensure the environment variables 'GENERATE_PROMPT' and 'HUMANIZE_PROMPT' are set securely in your deployment configuration.
        """)
    
    return app

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