#!/usr/bin/env python3
import os
import sys
import subprocess
import shutil
import tempfile
import json
import re
from datetime import datetime
from pathlib import Path

# List of required external tools
REQUIRED_TOOLS = [
    "ffmpeg", "ffprobe", "mkvmerge", "mkvpropedit",
    "sox", "opusenc", "mediainfo", "SvtAv1EncApp2", "HandBrakeCLI"
]

# Set the SVT-AV1 encoder binary here for easy switching
SVTAV1_BIN = "SvtAv1EncApp2"

# SVT-AV1-Essential (nekotrix) encoding parameters (edit here as needed)
SVT_AV1_PARAMS = {
    "speed": "slow",              # "slower", "slow", "medium", "fast", "faster"
    "quality": "high",            # "higher", "high", "medium", "low", "lower"
    "film-grain": 6,
    "color-primaries": 1,
    "transfer-characteristics": 1,
    "matrix-coefficients": 1,
    "scd": 1,                     # Scene change detection ON (recommended)
    "auto-tiling": 1,             # Auto tiling ON (recommended)
    "tune": 1,                    # 0 = VQ, 1 = PSNR, 2 = SSIM
    "progress": 2,                # Detailed progress output
}
SVT_AV1_PARAMS_LIST = []
for key, value in SVT_AV1_PARAMS.items():
    SVT_AV1_PARAMS_LIST.extend([f"--{key}", str(value)])

DIR_COMPLETED = Path("completed")
DIR_ORIGINAL = Path("original")
DIR_CONV_LOGS = Path("conv_logs")

REMUX_CODECS = {"aac", "opus"}

def check_tools():
    """Check if all required tools are available in PATH."""
    for tool in REQUIRED_TOOLS:
        if shutil.which(tool) is None:
            print(f"Required tool '{tool}' not found in PATH.")
            sys.exit(1)

def run_cmd(cmd, capture_output=False, check=True):
    """Run a subprocess command, optionally capturing output."""
    try:
        if capture_output:
            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=check, text=True)
            return result.stdout
        else:
            subprocess.run(cmd, check=check)
    except subprocess.CalledProcessError as e:
        print(f"Command failed: {' '.join(map(str, cmd))}")
        print(f"Return code: {e.returncode}")
        if capture_output:
            print(f"Output: {e.output}")
        raise

def convert_audio_track(index, ch, lang, audio_temp_dir, source_file, should_downmix):
    """Extract, normalize, and encode an audio track to Opus."""
    audio_temp_path = Path(audio_temp_dir)
    temp_extracted = audio_temp_path / f"track_{index}_extracted.flac"
    temp_normalized = audio_temp_path / f"track_{index}_normalized.flac"
    final_opus = audio_temp_path / f"track_{index}_final.opus"

    print(f"    - Extracting Audio Track #{index} to FLAC...")
    ffmpeg_args = [
        "ffmpeg", "-v", "quiet", "-stats", "-y", "-i", str(source_file), "-map", f"0:{index}"
    ]
    # Downmix if needed
    if should_downmix and ch >= 6:
        if ch == 6:
            ffmpeg_args += ["-af", "pan=stereo|c0=c2+0.30*c0+0.30*c4|c1=c2+0.30*c1+0.30*c5"]
        elif ch == 8:
            ffmpeg_args += ["-af", "pan=stereo|c0=c2+0.30*c0+0.30*c4+0.30*c6|c1=c2+0.30*c1+0.30*c5+0.30*c7"]
        else:
            ffmpeg_args += ["-ac", "2"]
    ffmpeg_args += ["-c:a", "flac", str(temp_extracted)]
    run_cmd(ffmpeg_args)

    print(f"    - Normalizing Audio Track #{index} with SoX...")
    run_cmd([
        "sox", str(temp_extracted), str(temp_normalized), "-S", "--temp", str(audio_temp_path), "--guard", "gain", "-n"
    ])

    # Set bitrate based on channel count and downmixing
    is_being_downmixed = should_downmix and ch >= 6
    if is_being_downmixed:
        bitrate = "128k"
    else:
        if ch == 2:
            bitrate = "128k"
        elif ch == 6:
            bitrate = "256k"
        elif ch == 8:
            bitrate = "384k"
        else:
            bitrate = "96k"

    print(f"    - Encoding Audio Track #{index} to Opus at {bitrate}...")
    run_cmd([
        "opusenc", "--vbr", "--bitrate", bitrate, str(temp_normalized), str(final_opus)
    ])
    return final_opus

