import os
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.tools import DuckDuckGoSearchRun


def generate_script(prompt, video_length, creativity, api_key):
    """
    Generate a YouTube video script based on the given topic, length, and creativity.
    
    Args:
        prompt (str): The topic of the video.
        video_length (float): Desired length of the video in minutes.
        creativity (float): Creativity level for the script (0.0 to 1.0).
        api_key (str): API key for Google Generative AI.
    
    Returns:
        tuple: A tuple containing the title, script, and search data.
    """
    # Setup Google Generative AI
    chat = ChatGoogleGenerativeAI(
        model="gemini-1.5-flash",
        google_api_key=api_key
    )

    # Title prompt template
    title_template = PromptTemplate(
        input_variables=["topic"],
        template="Generate a catchy YouTube video title for the topic: {topic}"
    )
    title_chain = LLMChain(llm=chat, prompt=title_template, verbose=True)
    
    # Generate the title
    title = title_chain.run({"topic": prompt})

    # Get relevant search data using DuckDuckGo
    search_tool = DuckDuckGoSearchRun()
    search_data = search_tool.run(prompt)

    # Script prompt template
    script_template = PromptTemplate(
        input_variables=["title", "duration", "search_data"],
        template=(
            "Write a detailed YouTube script for a video titled '{title}'. "
            "The video should last approximately {duration} minutes and incorporate "
            "the following search data: {search_data}"
        )
    )
    script_chain = LLMChain(llm=chat, prompt=script_template, verbose=True)

    # Generate the script
    script = script_chain.run({
        "title": title,
        "duration": video_length,
        "search_data": search_data
    })

    return title, script, search_data