compress.py

import sys
import subprocess
import json
import os
import re

def send_json(data):
    print(f"JSON:{json.dumps(data)}")
    sys.stdout.flush()

# --- CONFIGURATION ---
def resolve_bin_dir():
    """
    When running as a PyInstaller onefile EXE, __file__ points to a temp dir.
    Use the directory of the executable instead so we can find the external bin folder
    that Electron bundles alongside compress.exe.
    """
    if getattr(sys, 'frozen', False):
        base = os.path.dirname(sys.executable)
    else:
        base = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(base, "bin")

BIN_DIR = resolve_bin_dir()
FFMPEG_BIN = os.path.join(BIN_DIR, "ffmpeg.exe")
FFPROBE_BIN = os.path.join(BIN_DIR, "ffprobe.exe")
AUDIO_BITRATE = 128

if not os.path.exists(FFMPEG_BIN):
    send_json({'status': 'error', 'msg': f'ffmpeg not found at {FFMPEG_BIN}'})
    sys.exit(1)

def get_duration(file_path):
    cmd = [FFPROBE_BIN, "-v", "error", "-show_entries", "format=duration", "-of", "json", file_path]
    try:
        result = subprocess.run(cmd, stdout=subprocess.PIPE, text=True)
        return float(json.loads(result.stdout)['format']['duration'])
    except:
        return 0

def get_resolution(file_path):
    cmd = [FFPROBE_BIN, "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=height", "-of", "json", file_path]
    try:
        result = subprocess.run(cmd, stdout=subprocess.PIPE, text=True)
        data = json.loads(result.stdout)
        return int(data['streams'][0]['height'])
    except:
        return 0

def compress_video(input_path, output_path, target_mb, quality_level, fps=30):
    duration = get_duration(input_path)
    height = get_resolution(input_path)

    if duration == 0:
        send_json({'status': 'error', 'msg': 'Could not read video duration'})
        return

    # 1. Bitrate Math
    target_total_bits = target_mb * 8192 * 0.95 
    video_bitrate = int((target_total_bits / duration) - AUDIO_BITRATE)
    if video_bitrate < 100: video_bitrate = 100

    # 2. Output frame rate (30 or 60)
    fps = 60 if str(fps) == "60" else 30
    # At 30fps we use half the video bitrate so file size is ~half (same duration, half the frames)
    if fps == 30:
        video_bitrate = max(100, video_bitrate // 2)
    fps_filter = f"fps={fps}"

    # 3. Auto-Downscaling + combine with fps in a single -vf
    vf_parts = []
    if video_bitrate < 700 and height > 480:
        vf_parts.append("scale=-2:480")
        send_json({'status': 'scaling', 'msg': f'Bitrate too low ({video_bitrate}k). Downscaling to 480p.'})
    elif video_bitrate < 1500 and height > 720:
        vf_parts.append("scale=-2:720")
        send_json({'status': 'scaling', 'msg': f'Bitrate low ({video_bitrate}k). Downscaling to 720p.'})
    vf_parts.append(fps_filter)
    vf_str = ",".join(vf_parts)

    # 4. Quality Presets (0=Fast, 1=Medium, 2=Slow/Best)
    preset_map = { "0": "veryfast", "1": "medium", "2": "veryslow" }
    preset = preset_map.get(str(quality_level), "medium")

    # 5. Command: -vf does frame conversion, -r before output sets stream/timebase so container reports correct fps
    cmd = [
        FFMPEG_BIN, "-y", "-i", input_path,
        "-vf", vf_str,
        "-c:v", "libx264", "-pix_fmt", "yuv420p",
        "-b:v", f"{video_bitrate}k", "-maxrate", f"{video_bitrate}k", "-bufsize", f"{video_bitrate*2}k",
        "-preset", preset,
        "-r", str(fps),
        "-c:a", "aac", "-b:a", f"{AUDIO_BITRATE}k", "-ac", "2",
        output_path
    ]

    send_json({'status': 'start', 'target_bitrate': video_bitrate, 'fps': fps})

    # 5. Run & Monitor
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, universal_newlines=True)
    time_pattern = re.compile(r"time=(\d{2}:\d{2}:\d{2}\.\d{2})")

    for line in process.stdout:
        match = time_pattern.search(line)
        if match:
            h, m, s = match.group(1).split(':')
            seconds = int(h)*3600 + int(m)*60 + float(s)
            percent = min(round((seconds / duration) * 100, 1), 99.9)
            send_json({'status': 'progress', 'percent': percent})
    
    process.wait()
    if process.returncode == 0:
        send_json({'status': 'done', 'path': output_path})
    else:
        send_json({'status': 'error', 'msg': 'FFmpeg error'})

if __name__ == "__main__":
    try:
        # ARGS: input, output, target_mb, quality_level, fps (30 or 60)
        fps_arg = sys.argv[5] if len(sys.argv) > 5 else "30"
        compress_video(sys.argv[1], sys.argv[2], float(sys.argv[3]), sys.argv[4], fps_arg)
    except Exception as e:
        send_json({'status': 'error', 'msg': str(e)})