import gradio as gr
from pydantic import BaseModel
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_groq import ChatGroq
from langchain_tavily import TavilySearch
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver
import uuid
import os

# Set up LLM
llm = ChatGroq(
    temperature=0,
    groq_api_key=os.getenv("GROQ_API_KEY"),
    model_name="llama-3.3-70b-versatile"
)

# Define tools
search_tool = TavilySearch(
    max_results=5,
    tavily_api_key=os.getenv("TAVILY_API_KEY")
)

tools = [search_tool]

# Define system prompt
system_prompt = """
You are a professional content strategist and scriptwriter for Instagram Reels.
Your job is to help content creators generate high-performing short-form video scripts (Reels) based on a theme or rough idea they provide.

--------------------------------------

INPUT FROM USER:
Theme or Idea: {text}

--------------------------------------

Instructions:
1. Understand the topic and tone from the input. If the input is a vague idea, clarify the theme internally before writing.
2. The target audience is Instagram users (18–35), who enjoy short, engaging, and value-packed content.
3. Keep the script short, clear, and suitable for a 30–60 second video.
4. The tone should be friendly, motivational, and relatable (unless specified otherwise).
5. Use emojis in the caption and keep hashtags relevant to the theme.

--------------------------------------

OUTPUT FORMAT:
- Hook: A scroll-stopping first line (1 sentence only).
- 3-Point Structure or Mini Story:
1. [Point or part 1]
2. [Point or part 2]
3. [Point or part 3]
- Caption: A 1–2 sentence engaging caption.
- Hashtags: Include 3–5 relevant hashtags.
- Call-to-Action: End with a line like “Follow for more tips”, “Save this”, etc.

Only respond with the script in this format. Do not explain your reasoning.

--------------------------------------
"""

# Create checkpoint
checkpointer = InMemorySaver()

# Create LangGraph agent
agent = create_react_agent(
    model=llm,
    tools=tools,
    prompt=system_prompt,
    checkpointer=checkpointer
)

# Script generation function for Gradio
def generate_script(text: str):
    try:
        thread_id = str(uuid.uuid4())  # Generate a unique thread ID
        config = {"configurable": {"thread_id": thread_id}}
        response = agent.invoke(
            {"messages": [{"role": "user", "content": text}]},
            config
        )
        return response['messages'][-1].content
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
interface = gr.Interface(
    fn=generate_script,
    inputs=gr.Textbox(lines=3, placeholder="Enter your Reel idea or theme here..."),
    outputs=gr.Textbox(label="Generated Reel Script"),
    title="Instagram Reel Script Generator 🎬",
    description="Enter your idea or theme and get a full short-form video script tailored for Instagram Reels."
)

if __name__ == "__main__":
    interface.launch(share=True)