import os
from gemini_flash import Gemini
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.tools import DuckDuckGoSearchRun

def generate_script(prompt, video_length, creativity, api_key):
    # Setup Gemini API
    gemini = Gemini(api_key=api_key)
    
    # Prompt template for generating video title
    title_template = PromptTemplate(
        input_variables=["topic"], 
        template="Generate a YouTube video title based on the following topic: {topic}"
    )
    title_chain = LLMChain(llm=gemini, prompt=title_template, verbose=True)

    # Prompt template for generating video script
    script_template = PromptTemplate(
        input_variables=["title", "duration", "search_data"],
        template="Create a YouTube script with the title '{title}' that lasts {duration} minutes. Use the following search data: {search_data}"
    )
    script_chain = LLMChain(llm=gemini, prompt=script_template, verbose=True)

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

    # Generate the title
    title = title_chain.run({"topic": prompt})

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

    return title, script, search_data