main.py

# Copyright 2025 张昊
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import os
import json
import shutil
import subprocess
from pathlib import Path

# 初始化缓存字典
size_cache = {}

def check(prompt, default=False):
    """获取用户选择"""
    options = "[yes]/no" if default else "yes/[no]"
    prompt = f"{prompt} ({options})"

    while True:
        user_input = input(prompt).lower()

        if user_input.lower() in ("", "yes", "y"):
            return True
        elif user_input.lower() in ("no", "n"):
            return False
        else:
            print("无效输入！请重新输入")

def sanitize_filename(name):
    """移除文件名中的非法字符"""
    return name.replace(' ', '_').translate(str.maketrans('', '', '<>:"/\\|?*\'\t\n\r'))

def parse_selection(input_str, max_index):
    """解析用户的选择输入"""
    selected_indices = set()
    
    # 处理逗号分隔的多个选择
    parts = [part.strip() for part in input_str.split(',')]
    
    for part in parts:
        # 处理范围选择
        if '-' in part:
            try:
                start, end = map(int, part.split('-'))
                # 转换为0开头的索引
                start = max(1, min(start, max_index)) - 1
                end = max(1, min(end, max_index)) - 1
                # 确保范围有效
                if start <= end:
                    selected_indices.update(range(start, end + 1))
            except ValueError:
                print(f"忽略无效范围: {part}")
        # 处理单个选择
        else:
            try:
                index = int(part) - 1
                if 0 <= index < max_index:
                    selected_indices.add(index)
            except ValueError:
                print(f"忽略无效输入: {part}")
    
    return sorted(selected_indices)


def get_file_size(path):
    """获取文件大小"""
    bytes = get_file_bytes(path)
    if bytes == 0:
        return "0 B"
    
    size = bytes
    for unit in ['B', 'KB', 'MB', 'GB']:
        if size < 1024.0:
            return f"{size:.1f} {unit}"
        size /= 1024.0
    return f"{size:.1f} GB"
    
def run_command(command):
    """使用 rish 执行命令，如果失败尝试直接执行"""
    try:
        # 首先尝试使用 rish
        result = subprocess.run(['sh', '../rish', '-c', command], 
                               capture_output=True, text=True, check=True)
        print(f"\n[DEBUG]run:'sh ../rish -c {command}'\nresult:'{result}'\n")
        return result.returncode == 0
    except (FileNotFoundError, subprocess.CalledProcessError):
        try:
            # 如果 rish 不可用，尝试直接执行
            print("尝试直接执行命令...")
            result = subprocess.run(command, shell=True, capture_output=True, text=True, check=True)
            return result.returncode == 0
        except Exception as e:
            print(f"命令执行失败: {e}")
            return False

def copy_with_fallback(src, dest):
    """复制文件，先尝试 rish ，失败则尝试直接复制"""
    # 尝试使用 rish 命令复制
    if run_command(f"cp '{src}' '{dest}'"):
        return True
    
    # 如果 rish 命令失败，尝试直接复制
    try:
        print("尝试直接复制文件...")
        shutil.copy2(src, dest)
        return True
    except Exception as e:
        print(f"文件复制失败: {e}")
        return False

def process_entry(entry, output_base_dir, video_temp_dir, idx, total):
    """处理单个视频条目"""
    
    title, part, owner_name, type_tag, avid, cid, is_bangumi = entry
    
    print(f"\n{'='*40}")
    print(f"处理条目 [{idx+1}/{total}]: {title} - {part}")
    print(f"作者: {owner_name}")
    
    # 处理路径
    sanitized_title = sanitize_filename(title)
    sanitized_part = sanitize_filename(part)
    
    # 根据视频类型设置正确的源文件路径
    if is_bangumi:
        source_dir = f"/storage/emulated/0/Android/data/tv.danmaku.bili/download/s_{avid}/{cid}/{type_tag}"
    else:
        source_dir = f"/storage/emulated/0/Android/data/tv.danmaku.bili/download/{avid}/c_{cid}/{type_tag}"
    
    audio_src = os.path.join(source_dir, "audio.m4s")
    video_src = os.path.join(source_dir, "video.m4s")
    
    # 设置临时文件路径
    audio_temp = os.path.join(video_temp_dir, f"audio_{avid}_{cid}.m4s")
    video_temp = os.path.join(video_temp_dir, f"video_{avid}_{cid}.m4s")
    
    # 创建输出目录
    output_dir = os.path.join(output_base_dir, sanitized_title)
    os.makedirs(output_dir, exist_ok=True)
    
    # 设置输出文件路径
    output_file = os.path.join(output_dir, f"{sanitized_part}.mp4")
    
    # 显示文件大小
    print("\n文件大小:")
    print(f"音频文件: {get_file_size(audio_src)}")
    print(f"视频文件: {get_file_size(video_src)}")
    
    try:
        # 复制音频文件
        print("\n复制音频文件中...")
        if not copy_with_fallback(audio_src, audio_temp):
            print("音频文件复制失败，跳过此条目")
            return False
        
        # 复制视频文件
        print("复制视频文件中...")
        if not copy_with_fallback(video_src, video_temp):
            print("视频文件复制失败，跳过此条目")
            return False
        
        # 合并文件
        print("\n合并音视频中...")
        command = [
            'ffmpeg',
            '-i', video_temp,
            '-i', audio_temp,
            '-c:v', 'copy',
            '-c:a', 'copy',
            '-y', output_file
        ]
        
        try:
            # 运行ffmpeg
            process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                universal_newlines=True
            )
            
            # 打印输出
            for line in process.stdout:
                print(line)
            
            # 等待进程结束
            return_code = process.wait()
            
            if return_code == 0:
                print(f"\n✅ 合并成功: {output_file}")
                print(f"输出文件大小: {get_file_size(output_file)}")
                return True
            else:
                print(f"\n❌ 合并失败 (错误码: {return_code})")
                return False
                
        except Exception as e:
            print(f"合并过程中出错: {e}")
            return False
            
    except Exception as e:
        print(f"处理过程中出错: {e}")
        return False
    finally:
        # 清理临时文件
        if os.path.exists(audio_temp):
            os.remove(audio_temp)
        if os.path.exists(video_temp):
            os.remove(video_temp)
         
