#!/usr/bin/env python3
import argparse
import os
import subprocess
import re
import sys

SUPPORTED_EXTENSIONS = (".mp4", ".mkv", ".mov", ".avi")

def get_duration(file_path):
    result = subprocess.run([
        "ffprobe", "-v", "error", "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1", file_path
    ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    return float(result.stdout.strip())

def encode(input_file, codec, remove_source=False, save_log=False):
    if codec == "x265":
        ffmpeg_codec = "hevc_nvenc"
        cq = 32
        folder = "h265"
        extra_params = [
        ]

    elif codec == "av1":
        ffmpeg_codec = "av1_nvenc"
        cq = 32
        folder = "av1"
        extra_params = [
            "-pix_fmt", "yuv420p",
        ]

    else:
        raise ValueError(f"Unsupported codec: {codec}")

    try:
        duration = get_duration(input_file)
    except Exception:
        print(f"\n❌ Could not read duration for {input_file}")
        return

    filename = os.path.basename(input_file)
    name, _ = os.path.splitext(filename)
    outdir = os.path.join(os.path.dirname(input_file), folder)
    os.makedirs(outdir, exist_ok=True)
    output_file = os.path.join(outdir, f"{name}.mkv")

    if save_log:
        log_file = os.path.join(outdir, f"{name}.log")
        log = open(log_file, "w")
    else:
        log = subprocess.DEVNULL  # Pas de log

    cmd = [
        "ffmpeg",
        "-i", input_file,
        "-map", "0",
    ] + extra_params + [
        "-c:v", ffmpeg_codec,
        "-preset", "p4",
        "-cq", str(cq),
        "-rc", "vbr",
        "-b:v", "0",
        "-c:a", "copy",
        "-c:s", "copy",
        "-y", output_file
    ]

    print(f"\n🎬 Encoding {filename} → {folder}/{name}.mkv with codec [{ffmpeg_codec}]")
    process = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)

    time_re = re.compile(r"time=(\d+):(\d+):([\d.]+)")
    last_percent = -1

    while True:
        line = process.stderr.readline()
        if not line:
            break
        if save_log:
            log.write(line)
        match = time_re.search(line)
        if match:
            h, m, s = map(float, match.groups())
            current = h * 3600 + m * 60 + s
            percent = int((current / duration) * 100)
            if percent != last_percent:
                print(f"\r⏳ Progress: {percent}%", end='', flush=True)
                last_percent = percent

    process.wait()

    if save_log:
        log.close()

    if process.returncode != 0:
        print(f"\n❌ [Error] FFmpeg failed for {filename}.")
        if save_log:
            print(f"  ➔ Check the log file: {log_file}")
        if os.path.exists(output_file):
            os.remove(output_file)
    else:
        print("\n✅ Done.")
        if remove_source:
            try:
                os.remove(input_file)
                print(f"🗑️ Source file {input_file} removed.")
            except Exception as e:
                print(f"⚠️ Could not delete source file {input_file}: {e}")

def encode_batch(directory, codec, remove_source=False, save_log=False):
    if not os.path.isdir(directory):
        print(f"❌ Not a valid directory: {directory}")
        return

    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(SUPPORTED_EXTENSIONS):
                filepath = os.path.join(root, file)
                try:
                    encode(filepath, codec, remove_source, save_log)
                except Exception as e:
                    print(f"\n❌ Error with file {file}: {e}")
                    continue

def main():
    parser = argparse.ArgumentParser(description="Encode video(s) to x265 or AV1.")
    parser.add_argument("input", nargs="?", help="Path to input file")
    parser.add_argument("-d", "--directory", help="Path to a directory for batch encoding")
    parser.add_argument("--codec", choices=["x265", "av1"], default="x265", help="Codec to use")
    parser.add_argument("--remove-source", action="store_true", help="Remove the source file after successful encoding")
    parser.add_argument("--save-log", action="store_true", help="Save ffmpeg logs to a .log file")
    args = parser.parse_args()

    if args.directory:
        encode_batch(args.directory, args.codec, args.remove_source, args.save_log)
    elif args.input:
        if not os.path.isfile(args.input):
            print(f"❌ File not found: {args.input}")
            sys.exit(1)
        encode(args.input, args.codec, args.remove_source, args.save_log)
    else:
        print("❌ Please provide a file path or use -d for a directory.")
        parser.print_help()

if __name__ == "__main__":
    main()