# kukubuddy_ai_light.py

import gradio as gr
from transformers import pipeline
from TTS.api import TTS

# Small & memory-efficient models
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
story_gen = pipeline("text2text-generation", model="google/flan-t5-small")
tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=False)

# ---------- Feature 1: Daily Audio Digest ----------
def generate_audio_digest(topic):
    dummy_text = f"This is a sample Kuku FM podcast about {topic}. " * 10
    summary = summarizer(dummy_text, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
    
    audio_path = "digest.wav"
    tts.tts_to_file(text=summary, file_path=audio_path)
    
    return summary, audio_path

# ---------- Feature 2: Interactive Story Generator ----------
story_cache = {}

# def generate_story(genre, choice):
#     base_prompt = f"Start a short {genre} story. "
    
#     if genre not in story_cache:
#         story_cache[genre] = base_prompt
    
#     if choice:
#         story_cache[genre] += f"User chose: {choice}. Then, "
    
#     generated = story_gen(story_cache[genre])[0]['generated_text']
#     story_cache[genre] = generated
    
#     story_audio_path = "story.wav"
#     tts.tts_to_file(text=generated, file_path=story_audio_path)
    
#     return generated, story_audio_path

def generate_story(genre, scene_input):
    # Define genre-specific styles
    genre_styles = {
        "action": "Include fast-paced sequences, intense confrontations, and physical conflict. Add thrilling suspense and end with a dramatic cliffhanger.",
        "romance": "Highlight emotional tension, subtle gestures, and deep connection between characters. Use heartfelt dialogue and end with a tender or surprising moment.",
        "comedy": "Use witty dialogue, humorous misunderstandings, and light-hearted narration. Keep the tone playful and funny.",
        "thriller": "Add psychological tension, mysterious clues, and an unsettling atmosphere. End with a chilling or puzzling moment.",
        "fantasy": "Include magical elements, otherworldly creatures, and vivid scenery. Use epic tone and mystical storytelling.",
        "horror": "Use eerie descriptions, building dread, and jump-scare moments. End on a terrifying note.",
        "sci-fi": "Incorporate futuristic tech, space travel, or AI. Build an immersive setting and challenge the reader’s imagination.",
        "drama": "Emphasize emotional depth, personal conflict, and realistic dialogue. Add introspective moments.",
    }

    # Fallback style if genre not found
    style = genre_styles.get(genre.lower(), "Write in an immersive and engaging tone relevant to the genre.")

    # Build the prompt
    prompt = (
        f"[INST] You are a brilliant storyteller. Write a vivid and immersive scene in a {genre} story. "
        f"The scene starts like this: {scene_input}. {style} [/INST]"
    )

    # Generate the story using your pipeline (story_gen)
    response = story_gen(prompt, max_length=300, do_sample=True, temperature=0.8, top_p=0.95)[0]['generated_text']

    # Convert to speech
    story_audio_path = "story.wav"
    tts.tts_to_file(text=response, file_path=story_audio_path)

    return response, story_audio_path


# ---------- Gradio UI ----------

digest_ui = gr.Interface(
    fn=generate_audio_digest,
    inputs=gr.Textbox(label="Enter your topic of interest", placeholder="e.g. motivation, productivity"),
    outputs=[gr.Text(label="Summary"), gr.Audio(label="Audio Digest", type="filepath")],
    title="🎧 KukuBuddy: Daily Audio Digest"
)

story_ui = gr.Interface(
    fn=generate_story,
    inputs=[
        gr.Textbox(label="Genre", placeholder="e.g. sci-fi, myth, thriller"),
        gr.Textbox(label="Your last choice", placeholder="e.g. Enter the castle")
    ],
    outputs=[gr.Text(label="Next Scene"), gr.Audio(label="Narration", type="filepath")],
    title="📖 KukuBuddy: Interactive Story"
)

app = gr.TabbedInterface([digest_ui, story_ui], tab_names=["📌 Daily Digest", "🧠 Interactive Story"])

if __name__ == "__main__":
    app.launch()