import os
import yt_dlp
import gradio as gr
from faster_whisper import WhisperModel
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
import openai
import torch

# Optional: Set your OpenAI API key (use env var for security)
openai.api_key = "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"  # Replace with your real key or set as os.environ["OPENAI_API_KEY"]

def download_and_extract_audio(youtube_url):
    output_path = "downloads"
    os.makedirs(output_path, exist_ok=True)

    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': os.path.join(output_path, '%(id)s.%(ext)s'),
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
    }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(youtube_url, download=True)
        video_id = info_dict.get("id", None)
        filename = os.path.join(output_path, f"{video_id}.mp3")
        return filename

def transcribe_audio(audio_path):
    model = WhisperModel("base", compute_type="int8", device="cuda" if torch.cuda.is_available() else "cpu")
    segments, _ = model.transcribe(audio_path)
    transcript = " ".join([seg.text for seg in segments])
    return transcript

# Preload FLAN-T5 model offline
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
local_gen = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

def generate_response(transcript, user_prompt, use_online=False):
    prompt = f"""You are a helpful AI assistant. Based on the transcript of a video, please {user_prompt.strip().lower()}.

Transcript:
{transcript[:3000]}"""

    if use_online:
        try:
            response = openai.ChatCompletion.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=1000
            )
            return response.choices[0].message["content"]
        except Exception as e:
            return f"⚠️ Online API failed: {str(e)}"
    else:
        result = local_gen(prompt, max_length=1024, do_sample=False)
        return result[0]['generated_text']

def enhanced_ai_study_pipeline(video_source, youtube_url, upload_file, user_prompt, use_online_api):
    try:
        if video_source == "YouTube URL":
            audio_path = download_and_extract_audio(youtube_url)
        elif video_source == "Upload File" and upload_file is not None:
            audio_path = upload_file.name
        else:
            return "No valid input provided.", ""

        transcript = transcribe_audio(audio_path)
        ai_response = generate_response(transcript, user_prompt, use_online=use_online_api)
        return transcript, ai_response
    except Exception as e:
        return "Error occurred", str(e)

video_input = gr.Radio(["YouTube URL", "Upload File"], label="Video Source")
youtube_url = gr.Textbox(label="Enter YouTube URL")
upload_file = gr.File(label="Upload a Video File", file_types=[".mp4", ".mp3", ".wav"])
user_prompt = gr.Textbox(label="What do you want from the transcript?", placeholder="e.g., Prepare a diet plan based on this video")
use_online_api = gr.Checkbox(label="Use Online API (GPT-4)", value=False)

gr.Interface(
    fn=enhanced_ai_study_pipeline,
    inputs=[video_input, youtube_url, upload_file, user_prompt, use_online_api],
    outputs=[
        gr.Textbox(label="Transcript"),
        gr.Textbox(label="AI Response")
    ],
    title="📚 AI Transcription Assistant (Offline + Online GPT)",
    description="Upload or paste a YouTube video. Enter your goal and get a smart AI answer. Works offline with FLAN-T5 or online with GPT-4."
).launch(share=True)