import os
from groq import Groq
import gradio as gr
from gtts import gTTS
import re

# -------------------------------------------------------------
# GET API KEY
# -------------------------------------------------------------
groq_key = os.environ.get("Groq_API_Key")
if not groq_key:
    raise ValueError("Groq_API_Key not found in environment variables. Please add it to your Hugging Face Space secrets.")

client = Groq(api_key=groq_key)

# -------------------------------------------------------------
# STORY GENERATION LOGIC
# -------------------------------------------------------------
def generate_story(age, theme, language):
    system_msg = {
        "role": "system",
        "content": "You write age-appropriate stories with complete endings and a moral."
    }

    age_i = int(age)
    if age_i <= 6:
        length = "300 words"
    elif age_i <= 9:
        length = "600 words"
    else:
        length = "1000 words"

    user_msg = {
        "role": "user",
        "content": (
            f"Write a complete story for a {age}-year-old child. "
            f"Theme: {theme}. Language: {language}. Length about {length}. "
            "End with a clear moral of the story."
        )
    }

    response = client.chat.completions.create(
        messages=[system_msg, user_msg],
        model="llama-3.3-70b-versatile",
        max_tokens=2000,
        temperature=0.8
    )

    return response.choices[0].message.content.strip()

# -------------------------------------------------------------
# TEXT TO SPEECH
# -------------------------------------------------------------
def story_to_speech(story, language_code="en"):
    tts = gTTS(text=story, lang=language_code)
    filename = "story.mp3"
    tts.save(filename)
    return filename

# -------------------------------------------------------------
# PROCESS STORY + AUDIO
# -------------------------------------------------------------
def create_story_interface(age, theme, language, tts_option):
    story = generate_story(age, theme, language)
    story = re.sub(r'[*_#`~]', '', story)

    lines = story.split("\n", 1)
    if len(lines) > 1:
        title, rest_of_story = lines[0], lines[1]
    else:
        title, rest_of_story = lines[0], ""

    if language.lower().startswith("u"):
        story_html = (
            f'<div style="text-align:center; font-size:1.3em; direction:rtl; margin-bottom:5px;">{title}</div>'
            f'<div style="font-size:1.2em; direction:rtl; text-align:right;">{rest_of_story}</div>'
        )
        lang_code = "ur"
    else:
        story_html = (
            f'<div style="text-align:center; font-size:1.3em; direction:ltr; margin-bottom:5px;">{title}</div>'
            f'<div style="font-size:1.2em; direction:ltr; text-align:left;">{rest_of_story}</div>'
        )
        lang_code = "en"

    audio_file = story_to_speech(story, lang_code) if tts_option else None
    return story_html, audio_file

# -------------------------------------------------------------
# UI OPTIONS
# -------------------------------------------------------------
age_options = [str(i) for i in range(3, 13)]
theme_options = [
    "Adventure", "Animals", "Fantasy", "Educational",
    "Friendship", "Magic", "Science", "Mystery",
    "Space", "Nature", "Kindness", "Courage"
]
language_options = ["English", "Urdu"]

# -------------------------------------------------------------
# CUSTOM CSS
# -------------------------------------------------------------
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Comic+Neue:wght@400;700&display=swap');
.gradio-container {
    font-family: 'Comic Neue', cursive !important;
    background: linear-gradient(135deg, #ffafbd, #ffc3a0) !important;
    padding-top: 10px !important;
    margin-top: -70px !important;
}
.main-header {
    text-align: center;
    font-weight: bold !important;
    font-size: 3.2em !important;
    margin-top: 40px !important;
    margin-bottom: 20px !important;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
    color: #12979B;
}
.story-container {
    background:#fff5ee;
    border-radius: 20px;
    padding: 25px;
    border: 5px solid #FFD166;
    box-shadow: 0 8px 25px rgba(0,0,0,0.1);
    min-height: 400px;
    margin-top: -20px !important;
}
.control-panel {
    background: #fff5ee;
    border-radius: 15px;
    padding: 20px;
    border: 3px dashed #4ECDC4;
    font-size: 1.2em !important;
    font-weight: 600;
    margin-top: -20px !important;
}
.generate-btn {
    background: linear-gradient(45deg, #FF6B6B, #FF9E7D) !important;
    border: none !important;
    border-radius: 50px !important;
    padding: 10px 22px !important;
    font-size: 1em !important;
    font-weight: bold !important;
    color: white !important;
    margin-top: 20px !important;
}
.generate-btn:hover {
    transform: scale(1.05);
    box-shadow: 0 5px 15px rgba(255, 107, 107, 0.4);
}
.dropdown, .checkbox {
    background: white !important;
    border-radius: 10px !important;
    border: 2px solid #96CEB4 !important;
}
.audio-player {
    border-radius: 15px !important;
    margin-top: 20px;
}
.subtitle {
    text-align: left;
    font-size: 1.5em !important;
    font-weight: 600;
    color: white;
    margin-top: -10px !important;
    margin-bottom: 20px;
    text-shadow: 1px 1px 2px rgba(0,0,0,0.1);
}
.footer {
    text-align: center;
    margin-top: 20px;
    color: #666;
    font-style: italic;
    font-size: 1.5em !important;
}
.main-header + div {
    margin-top: -145px !important;
}
"""

# -------------------------------------------------------------
# BUILD GRADIO INTERFACE
# -------------------------------------------------------------
with gr.Blocks() as iface:  # Removed theme argument
    gr.HTML(f"<style>{custom_css}</style>")

    gr.HTML('<div class="main-header">🌈 StoryTime AI🎪</div>')
    gr.HTML('<div class="subtitle">🎯 Create amazing stories that spark imagination and teach valuable lessons!</div>')

    with gr.Row():
        # LEFT PANEL
        with gr.Column(scale=1, min_width=300):
            with gr.Group(elem_classes="control-panel"):
                gr.Markdown("### 🎨 Story Settings")

                with gr.Group():
                    gr.Markdown("👶 *Child's Age*")
                    age_input = gr.Dropdown(age_options, value="5", show_label=False)

                with gr.Group():
                    gr.Markdown("🎭 *Story Theme*")
                    theme_input = gr.Dropdown(theme_options, value="Adventure", show_label=False)

                with gr.Group():
                    gr.Markdown("🌍 *Language*")
                    language_input = gr.Dropdown(language_options, value="English", show_label=False)

                with gr.Group():
                    gr.Markdown("🔊 *Audio Options*")
                    tts_input = gr.Checkbox(label="🎵 Include Audio Story", value=True)

                generate_btn = gr.Button("✨ Create Magical Story! ✨", elem_classes="generate-btn")

        # RIGHT PANEL
        with gr.Column(scale=2):
            with gr.Group(elem_classes="story-container"):
                story_output = gr.Markdown(
                    "Your story will appear here! Choose settings and click the magic button! 🎇",
                    show_label=False
                )

            audio_output = gr.Audio(elem_classes="audio-player", show_label=False)

    gr.HTML('<div class="footer">Made with ❤ for young readers everywhere! | Watch your stories come to life! ✨</div>')

    generate_btn.click(
        create_story_interface,
        [age_input, theme_input, language_input, tts_input],
        [story_output, audio_output]
    )

# -------------------------------------------------------------
# LAUNCH APP
# -------------------------------------------------------------
if __name__ == "__main__":
    iface.launch()