import streamlit as st
import yt_dlp
import os
from pathlib import Path
import random
import time

# 페이지 설정
st.set_page_config(page_title="Simple YouTube Downloader", page_icon="📺")
st.title("Simple YouTube Downloader 📺")

# 스타일 설정
st.markdown("""
    <style>
        .stButton > button {
            width: 100%;
            height: 60px;
        }
    </style>
""", unsafe_allow_html=True)

# 다운로드 폴더 생성
output_dir = Path("downloads")
output_dir.mkdir(exist_ok=True)

def download_video(url, progress_bar, status_text):
    try:
        user_agents = [
            'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
            'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.2.1 Safari/605.1.15',
            'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:122.0) Gecko/20100101 Firefox/122.0'
        ]
        
        ydl_opts = {
            'format': 'best',
            'outtmpl': str(output_dir / '%(title)s.%(ext)s'),
            'quiet': True,
            'no_warnings': True,
            'cookiefile': 'youtube.com_cookies.txt',
            'user_agent': random.choice(user_agents),
            'referer': 'https://www.youtube.com/',
            'http_chunk_size': random.randint(10000000, 15000000),
            'retries': 1,
            'sleep_interval': 3,
            'max_sleep_interval': 8,
            'headers': {
                'Accept-Language': 'en-US,en;q=0.9',
                'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
            },
        }

        def progress_hook(d):
            if d['status'] == 'downloading':
                try:
                    progress = d['downloaded_bytes'] / d['total_bytes']
                    progress_bar.progress(progress)
                    status_text.text(f"다운로드 중: {progress:.1%}")
                except:
                    status_text.text("다운로드 중...")
            elif d['status'] == 'finished':
                progress_bar.progress(1.0)
                status_text.text("처리 완료!")

        ydl_opts['progress_hooks'] = [progress_hook]
        
        for attempt in range(3):
            try:
                with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                    info = ydl.extract_info(url, download=True)
                    filename = ydl.prepare_filename(info)
                    return filename
            except yt_dlp.utils.ExtractorError as e:
                if "Sign in to confirm you're not a bot" in str(e):
                    status_text.text(f"재시도 중... ({attempt + 1}/3)")
                    time.sleep(random.uniform(3, 5))
                    continue
                raise e
            except Exception as e:
                if attempt < 2:
                    status_text.text(f"재시도 중... ({attempt + 1}/3)")
                    time.sleep(2)
                    continue
                raise e

    except Exception as e:
        st.error(f"오류 발생: {str(e)}")
        return None

# 쿠키 파일 체크
if not os.path.exists('youtube.com_cookies.txt'):
    st.warning("⚠️ 'youtube.com_cookies.txt' 파일이 필요합니다. YouTube에 로그인된 쿠키 파일을 업로드해주세요.")
    uploaded_file = st.file_uploader("쿠키 파일 업로드", type=['txt'])
    if uploaded_file is not None:
        with open('youtube.com_cookies.txt', 'wb') as f:
            f.write(uploaded_file.getvalue())
        st.success("✅ 쿠키 파일이 업로드되었습니다!")
        st.rerun()

# 메인 UI
if 'video_url' not in st.session_state:
    st.session_state.video_url = ""

video_url = st.text_input("YouTube URL 입력:", 
                         value=st.session_state.video_url,
                         placeholder="https://www.youtube.com/watch?v=...")

# 버튼을 같은 줄에 배치
col1, col2 = st.columns(2)

with col1:
    download_button = st.button("다운로드", type="primary", key="download_btn", use_container_width=True)

with col2:
    reset_button = st.button("초기화", key="reset_btn", use_container_width=True)

# 다운로드 로직
if download_button:
    if video_url:
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        downloaded_file = download_video(video_url, progress_bar, status_text)
        
        if downloaded_file and os.path.exists(downloaded_file):
            with open(downloaded_file, 'rb') as file:
                st.download_button(
                    label="⬇️ 파일 다운로드",
                    data=file,
                    file_name=os.path.basename(downloaded_file),
                    mime="video/mp4",
                    use_container_width=True
                )
    else:
        st.warning("⚠️ YouTube URL을 입력해주세요.")

# 초기화 로직
if reset_button:
    st.session_state.video_url = ""
    st.rerun()