import streamlit as st
import openai
import pandas as pd
from youtube_transcript_api import YouTubeTranscriptApi
from io import StringIO
from docx import Document

# Function to get YouTube transcript
def get_transcript(video_id):
    transcript = YouTubeTranscriptApi.get_transcript(video_id)
    full_text = " ".join([entry['text'] for entry in transcript])
    return full_text

# Function to generate a script using OpenAI
def generate_shorts_script(transcript, ai_system_content, user_content):
    openai.api_key = st.secrets["OPENAI_API_KEY"]  # API Key from Streamlit secrets
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": ai_system_content},
            {"role": "user", "content": f"{user_content}: {transcript}"}
        ]
    )
    return response.choices[0].message["content"]

# Streamlit UI
st.title("YouTube Shorts Script Generator")

# Input fields
openai_api_key = st.text_input("OpenAI API Key", type="password")
video_id = st.text_input("YouTube Video ID")
ai_system_content = st.text_area("AI System Content")
user_content = st.text_area("User Content")

if st.button("Generate Script"):
    if openai_api_key and video_id and ai_system_content and user_content:
        try:
            # Set OpenAI API key
            openai.api_key = openai_api_key

            # Get transcript
            transcript = get_transcript(video_id)

            # Generate script
            shorts_script = generate_shorts_script(transcript, ai_system_content, user_content)

            # Display script
            st.write("Generated Script:")
            st.write(shorts_script)

            # Options to download
            if st.button("Download as CSV"):
                df = pd.DataFrame({"Script": [shorts_script]})
                csv = df.to_csv(index=False)
                st.download_button("Download CSV", csv, "script.csv", "text/csv")

            if st.button("Download as Text"):
                st.download_button("Download Text", shorts_script, "script.txt", "text/plain")

            if st.button("Download as Word"):
                doc = Document()
                doc.add_paragraph(shorts_script)
                buffer = StringIO()
                doc.save(buffer)
                docx = buffer.getvalue()
                st.download_button("Download Word", docx, "script.docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document")
        except Exception as e:
            st.error(f"An error occurred: {e}")
    else:
        st.error("Please fill all the input fields.")