from pytube import YouTube
from google.cloud import speech_v1p1beta1 as speech
from pydub import AudioSegment
import io
import openai
import os
import gradio as gr

# 유튜브 비디오 ID 추출 함수
def get_yt_video_id(url):
    from urllib.parse import urlparse, parse_qs

    if url.startswith(('youtu', 'www')):
        url = 'http://' + url
        
    query = urlparse(url)
    
    if 'youtube' in query.hostname:
        if query.path == '/watch':
            return parse_qs(query.query)['v'][0]
        elif query.path.startswith(('/embed/', '/v/')):
            return query.path.split('/')[2]
    elif 'youtu.be' in query.hostname:
        return query.path[1:]
    else:
        raise ValueError("유효한 유튜브 링크가 아닙니다.")

# 오디오 추출 및 변환 함수 (WAV 형식으로 변환)
def download_and_convert_audio(youtube_url):
    yt = YouTube(youtube_url)
    stream = yt.streams.filter(only_audio=True).first()
    audio_path = stream.download(filename="audio.mp4")
    
    # 오디오 파일을 WAV로 변환 (16000Hz 샘플 레이트)
    audio = AudioSegment.from_file(audio_path)
    wav_audio_path = "converted_audio.wav"
    audio.set_frame_rate(16000).set_channels(1).export(wav_audio_path, format="wav")
    
    return wav_audio_path

# 오디오를 청크로 나누는 함수
def split_audio(audio_path, chunk_length_ms=60000):
    audio = AudioSegment.from_wav(audio_path)
    chunks = [audio[i:i + chunk_length_ms] for i in range(0, len(audio), chunk_length_ms)]
    return chunks

# Google Speech-to-Text API를 사용하여 오디오를 텍스트로 변환
def speech_to_text(audio_path):
    client = speech.SpeechClient()
    chunks = split_audio(audio_path)  # 오디오를 청크로 나눔

    transcript = ""
    for chunk in chunks:
        with io.BytesIO() as audio_file:
            chunk.export(audio_file, format="wav")
            audio_file.seek(0)
            content = audio_file.read()

            audio = speech.RecognitionAudio(content=content)
            config = speech.RecognitionConfig(
                encoding=speech.RecognitionConfig.AudioEncoding.LINEAR16,
                sample_rate_hertz=16000,  # 16000Hz 샘플 레이트
                language_code="ko-KR"  # 한국어 인식
            )

            response = client.recognize(config=config, audio=audio)
            for result in response.results:
                transcript += result.alternatives[0].transcript + " "
    
    return transcript.strip()

# 텍스트를 요약하는 함수 (OpenAI API 사용)
def textToSummary(text):
    openai.api_key = os.getenv("OPENAI_API_KEY")  # 환경 변수에서 OpenAI API 키 가져오기
    response = openai.Completion.create(
        model="text-davinci-003",
        prompt="Summarize this in 200 words or less:\n\n" + text,
        temperature=0.7,
        max_tokens=400,
        top_p=1.0,
        frequency_penalty=0.0,
        presence_penalty=1
    )
    return response["choices"][0]["text"].replace("\n", " ").strip()

# 전체 요약 프로세스를 처리하는 함수
def summarize(url):
    try:
        # 유튜브 오디오 다운로드 및 변환
        audio_path = download_and_convert_audio(url)
        # 음성을 텍스트로 변환
        transcript = speech_to_text(audio_path)
        # 텍스트 요약
        summary = textToSummary(transcript)
        return summary
    except Exception as e:
        return f"요약에 실패했습니다: {str(e)}"

# Gradio 인터페이스 설정
description = "유튜브 동영상의 자막이 없더라도 음성 인식 기능을 사용해 요약합니다."

gr.Interface(fn=summarize,
             inputs="text",
             outputs="textbox", 
             description=description
            ).launch()