from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_community.utilities import GoogleSerperAPIWrapper
import os

def generate_script(prompt, video_length, creativity, openai_api_key, serper_api_key):
    """
    Generate a YouTube video title and script based on the provided prompt, duration, and creativity level.
    """

    # Ensure the OpenAI API key is set
    if not openai_api_key:
        raise ValueError("OpenAI API Key is missing.")
    os.environ["OPENAI_API_KEY"] = openai_api_key

    # Ensure the Serper API key is set
    if not serper_api_key:
        raise ValueError("Serper API Key is missing.")
    os.environ["SERPER_API_KEY"] = serper_api_key

    # Initialize the LLM
    llm = ChatOpenAI(
        model='gpt-4o-mini',
        temperature=creativity,
        max_tokens=None,
        timeout=None,
        max_retries=2
    )

    # Define the title generation prompt
    title_template = PromptTemplate(
        input_variables=["topic"],
        template="Generate a YouTube video title for the topic '{topic}'."
    )
    title_chain = LLMChain(llm=llm, prompt=title_template, verbose=False)

    # Define the script generation prompt
    script_template = PromptTemplate(
        input_variables=["title", "duration", "search_data"],
        template=(
            "Create a YouTube script with the title '{title}' that lasts approximately {duration} minutes. "
            "Incorporate the following research information: {search_data}."
        )
    )
    script_chain = LLMChain(llm=llm, prompt=script_template, verbose=False)

    # Fetch relevant search data using Google Serper
    search_tool = GoogleSerperAPIWrapper()
    search_data = search_tool.results(prompt)

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

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

    return title, script, search_data