# -*- coding: utf-8 -*-
import sys
import requests
import re
import os
import time
import json
from datetime import datetime
from bs4 import BeautifulSoup

BASE_PATH = ""

def main():
    global BASE_PATH
    BASE_PATH = "./file/result/file_" + str(int(round(time.time() * 1000)))
    
    while True:
        r = get_dy_video()
        if r == 0:
            break

def get_dy_video():
    url = input("\n请输入抖音视频链接(输入0退出): ")
    if url == "0":
        return 0
    
    print(f"\n=== 开始处理 ===")
    print(f"输入链接: {url}")

    # 获取视频信息
    print("\n[1/4] 正在解析视频信息...")
    v_info = get_dy_video_info(url)
    if not v_info:
        print("❗ 获取视频信息失败")
        return 1
    
    # 构建下载链接
    print("\n[2/4] 构建下载地址...")
    try:
        v_id = v_info["video"]["play_addr"]["uri"]
        video_url = f'https://aweme.snssdk.com/aweme/v1/play/?video_id={v_id}&line=0&ratio=720p'
        print(f"原始下载地址: {video_url}")
    except KeyError:
        print("❗ 视频信息结构异常")
        return 1

    # 获取最终下载地址
    print("\n[3/4] 解析真实下载地址...")
    res_url = get_redirect_url(video_url)
    if not res_url:
        print("❗ 无法获取下载地址")
        return 1
    print(f"最终下载地址: {res_url[:80]}...")  # 显示部分地址避免刷屏

    # 下载视频
    print("\n[4/4] 开始下载文件...")
    save_path = BASE_PATH + ".mp4"
    download_media(res_url, save_path)
    
    print("\n------ 🎉 下载完成 ------")
    print(f"保存路径: {save_path}\n")
    return 1

def get_dy_video_info(url):
    print("\n[解析步骤] 正在处理短链接...")
    redirect_url = get_redirect_url(url)
    if not redirect_url:
        print("❗ 短链接解析失败")
        return None
    print(f"重定向地址: {redirect_url}")

    # 提取视频ID
    print("\n[解析步骤] 提取视频ID...")
    video_id = None
    patterns = [
        r'/video/(\d+)', 
        r'/(\d+)(?:/|\?|$)',
        r'(\d{19})'
    ]
    for pattern in patterns:
        match = re.search(pattern, redirect_url)
        if match:
            video_id = match.group(1)
            print(f"使用正则 {pattern} 匹配到ID: {video_id}")
            break
    if not video_id:
        print("❗ 无法提取视频ID")
        return None

    # 请求视频页面
    print("\n[解析步骤] 请求视频页面...")
    headers = {
        "User-Agent": "Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 Mobile/15E148 Safari/604.1",
        "Referer": "https://www.douyin.com/",
        # "Cookie": "在此处添加你的Cookie"  # 如果需要请取消注释
    }
    try:
        response = requests.get(redirect_url, headers=headers, timeout=15)
        print(f"响应状态码: {response.status_code}")
        print(f"响应长度: {len(response.text)}字节")
    except Exception as e:
        print(f"❗ 请求异常: {str(e)}")
        return None

    # 解析页面数据
    print("\n[解析步骤] 提取RENDER_DATA...")
    soup = BeautifulSoup(response.text, 'html.parser')
    script_tag = soup.find('script', id='RENDER_DATA')
    if not script_tag:
        print("❗ 未找到RENDER_DATA")
        return None
    
    # 解析JSON数据
    print("\n[解析步骤] 解析视频数据...")
    try:
        decoded_data = requests.utils.unquote(script_tag.string)
        video_data = json.loads(decoded_data)
        print("找到有效视频数据")
    except Exception as e:
        print(f"❗ 数据解析失败: {str(e)}")
        return None

    # 提取视频地址
    print("\n[解析步骤] 定位下载地址...")
    try:
        uri = (
            video_data.get('aweme', {}).get('detail', {}).get('video', {})
            .get('playAddr', [{}])[0].get('src', '').split('/')[-1].split('.')[0]
        )
        if not uri:
            print("❗ 未找到视频URI")
            return None
        print(f"成功提取视频ID: {uri}")
        return {"video": {"play_addr": {"uri": uri}}
    except Exception as e:
        print(f"❗ 地址提取失败: {str(e)}")
        return None

def download_media(url, save_path):
    try:
        print(f"开始下载: {url[:60]}...")
        headers = {"User-Agent": "Mozilla/5.0"}
        response = requests.get(url, stream=True, headers=headers)
        
        if response.status_code != 200:
            print(f"下载失败，状态码: {response.status_code}")
            return False
        
        total_size = int(response.headers.get('content-length', 0))
        print(f"文件大小: {total_size//1024} KB")
        
        with open(save_path, "wb") as file:
            downloaded = 0
            start_time = time.time()
            for chunk in response.iter_content(chunk_size=1024*1024):
                if chunk:
                    file.write(chunk)
                    downloaded += len(chunk)
                    # 显示下载进度
                    speed = downloaded / (time.time() - start_time) / 1024
                    print(f"\r下载进度: {downloaded//1024}KB [{speed:.1f}KB/s]", end="")
            print()
        return True
    except Exception as e:
        print(f"\n下载异常: {str(e)}")
        return False

def get_redirect_url(url):
    try:
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
            "Referer": "https://www.douyin.com/"
        }
        response = requests.get(url, headers=headers, allow_redirects=True, timeout=10)
        print(f"重定向次数: {len(response.history)}次")
        return response.url
    except Exception as e:
        print(f"重定向解析失败: {str(e)}")
        return None

if __name__ == "__main__":
    main()