import gradio as gr
import torch
import os
from langchain.document_loaders import YoutubeLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.chains import RetrievalQA
from langchain import HuggingFaceHub
from urllib.parse import urlparse, parse_qs

def extract_video_id(youtube_url):
    try:
        parsed_url = urlparse(youtube_url)
        query_params = parse_qs(parsed_url.query)
        video_id = query_params.get('v', [None])[0]
        return video_id
    except Exception as e:
        return f"Error extracting video ID: {e}"

def process_video(youtube_url, question):
    video_id = extract_video_id(youtube_url)
    if not video_id:
        return 'Invalid YouTube URL'

    try:
        # Initialize the YouTube Loader
        loader = YoutubeLoader(video_id)
        documents = loader.load()

        # Process the documents
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
        documents = text_splitter.split_documents(documents)

        # Initialize Vector Store
        model_name = "BAAI/bge-base-en"
        encode_kwargs = {'normalize_embeddings': True}
        vectordb = Chroma.from_documents(
            documents,
            embedding=HuggingFaceBgeEmbeddings(model_name=model_name, 
                                               model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'}, 
                                               encode_kwargs=encode_kwargs)
        )

        # Setup the QA Chain
        HUGGINGFACE_API_TOKEN = os.environ['HUGGINGFACE_API_TOKEN']
        repo_id = "tiiuae/falcon-7b-instruct"
        qa_chain = RetrievalQA.from_chain_type(
            llm=HuggingFaceHub(huggingfacehub_api_token=HUGGINGFACE_API_TOKEN,
                               repo_id=repo_id,
                               model_kwargs={"temperature":0.1, "max_new_tokens":1000}),
            retriever=vectordb.as_retriever(),
            return_source_documents=False,
            verbose=False
        )

        # Process the question
        llm_response = qa_chain(question)
        return llm_response['result']
    except Exception as e:
        return f"Error processing video: {e}"

iface = gr.Interface(
    fn=process_video,
    inputs=["text", "text"],
    outputs="text",
    title="YouTube Video AI Assistant",
    description="Enter a YouTube URL and a question to get AI-generated answers based on the video."
)

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