def get_file_bytes(path):
    """获取文件的大小"""
    # 优先在缓存中获取
    if path in size_cache:
        return size_cache[path]
    
    try:
        # 尝试使用 rish 命令获取
        command = f'stat -c "%s" "{path}"'
        result = subprocess.run(['sh', '../rish', '-c', command], 
                               capture_output=True, text=True)
        
        """
        print(f"\n[DEBUG]run:'sh ../rish -c {command}'\n"
              f"returncode: {result.returncode}\n"
              f"stdout: '{result.stdout.strip()}'\n"
              f"stderr: '{result.stderr.strip()}'\n")
        """
        
        # 此处未知原因导致文件大小的输出在stderr中
        # 优先检查stdout，若为空则检查stderr
        output = result.stdout.strip() or result.stderr.strip()
        
        if result.returncode == 0 and output.isdigit():
            size = int(output)
            size_cache[path] = size  # 更新缓存
            return size
        elif os.path.exists(path):
            # 尝试直接获取大小
            size = os.path.getsize(path)
            size_cache[path] = size
            return size
        else:
            return 0
    except Exception as e:
        print(f"获取文件大小出错: {e}")
        return 0
        
        
def calculate_total_size(entries, selected_indices):
    """计算选中条目的总大小"""
    total_bytes = 0
    
    for idx in selected_indices:
        title, part, owner_name, type_tag, avid, cid, is_bangumi = entries[idx]
        
        # 根据视频类型设置源路径
        if is_bangumi:
            source_dir = f"/storage/emulated/0/Android/data/tv.danmaku.bili/download/s_{avid}/{cid}/{type_tag}"
        else:
            source_dir = f"/storage/emulated/0/Android/data/tv.danmaku.bili/download/{avid}/c_{cid}/{type_tag}"
        
        audio_src = os.path.join(source_dir, "audio.m4s")
        video_src = os.path.join(source_dir, "video.m4s")
        
        total_bytes += get_file_bytes(audio_src)
        total_bytes += get_file_bytes(video_src)
    
    # 将文件大小转换为带单位的形式
    if total_bytes == 0:
        return "0 B"
    
    size = total_bytes * 2 # 临时文件与最终输出总共的两倍大小
    for unit in ['B', 'KB', 'MB', 'GB']:
        if size < 1024.0:
            return f"{size:.1f} {unit}"
        size /= 1024.0
    return f"{size:.1f} GB"
    
