import streamlit as st
import yt_dlp
import os

# Streamlit 제목
st.title('유튜브 영상 다운로드')
st.write("<br>", unsafe_allow_html=True)  # 1줄 띄어쓰기
st.write("사용방법")
st.write("<br>", unsafe_allow_html=True)  # 1줄 띄어쓰기
st.write("1. 다운받으려는 유튜브 영상의 링크를 입력해주세요")
st.write("2. 링크를 입력하고 확인버튼을 누르면 해당 영상의 썸네일이 등장합니다")
st.write("3. 썸네일 밑의 선택란을 통해 해상도와 영상 용량을 확인한 후 원하는 해상도를 선택하고 추출 버튼을 눌러주세요")
st.write("4. 잠시 후 다운받으려는 영상이 밑에 등장하고 재생하거나 다운로드 받을 수 있습니다 (영상의 용량이 클수록 추출하는데 시간이 더 걸립니다. 20분 영상 추출하는데 1분내로 소요)")

# 유저로부터 유튜브 링크 입력받기
col1, col2 = st.columns([4, 1])
with col1:
    youtube_url = st.text_input('Enter YouTube video URL')

def fetch_video_info(youtube_url):
    ydl_opts = {
        'nocheckcertificate': True,
        'cookiefile': 'cookies.txt'  # 쿠키 파일 경로
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(youtube_url, download=False)
        return info_dict

def get_available_formats(info_dict):
    formats = {}
    for fmt in info_dict['formats']:
        if fmt.get('height') and fmt.get('ext') == 'mp4':
            filesize = fmt.get('filesize', 0)
            if filesize:
                filesize_mb = round(filesize / (1024 * 1024), 2)
                resolution = f"{fmt['height']}p"
                formats[resolution] = {
                    'id': fmt['format_id'],
                    'info': f"{resolution} - {filesize_mb} MB"
                }
    return formats

def reset_session_state():
    st.session_state.pop('info_dict', None)
    st.session_state.pop('thumbnail_url', None)
    st.session_state.pop('formats', None)
    st.session_state.pop('youtube_url', None)
    st.session_state.pop('confirmed', None)
    st.session_state.pop('selected_format_id', None)

def delete_downloaded_video():
    if os.path.exists('downloaded_video.mp4'):
        os.remove('downloaded_video.mp4')

# "확인" 버튼을 누르면 처리할 로직
if st.button('확인'):
    if youtube_url:
        try:
            # 기존 다운로드된 영상 삭제
            delete_downloaded_video()

            # 세션 상태 초기화 및 URL 변경 확인
            if 'youtube_url' in st.session_state and st.session_state.youtube_url != youtube_url:
                reset_session_state()

            st.session_state.youtube_url = youtube_url
            info_dict = fetch_video_info(youtube_url)
            st.session_state.info_dict = info_dict
            st.session_state.thumbnail_url = info_dict.get('thumbnail')
            st.session_state.formats = get_available_formats(info_dict)
            st.session_state.confirmed = True
        except Exception as e:
            st.error(f"An error occurred: {e}")
    else:
        st.warning('Please enter a valid YouTube video URL')

if 'confirmed' in st.session_state and st.session_state.confirmed:
    st.image(st.session_state.thumbnail_url, caption='Video Thumbnail', use_column_width=True)

    # 해상도 선택
    format_options = [value['info'] for value in st.session_state.formats.values()]
    format_selection = st.selectbox('Select resolution', format_options)

    # 선택된 해상도의 format_id 저장
    st.session_state.selected_format_id = next(
        (value['id'] for key, value in st.session_state.formats.items() if value['info'] == format_selection), None
    )

    def download_video(youtube_url, format_id):
        ydl_opts = {
            'format': f"{format_id}+bestaudio/best",
            'merge_output_format': 'mp4',
            'outtmpl': 'downloaded_video.%(ext)s',
            'nocache': True,  # 캐시 비활성화
            'cookiefile': 'cookies.txt'  # 쿠키 파일 경로
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([youtube_url])

    # "추출" 버튼 추가
    if st.button('추출'):
        format_id = st.session_state.selected_format_id
        if format_id is None:
            st.error('Selected format not found.')
        else:
            try:
                # 세션 상태의 youtube_url 사용
                download_video(st.session_state.youtube_url, format_id)
                st.success('Video downloaded successfully!')
                st.video('downloaded_video.mp4')

                # 다운로드된 파일 제공 (optional)
                with open('downloaded_video.mp4', 'rb') as file:
                    st.download_button(
                        label="Download Video File",
                        data=file,
                        file_name='downloaded_video.mp4',
                        mime='video/mp4'
                    )

                # 추출된 영상 정보 출력
                st.write(f"**제목:** {st.session_state.info_dict['title']}")
                st.write(f"**채널명:** {st.session_state.info_dict['uploader']}")
                st.write(f"**업로드 날짜:** {st.session_state.info_dict['upload_date']}")

                # 다운로드 후 세션 상태 리셋
                reset_session_state()
            except Exception as e:
                st.error(f"An error occurred during download: {e}")