import os
import requests
import json
import gradio as gr
import re

# Hugging Face 환경 변수로부터 RapidAPI 키와 호스트 가져오기
AA_KEY = os.getenv("AA_KEY")
AA_HOST = "youtube-transcriptor.p.rapidapi.com"

# 자막 언어 우선순위 리스트 
LANGUAGE_PRIORITY = ['ko', 'en', 'ja', 'zh']

# 유튜브 URL에서 비디오 ID를 추출하는 함수
def get_video_id(youtube_url):
   # 유튜브 URL 또는 youtu.be 단축 URL에서 video_id 추출
   video_id_match = re.search(r"(?<=v=)[^#&?]*", youtube_url) or re.search(r"(?<=youtu.be/)[^#&?]*", youtube_url)
   
   # YouTube Shorts URL 처리
   if not video_id_match:
       video_id_match = re.search(r"(?<=shorts/)[^#&?]*", youtube_url)
   
   return video_id_match.group(0) if video_id_match else None

# 유튜브 자막을 요청하는 함수 (언어 우선순위를 적용하여 시도)
def get_youtube_transcript(youtube_url):
    try:
        # 비디오 ID 추출
        video_id = get_video_id(youtube_url)
        if video_id is None:
            return {"error": "잘못된 유튜브 URL입니다."}
        
        url = "https://youtube-transcriptor.p.rapidapi.com/transcript"
        headers = {
            "x-rapidapi-key": AA_KEY,
            "x-rapidapi-host": AA_HOST
        }
        
        # 1. 우선순위 언어로 시도
        for lang in LANGUAGE_PRIORITY:
            querystring = {"video_id": video_id, "lang": lang}
            response = requests.get(url, headers=headers, params=querystring)
            
            if response.status_code == 200:
                data = response.json()
                if data and not isinstance(data, str) and "error" not in data:
                    return {"language": lang, "data": data}
                if isinstance(data, dict) and "availableLangs" in data:
                    available_langs = data["availableLangs"]
                    # 2. 가능한 언어가 있다면 해당 언어로 시도
                    for available_lang in available_langs:
                        querystring = {"video_id": video_id, "lang": available_lang}
                        response = requests.get(url, headers=headers, params=querystring)
                        if response.status_code == 200:
                            data = response.json()
                            if data and not isinstance(data, str) and "error" not in data:
                                return {"language": available_lang, "data": data}
        
        # 3. 모든 시도 실패시
        return {"error": "자막을 찾을 수 없습니다."}
        
    except Exception as e:
        return {"error": "자막을 불러오는데 실패했습니다."}

# Gradio 인터페이스 함수
def youtube_transcript_interface(youtube_url):
    transcript_data = get_youtube_transcript(youtube_url)
    return json.dumps(transcript_data, ensure_ascii=False, indent=2)

# Gradio 인터페이스 생성
interface = gr.Interface(
   fn=youtube_transcript_interface,
   inputs="text",
   outputs="text",
   title="YouTube 자막 추출기",
   description="유튜브 URL을 입력하세요."
)

# Gradio 인터페이스 실행
interface.launch()