def main():
    # 设置路径
    json_dir = "/storage/emulated/0/RishBiliCacheMerge/json"
    output_base_dir = "/storage/emulated/0/RishBiliCacheMerge/output"
    video_temp_dir = "/storage/emulated/0/RishBiliCacheMerge/video"
    
    entries = []
    
    print("正在扫描JSON文件...")
    # 遍历目录中的所有JSON文件
    for filename in os.listdir(json_dir):
        if filename.endswith(".json"):
            filepath = os.path.join(json_dir, filename)
            try:
                with open(filepath, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    
                    # 检查视频类型
                    if "ep" in data and isinstance(data["ep"], dict) and "season_id" in data:
                        title = data.get("title", "")
                        part = data["ep"].get("index", "")  
                        if not part:
                            part = data["ep"].get("index_title", "")
                        owner_name = data["ep"].get("season_title", "未知")
                        type_tag = data.get("type_tag", "")
                        season_id = data.get("season_id", "")
                        episode_id = data["ep"].get("episode_id", "")
                        
                        # 使用season_id和episode_id作为avid和cid
                        avid = str(season_id)
                        cid = str(episode_id)
                        is_bangumi = True
                    else:
                        title = data.get("title", "")
                        part = data.get("page_data", {}).get("part", "")
                        owner_name = data.get("owner_name", "")
                        type_tag = data.get("type_tag", "")
                        avid = data.get("avid", "")
                        cid = data.get("page_data", {}).get("cid", "")
                        is_bangumi = False
                    
                    # 确保所有值都是字符串
                    title = str(title)
                    part = str(part)
                    owner_name = str(owner_name)
                    type_tag = str(type_tag)
                    avid = str(avid)
                    cid = str(cid)
                    
                    # 将提取的数据存入元组并添加到列表
                    # 添加is_bangumi标志以区分视频类型
                    entries.append((title, part, owner_name, type_tag, avid, cid, is_bangumi))
                    
            except (json.JSONDecodeError, KeyError, TypeError) as e:
                print(f"解析文件 {filename} 时出错: {e}")
                continue

    # 检查是否找到有效JSON
    if not entries:
        print("未找到有效的JSON文件")
        return

    # 显示选择菜单
    print("\n" + "=" * 50)
    print("请选择视频条目：")
    print("=" * 50)
    for i, entry in enumerate(entries, 1):
        title, part, owner_name, type_tag, avid, cid, is_bangumi = entry
        
        # 根据视频类型设置源文件路径
        if is_bangumi:
            source_dir = f"/storage/emulated/0/Android/data/tv.danmaku.bili/download/s_{avid}/{cid}/{type_tag}"
        else:
            source_dir = f"/storage/emulated/0/Android/data/tv.danmaku.bili/download/{avid}/c_{cid}/{type_tag}"
        
        audio_src = os.path.join(source_dir, "audio.m4s")
        video_src = os.path.join(source_dir, "video.m4s")
        
        # 获取文件大小
        audio_size = get_file_size(audio_src)
        video_size = get_file_size(video_src)
        total_size = f"音频: {audio_size}, 视频: {video_size}"
        
        print(f"{i}. {title[:30]}{'...' if len(title) > 30 else ''}")
        print(f"   ├─ 分集: {part[:20]}{'...' if len(part) > 20 else ''}")
        print(f"   ├─ 作者: {owner_name}")
        print(f"   └─ 大小: {total_size}")
        if i < len(entries):
            print("-" * 50)

    # 获取用户选择
    print("\n输入说明:")
    print("  - 单个条目: 输入序号 (例如: 3)")
    print("  - 多个条目: 用英文逗号分隔 (例如: 1,3,5)")
    print("  - 范围选择: 用连字符 (例如: 2-5)")
    print("  - 全部选择: 输入 'all'")
    print()
    
    user_input = input("请输入选择: ").strip()
    
    if user_input.lower() == 'all':
        selected_indices = list(range(len(entries)))
    else:
        selected_indices = parse_selection(user_input, len(entries))
    
    if not selected_indices:
        print("未选择任何有效条目")
        return
    
    # 显示选中的条目
    print("\n已选择以下条目:")
    for idx in selected_indices:
        title, part, owner_name, *_ = entries[idx]
        print(f"\n  - {title}\n  - {part}\n  [作者: {owner_name}]")
    
    print("\n计算大小中...")
    # 计算并显示总大小
    total_size = calculate_total_size(entries, selected_indices)
    print(f"\n预计需要预留空间: {total_size}")
    print("请确保设备有足够存储空间")
    
    if not check("\n确认开始处理?", default=True):
        print("操作已取消")
        return
    
    # 处理所有选中的条目
    success_count = 0
    total_count = len(selected_indices)
    
    for i, idx in enumerate(selected_indices):
        entry = entries[idx]
        if process_entry(entry, output_base_dir, video_temp_dir, i, total_count):
            success_count += 1
    
    print(f"\n处理完成! 成功: {success_count}/{total_count}")
    print(f"文件输出目录：{output_base_dir}")
    
    # 清理操作
    if success_count > 0:
        if check("\n是否删除JSON临时文件?", default=True):
            print("删除JSON文件...")
            shutil.rmtree(json_dir, ignore_errors=True)
            print("JSON目录已删除")
        
        if check("是否删除视频临时目录?", default=True):
            print("删除视频临时目录...")
            shutil.rmtree(video_temp_dir, ignore_errors=True)
            print("视频临时目录已删除")
    else:
        print("没有成功处理的条目，跳过清理")

if __name__ == "__main__":
    main()