import streamlit as st
from crewai import Agent, Task, Crew
from transformers import pipeline
import torch
import whisper
import os

st.title("AI Blog Writer from YouTube")

# Whisper model setup (Base Model for better speed)
whisper_model = whisper.load_model("base")

def get_video_transcript(video_url):
    """Fetch video transcript using CrewAI's YouTubeChannelSearch tool."""
    from crewai.tools.youtubechannelsearch import YouTubeChannelSearch

    # Initialize the YouTube Channel Search tool
    search_tool = YouTubeChannelSearch()

    # Search for video and fetch transcript
    video_data = search_tool.search(video_url)
    
    if video_data and 'transcript' in video_data:
        return video_data['transcript']
    else:
        raise Exception("Transcript not found for this video.")

def transcribe_audio(audio_path):
    """Transcribes the audio using Whisper AI."""
    transcription = whisper_model.transcribe(audio_path)
    return transcription['text']

# CrewAI Setup
def create_crew_agents(transcript_text):
    research_agent = Agent(name="Research Agent", goal="Extract video transcript.")
    blog_writer_agent = Agent(name="Blog Writer", goal="Generate blog content.")
    editor_agent = Agent(name="Editor", goal="Refine blog for clarity.")

    research_task = Task(description="Fetch transcript.", agent=research_agent)
    blog_task = Task(description="Write blog using transcript.", agent=blog_writer_agent, context=transcript_text)
    editor_task = Task(description="Edit the blog.", agent=editor_agent)

    crew = Crew(agents=[research_agent, blog_writer_agent, editor_agent], tasks=[research_task, blog_task, editor_task])
    return crew

# User Input for Video URL
video_url = st.text_input("Enter YouTube Video URL:")

if st.button("Generate Blog"):
    if video_url:
        try:
            st.info("Fetching transcript...")
            transcript_text = get_video_transcript(video_url)  # Fetch transcript using CrewAI
            st.text_area("Transcript Extracted", transcript_text, height=200)

            # Generate Blog with CrewAI
            crew = create_crew_agents(transcript_text)
            results = crew.kickoff()
            st.subheader("Generated Blog")
            st.write(results)

        except Exception as e:
            st.error(f"Error: {e}")
    else:
        st.warning("Please enter a valid YouTube URL.")