# app.py
import gradio as gr
import sys
from tenacity import retry, stop_after_attempt, wait_exponential
import os

# Workaround for sqlite3 version issue
__import__('pysqlite3')
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

from crewai import Agent, Task, Crew
from langchain_community.llms import HuggingFaceHub

def generate_jingle(theme):
    # Hugging Face API Key handling
    hf_api_key = os.getenv("HUGGINGFACEHUB_API_TOKEN")
    if not hf_api_key:
        return "Error: Hugging Face API key not found. Please set 'HUGGINGFACEHUB_API_TOKEN' in Space secrets."
    if not hf_api_key.startswith("hf_"):
        return "Error: Invalid Hugging Face API token. It should start with 'hf_'. Please check and update the token in Space secrets."

    # LLM setup with HuggingFaceHub
    try:
        llm = HuggingFaceHub(
            repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1",
            huggingfacehub_api_token=hf_api_key,
            model_kwargs={
                "max_new_tokens": 500,
                "temperature": 0.7,
                "top_p": 0.9
            }
        )
        # Test the endpoint
        try:
            test_response = llm.invoke("Test connection to Hugging Face API")
            return f"Hugging Face API connection successful! Response: {test_response[:100]}...\n\nStarting jingle generation..."
        except Exception as test_e:
            return f"Warning: Test API call failed: {str(test_e)}. Proceeding anyway, but jingle generation may fail."
    except Exception as e:
        return f"Error: Failed to initialize SLM: {str(e)}"

    try:
        # Define Agents
        researcher = Agent(
            role="Researcher",
            goal="Research key facts, trends, and appealing elements about the theme to inspire the jingle.",
            backstory="You are an expert researcher with access to vast knowledge on various topics. Focus on fun, engaging, and relevant info for radio ads.",
            verbose=True,
            llm=llm
        )

        creator = Agent(
            role="Jingle Creator",
            goal="Create a short, catchy radio jingle based on research, including lyrics and simple structure suggestions.",
            backstory="You are a creative genius specializing in audio ads. Make it rhythmic, memorable, and under 30 seconds worth of content.",
            verbose=True,
            llm=llm
        )

        copywriter = Agent(
            role="Copywriter",
            goal="Refine and polish the jingle content for clarity, impact, and radio-friendliness.",
            backstory="You are a professional copywriter with experience in advertising. Ensure it's persuasive, error-free, and optimized for spoken delivery.",
            verbose=True,
            llm=llm
        )

        # Define Tasks
        research_task = Task(
            description=f"Research '{theme}'. List 5 key points for an engaging jingle.",
            expected_output="A bullet-point list of research findings.",
            agent=researcher
        )

        create_task = Task(
            description=f"Using the research, create a short radio jingle for '{theme}'. Include lyrics and notes on rhythm/timing.",
            expected_output="The jingle lyrics with structure (e.g., verse, chorus). Keep it concise.",
            agent=creator,
            context=[research_task]
        )

        copywrite_task = Task(
            description=f"Refine the created jingle for '{theme}'. Improve flow, add punch, ensure it's radio-ready.",
            expected_output="The final polished jingle script.",
            agent=copywriter,
            context=[create_task]
        )

        # Assemble Crew
        crew = Crew(
            agents=[researcher, creator, copywriter],
            tasks=[research_task, create_task, copywrite_task],
            verbose=True
        )

        # Run the crew with retry logic
        @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
        def run_crew():
            try:
                return crew.kickoff()
            except Exception as inner_e:
                raise Exception(f"Crew execution failed: {str(inner_e)}")

        result = run_crew()

        # Format output
        output = f"**Final Polished Jingle**\n{result}\n\n"
        output += "**Research Findings**\n" + (crew.tasks[0].output.raw_output or "Not available")
        output += "\n\n**Initial Creation**\n" + (crew.tasks[1].output.raw_output or "Not available")
        return output

    except Exception as e:
        return f"Error: An error occurred while generating the jingle: {str(e)}"

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("""
    # Radio Jingle Generator using CrewAI and SLM
    This app uses CrewAI with a Small Language Model (Mistral 7B) to generate short radio jingles. It includes:
    - **Researcher AI**: Researches the theme or product.
    - **Creator AI**: Creates the initial jingle lyrics and structure.
    - **Copywriter AI**: Polishes and refines the content for radio.

    Enter a theme (e.g., "coffee shop promotion") and generate!
    """)
    theme_input = gr.Textbox(label="Theme or Product for the Jingle", placeholder="Example: Summer Beach Party")
    generate_button = gr.Button("Generate Jingle")
    output = gr.Textbox(label="Generated Jingle and Details", interactive=False)
    generate_button.click(fn=generate_jingle, inputs=theme_input, outputs=output)

demo.launch()