import streamlit as st
import requests
from tqdm import tqdm
import time

def download(url, filename, path, resolution):
    '''
    下载函数
    url: str 需要下载的文件链接
    filename: 储存的文件名
    path: 下载的文件储存地址
    resolution: 选择的分辨率（240, 480, 720, 1080）
    
    return: None
    '''
    headers = {
        'Referer': 'https://music.163.com/',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'
    }
    res = requests.get(url, headers=headers, stream=True)
    size = int(res.headers.get('Content-Length'))
    
    # 根据分辨率选择视频文件
    resol_map = {
        '240': '240',
        '480': '480',
        '720': '720',
        '1080': '1080',
    }
    resol_name = resol_map.get(resolution, None)
    if not resol_name:
        resol_name = '1080'

    with open(f'{path}/{resol_name}_{filename}', 'wb') as f:
        pbar = tqdm(
            total=size/1024/1024,
            unit='MB',
            desc=f'正在下载 {resolution}P 视频至 {path}/{resol_name}_{filename}')
        for chunk in res.iter_content(chunk_size=1024):
            f.write(chunk)
            pbar.update(1024/1024)
            time.sleep(0.001)
        pbar.close()

def check_url(url):
    '''
    检查链接是否为mtv链接
    url: str 视频链接
    
    return: bool
    '''
    if 'music.163.com/#/mv?id=' in url:
        return True
    return False

def main():
    st.title('网易云mtv下载工具')

    # 输入链接
    url = st.text_input('请输入mtv链接')

    # 检查链接是否正确
    if not check_url(url):
        st.error('输入链接不是网易云mtv链接')
        return

    # 获取视频名称
    res = requests.get(url)
    title = res.text.split('<title>')[1].split(' - ')[0]
    st.write(f'视频名称：{title}')

    # 选择文件储存路径
    path = st.text_input('请选择储存路径', 'C:/Users/xxx/Downloads')
    st.write(f'储存路径：{path}')

    # 选择分辨率
    resolution = st.selectbox('请选择分辨率', ('240', '480', '720', '1080'))
    st.write(f'已选择分辨率：{resolution}P')

    # 开始下载
    if st.button('开始下载'):
        download(url, title, path, resolution)
        st.success('下载完成')

if __name__ == '__main__':
    main()