def convert_video(source_file_base, source_file_full, is_vfr, target_cfr_fps_for_handbrake, autocrop_filter=None):
    """
    Convert video to AV1 using ffmpeg and SvtAv1EncApp.
    If VFR, use HandBrakeCLI to convert to CFR first.
    """
    print("  --- Starting Video Processing ---")
    encoded_video_file = Path(f"temp-{source_file_base}.ivf")
    handbrake_cfr_intermediate_file = None
    current_input = Path(source_file_full)

    # VFR-to-CFR conversion using HandBrakeCLI if needed
    if is_vfr and target_cfr_fps_for_handbrake:
        print(f"    - Source is VFR. Converting to CFR ({target_cfr_fps_for_handbrake}) with HandBrakeCLI...")
        handbrake_cfr_intermediate_file = Path(f"{source_file_base}.cfr_temp.mkv")
        handbrake_args = [
            "HandBrakeCLI",
            "--input", str(source_file_full),
            "--output", str(handbrake_cfr_intermediate_file),
            "--cfr",
            "--rate", str(target_cfr_fps_for_handbrake),
            "--encoder", "x264_10bit",
            "--quality", "0",
            "--encoder-preset", "superfast",
            "--encoder-tune", "fastdecode",
            "--audio", "none",
            "--subtitle", "none",
            "--crop-mode", "none"
        ]
        print(f"    - Running HandBrakeCLI: {' '.join(handbrake_args)}")
        try:
            run_cmd(handbrake_args)
            if handbrake_cfr_intermediate_file.exists() and handbrake_cfr_intermediate_file.stat().st_size > 0:
                print(f"    - HandBrake VFR to CFR conversion successful: {handbrake_cfr_intermediate_file}")
                current_input = handbrake_cfr_intermediate_file
            else:
                print(f"    - Warning: HandBrakeCLI VFR-to-CFR conversion failed or produced an empty file. Proceeding with original source.")
                handbrake_cfr_intermediate_file = None
        except subprocess.CalledProcessError as e:
            print(f"    - Error during HandBrakeCLI execution: {e}")
            print(f"    - Proceeding with original source.")
            handbrake_cfr_intermediate_file = None

    print("    - Starting AV1 encode with ffmpeg and SvtAv1EncApp...")

    # Probe video info for width, height, and frame rate
    ffprobe_cmd = [
        "ffprobe", "-v", "error", "-select_streams", "v:0",
        "-show_entries", "stream=width,height,r_frame_rate", "-of", "json",
        str(current_input)
    ]
    video_info_json = run_cmd(ffprobe_cmd, capture_output=True, check=True)
    video_info = json.loads(video_info_json)["streams"][0]
    width = video_info["width"]
    height = video_info["height"]
    frame_rate = video_info["r_frame_rate"]

    # Prepare ffmpeg command for yuv4mpegpipe
    ffmpeg_args = [
        "ffmpeg", "-hide_banner", "-v", "quiet", "-stats", "-y", "-i", str(current_input),
        "-pix_fmt", "yuv420p10le", "-f", "yuv4mpegpipe", "-"
    ]
    if autocrop_filter:
        match = re.match(r"crop=(\d+):(\d+):(\d+):(\d+)", autocrop_filter)
        if match:
            width = int(match.group(1))
            height = int(match.group(2))
        ffmpeg_args.insert(-2, "-vf")
        ffmpeg_args.insert(-2, autocrop_filter)

    svt_enc_args = [
        SVTAV1_BIN, "-i", "stdin",
        "--width", str(width), "--height", str(height),
        "--fps-num", frame_rate.split('/')[0], "--fps-den", frame_rate.split('/')[1],
        "-b", str(encoded_video_file)
    ] + SVT_AV1_PARAMS_LIST

    print(f"    - FFmpeg command: {' '.join(ffmpeg_args)}")
    print(f"    - SvtAv1EncApp command: {' '.join(svt_enc_args)}")

    # Pipe ffmpeg output to SvtAv1EncApp
    ffmpeg_process = subprocess.Popen(ffmpeg_args, stdout=subprocess.PIPE)
    svt_process = subprocess.Popen(svt_enc_args, stdin=ffmpeg_process.stdout)

    if ffmpeg_process.stdout:
        ffmpeg_process.stdout.close()

    svt_process.wait()
    ffmpeg_process.wait()

    if svt_process.returncode != 0:
        raise subprocess.CalledProcessError(svt_process.returncode, svt_enc_args)

    print("  --- Finished Video Processing ---")
    return encoded_video_file, handbrake_cfr_intermediate_file

