import os
from flask_sqlalchemy import SQLAlchemy
from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from gradio import Interface, Textbox, Dropdown, Markdown
from textwrap import dedent
import nest_asyncio

# Apply nested asyncio
nest_asyncio.apply()

# Set the API key for Groq
os.environ["GROQ_API_KEY"] = "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"

# Helper function to create agents
def create_agent(system_prompt: str, model_name: str) -> ChatGroq:
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", "{input}")
    ])
    llm = ChatGroq(model=model_name)
    return prompt | llm | StrOutputParser()

# Combined function that handles all tasks
def handle_task(task_type, query):
    if task_type == "Trends":
        system_prompt = """
        You are an expert at identifying trending topics on TikTok, YouTube, Instagram, and Snapchat.
        Analyze the query and provide the most relevant trending topics.
        """
        agent = create_agent(system_prompt, model_name="llama3-8b-8192")
        return agent.invoke({"input": query})

    elif task_type == "Script Generation":
        system_prompt = """
        You are a creative expert who writes scripts with the perfect formula for TikTok virality.
        Generate a detailed, engaging script based on the query.
        """
        agent = create_agent(system_prompt, model_name="llama3-8b-8192")
        return agent.invoke({"input": query})

    elif task_type == "Hashtag Generation":
        system_prompt = """
        You are skilled at generating hashtags and tags for social media platforms.
        Based on the query, provide the following:
        - 30 unique TikTok viral tags
        - 50 most popular hashtags
        - 50 FYP-related tags
        - 25 YouTube viral keyword tags
        - A clickbait title with emojis
        """
        agent = create_agent(system_prompt, model_name="llama3-8b-8192")
        return agent.invoke({"input": query})

# Create a dropdown for task selection
task_options = ["Trends", "Script Generation", "Hashtag Generation"]

# Gradio interface
interface = Interface(
    fn=handle_task,
    inputs=[
        Dropdown(label="Select Task", choices=task_options),
        Textbox(label="Enter Query", placeholder="Enter your query here...")
    ],
    outputs=Markdown(label="Output"),
    title="Viral Content Creation tool, by iLL-Ai Aaron Allton",
    description="Choose a task (Trends, Script Generation, or Hashtag Generation) and enter your query to get tailored responses."
)

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