def is_ffmpeg_decodable(file_path):
    """Quickly check if ffmpeg can decode the input file."""
    try:
        subprocess.run([
            "ffmpeg", "-v", "error", "-i", str(file_path), "-map", "0:a:0", "-t", "1", "-f", "null", "-"
        ], check=True)
        return True
    except subprocess.CalledProcessError:
        return False

# --- Cropdetect logic omitted for brevity (unchanged) ---

def main(no_downmix=False, autocrop=False):
    """Main batch-processing loop for MKV files."""
    check_tools()
    current_dir = Path(".")
    files_to_process = sorted(
        f for f in current_dir.glob("*.mkv")
        if not (f.name.endswith(".ut.mkv") or f.name.startswith("temp-") or f.name.startswith("output-") or f.name.endswith(".cfr_temp.mkv"))
    )
    if not files_to_process:
        print("No MKV files found to process. Exiting.")
        return
    DIR_COMPLETED.mkdir(exist_ok=True, parents=True)
    DIR_ORIGINAL.mkdir(exist_ok=True, parents=True)
    DIR_CONV_LOGS.mkdir(exist_ok=True, parents=True)
    while True:
        files_to_process = sorted(
            f for f in current_dir.glob("*.mkv")
            if not (f.name.endswith(".ut.mkv") or f.name.startswith("temp-") or f.name.startswith("output-") or f.name.endswith(".cfr_temp.mkv"))
        )
        if not files_to_process:
            print("No more .mkv files found to process in the current directory. The script will now exit.")
            break
        file_path = files_to_process[0]
        if not is_ffmpeg_decodable(file_path):
            print(f"ERROR: ffmpeg cannot decode '{file_path.name}'. Skipping this file.", file=sys.stderr)
            shutil.move(str(file_path), DIR_ORIGINAL / file_path.name)
            continue
        print("-" * shutil.get_terminal_size(fallback=(80, 24)).columns)
        log_file_name = f"{file_path.stem}.log"
        log_file_path = DIR_CONV_LOGS / log_file_name
        original_stdout_console = sys.stdout
        original_stderr_console = sys.stderr
        print(f"Processing: {file_path.name}", file=original_stdout_console)
        print(f"Logging output to: {log_file_path}", file=original_stdout_console)
        log_file_handle = None
        processing_error_occurred = False
        date_for_runtime_calc = datetime.now()
        try:
            log_file_handle = open(log_file_path, 'w', encoding='utf-8')
            sys.stdout = log_file_handle
            sys.stderr = log_file_handle
            print(f"STARTING LOG FOR: {file_path.name}")
            print(f"Processing started at: {date_for_runtime_calc}")
            print(f"Full input file path: {file_path.resolve()}")
            print("-" * shutil.get_terminal_size(fallback=(80, 24)).columns)
            input_file_abs = file_path.resolve()
            intermediate_output_file = current_dir / f"output-{file_path.name}"
            audio_temp_dir = None
            handbrake_intermediate_for_cleanup = None
            try:
                audio_temp_dir = tempfile.mkdtemp(prefix="anime_audio_")
                print(f"Audio temporary directory created at: {audio_temp_dir}")
                print(f"Analyzing file: {input_file_abs}")
                ffprobe_info_json = run_cmd([
                    "ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", "-show_format", str(input_file_abs)
                ], capture_output=True)
                ffprobe_info = json.loads(ffprobe_info_json)
                mkvmerge_info_json = run_cmd([
                    "mkvmerge", "-J", str(input_file_abs)
                ], capture_output=True)
                mkv_info = json.loads(mkvmerge_info_json)
                mediainfo_json = run_cmd([
                    "mediainfo", "--Output=JSON", "-f", str(input_file_abs)
                ], capture_output=True)
                media_info = json.loads(mediainfo_json)
                is_vfr = False
                target_cfr_fps_for_handbrake = None
                video_track_info = None
                # Find the video track in MediaInfo
                if media_info.get("media") and media_info["media"].get("track"):
                    for track in media_info["media"]["track"]:
                        if track.get("@type") == "Video":
                            video_track_info = track
                            break
                if video_track_info:
                    frame_rate_mode = video_track_info.get("FrameRate_Mode")
                    if frame_rate_mode and frame_rate_mode.upper() in ["VFR", "VARIABLE"]:
                        is_vfr = True
                        print(f"    - Detected VFR based on MediaInfo FrameRate_Mode: {frame_rate_mode}")
                        original_fps_str = video_track_info.get("FrameRate_Original_String")
                        if original_fps_str:
                            match = re.search(r'\((\d+/\d+)\)', original_fps_str)
                            if match:
                                target_cfr_fps_for_handbrake = match.group(1)
                            else:
                                target_cfr_fps_for_handbrake = video_track_info.get("FrameRate_Original")
                        if not target_cfr_fps_for_handbrake:
                            target_cfr_fps_for_handbrake = video_track_info.get("FrameRate_Original")
                        if not target_cfr_fps_for_handbrake:
                            target_cfr_fps_for_handbrake = video_track_info.get("FrameRate")
                            if target_cfr_fps_for_handbrake:
                                print(f"    - Using MediaInfo FrameRate ({target_cfr_fps_for_handbrake}) as fallback for HandBrake target FPS.")
                        if target_cfr_fps_for_handbrake:
                            print(f"    - Target CFR for HandBrake: {target_cfr_fps_for_handbrake}")
                            if isinstance(target_cfr_fps_for_handbrake, str) and "/" in target_cfr_fps_for_handbrake:
                                try:
                                    num, den = map(float, target_cfr_fps_for_handbrake.split('/'))
                                    target_cfr_fps_for_handbrake = f"{num / den:.3f}"
                                    print(f"    - Converted fractional FPS to decimal for HandBrake: {target_cfr_fps_for_handbrake}")
                                except ValueError:
                                    print(f"    - Warning: Could not parse fractional FPS '{target_cfr_fps_for_handbrake}'. HandBrakeCLI might fail.")
                                    is_vfr = False
                        else:
                            print("    - Warning: VFR detected, but could not determine target CFR from MediaInfo. Will attempt encode without HandBrake.")
                            is_vfr = False
                    else:
                        print(f"    - Video appears to be CFR or FrameRate_Mode not specified as VFR/Variable by MediaInfo.")
                autocrop_filter = None
                if autocrop:
                    print("--- Running autocrop detection ---")
                    autocrop_filter = detect_autocrop_filter(str(input_file_abs))
                    if autocrop_filter:
                        print(f"    - Autocrop filter detected: {autocrop_filter}")
                    else:
                        print("    - No crop needed or detected.")
                encoded_video_file, handbrake_intermediate_for_cleanup = convert_video(
                    file_path.stem, str(input_file_abs), is_vfr, target_cfr_fps_for_handbrake, autocrop_filter=autocrop_filter
                )

                print("--- Starting Audio Processing ---")
                processed_audio_files = []
                audio_tracks_to_remux = []
                audio_streams = [s for s in ffprobe_info.get("streams", []) if s.get("codec_type") == "audio"]

                mkv_audio_tracks = {t["id"]: t for t in mkv_info.get("tracks", []) if t.get("type") == "audio"}

                media_tracks_data = media_info.get("media", {}).get("track", [])
                mediainfo_audio_tracks = {int(t.get("StreamOrder", -1)): t for t in media_tracks_data if t.get("@type") == "Audio"}

                for stream in audio_streams:
                    stream_index = stream["index"]
                    codec = stream.get("codec_name")
                    channels = stream.get("channels", 2)
                    language = stream.get("tags", {}).get("language", "und")

                    mkv_track = None
                    for t in mkv_info.get("tracks", []):
                        if t.get("type") == "audio" and t.get("properties", {}).get("stream_id") == stream_index:
                            mkv_track = t
                            break
                    if not mkv_track:
                        mkv_track = mkv_info.get("tracks", [])[stream_index] if stream_index < len(mkv_info.get("tracks", [])) else {}

                    track_id = mkv_track.get("id", -1)
                    track_title = mkv_track.get("properties", {}).get("track_name", "")

                    audio_track_info = mediainfo_audio_tracks.get(stream_index)
                    track_delay = 0
                    delay_raw = audio_track_info.get("Video_Delay") if audio_track_info else None
                    if delay_raw is not None:
                        try:
                            delay_val = float(delay_raw)
                            if delay_val < 1:
                                track_delay = int(round(delay_val * 1000))
                            else:
                                track_delay = int(round(delay_val))
                        except Exception:
                            track_delay = 0

                    print(f"Processing Audio Stream #{stream_index} (TID: {track_id}, Codec: {codec}, Channels: {channels})")
                    if codec in REMUX_CODECS:
                        audio_tracks_to_remux.append(str(track_id))
                    else:
                        opus_file = convert_audio_track(
                            stream_index, channels, language, audio_temp_dir, str(input_file_abs), not no_downmix
                        )
                        processed_audio_files.append({
                            "Path": opus_file,
                            "Language": language,
                            "Title": track_title,
                            "Delay": track_delay
                        })

                print("--- Finished Audio Processing ---")

                print("Assembling final file with mkvmerge...")
                mkvmerge_args = ["mkvmerge", "-o", str(intermediate_output_file), str(encoded_video_file)]
                for file_info in processed_audio_files:
                    sync_switch = ["--sync", f"0:{file_info['Delay']}"] if file_info["Delay"] else []
                    mkvmerge_args += [
                        "--language", f"0:{file_info['Language']}",
                        "--track-name", f"0:{file_info['Title']}"
                    ] + sync_switch + [str(file_info["Path"])]

                source_copy_args = ["--no-video"]
                if audio_tracks_to_remux:
                    source_copy_args += ["--audio-tracks", ",".join(audio_tracks_to_remux)]
                else:
                    source_copy_args += ["--no-audio"]
                mkvmerge_args += source_copy_args + [str(input_file_abs)]
                run_cmd(mkvmerge_args)

                print("Moving files to final destinations...")
                shutil.move(str(file_path), DIR_ORIGINAL / file_path.name)
                shutil.move(str(intermediate_output_file), DIR_COMPLETED / file_path.name)

                print("Cleaning up persistent video temporary files (after successful processing)...")
                video_temp_files_on_success = [encoded_video_file]
                if handbrake_intermediate_for_cleanup and handbrake_intermediate_for_cleanup.exists():
                    video_temp_files_on_success.append(handbrake_intermediate_for_cleanup)

                for temp_vid_file in video_temp_files_on_success:
                    if temp_vid_file.exists():
                        print(f"    Deleting: {temp_vid_file}")
                        temp_vid_file.unlink(missing_ok=True)
                    else:
                        print(f"    Skipping (not found): {temp_vid_file}")

            except Exception as e:
                print(f"ERROR: An error occurred while processing '{file_path.name}': {e}", file=sys.stderr)
                original_stderr_console.write(f"ERROR during processing of '{file_path.name}': {e}\nSee log '{log_file_path}' for details.\n")
                processing_error_occurred = True
            finally:
                print("--- Starting Universal Cleanup (for this file) ---")
                print("  - Cleaning up disposable audio temporary directory...")
                if audio_temp_dir and Path(audio_temp_dir).exists():
                    shutil.rmtree(audio_temp_dir, ignore_errors=True)
                    print(f"    - Deleted audio temp dir: {audio_temp_dir}")
                elif audio_temp_dir:
                    print(f"    - Audio temp dir not found or already cleaned: {audio_temp_dir}")
                else:
                    print(f"    - Audio temp dir was not created.")

                print("  - Cleaning up intermediate output file (if it wasn't moved on success)...")
                if intermediate_output_file.exists():
                    if processing_error_occurred:
                        print(f"    - WARNING: Processing error occurred. Intermediate output file '{intermediate_output_file}' is being preserved at its original path for inspection.")
                    else:
                        print(f"    - INFO: Intermediate output file '{intermediate_output_file}' found at original path despite no errors (expected to be moved). Cleaning up.")
                        intermediate_output_file.unlink(missing_ok=True)
                        print(f"    - Deleted intermediate output file from original path: {intermediate_output_file}")
                else:
                    if not processing_error_occurred:
                        print(f"    - Intermediate output file successfully moved (not found at original path, as expected): {intermediate_output_file}")
                    else:
                        print(f"    - Processing error occurred, and intermediate output file '{intermediate_output_file}' not found at original path (likely not created or cleaned by another step).")
            
            print(f"FINISHED LOG FOR: {file_path.name}")

        finally:
            runtime = datetime.now() - date_for_runtime_calc
            runtime_str = str(runtime).split('.')[0]
            
            print(f"\nTotal runtime for this file: {runtime_str}")
            
            if sys.stdout != original_stdout_console:
                sys.stdout = original_stdout_console
            if sys.stderr != original_stderr_console:
                sys.stderr = original_stderr_console
            if log_file_handle:
                log_file_handle.close()
            
            if processing_error_occurred:
                original_stderr_console.write(f"File: {file_path.name}\n")
                original_stderr_console.write(f"Log: {log_file_path}\n")
                original_stderr_console.write(f"Runtime: {runtime_str}\n")
            else:
                original_stdout_console.write(f"File: {file_path.name}\n")
                original_stdout_console.write(f"Log: {log_file_path}\n")
                original_stdout_console.write(f"Runtime: {runtime_str}\n")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Batch-process MKV files with ffmpeg+SvtAv1EncApp, audio downmixing, per-file logging, and optional autocrop.")
    parser.add_argument("--no-downmix", action="store_true", help="Preserve original audio channel layout.")
    parser.add_argument("--autocrop", action="store_true", help="Automatically detect and crop black bars from video using cropdetect.")
    args = parser.parse_args()
    main(no_downmix=args.no_downmix, autocrop=args.autocrop)