import os
import shutil
import subprocess
from pathlib import Path
import argparse

def get_audio_info(filepath):
    """Return (stream_index, channels, bitrate_in_bits)."""
    cmd = [
        "ffprobe", "-v", "error", "-select_streams", "a:0",
        "-show_entries", "stream=index,channels,bit_rate",
        "-of", "csv=p=0", filepath
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if not result.stdout:
        return None, 2, 128000
    parts = result.stdout.strip().split(',')
    if len(parts) < 3:
        return None, 2, 128000
    _, channels, bitrate = parts
    channels = int(channels)
    bitrate = int(bitrate) // 1000  # kbps approx

    # Bucket logic
    if channels == 2:
        if bitrate < 80:
            br = 64000
        elif bitrate < 112:
            br = 96000
        else:
            br = 128000
    else:
        if bitrate < 176:
            br = 160000
        else:
            br = 192000
    return 0, channels, br

def get_language_tag(filepath):
    cmd = [
        "ffprobe", "-v", "error", "-select_streams", "a:0",
        "-show_entries", "stream_tags=language",
        "-of", "default=nokey=1:noprint_wrappers=1", filepath
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    return result.stdout.strip() or None


def main():
    parser = argparse.ArgumentParser(
        description="Batch AV1 transcoder (FFmpeg-based) for anime/TV folders."
    )
    parser.add_argument("origination", help="Path to origination folder (e.g., P:\\Anime\\Show)")
    parser.add_argument("-p", "--processing", default=r"C:\Users\Tyler\Videos\Video Conversion\temp",
                        help="Temporary processing folder (default: %(default)s)")
    parser.add_argument("-c", "--completed", default=None,
                        help="Completed output folder (default: back to original)")
    parser.add_argument("-r", "--resolution", type=int, choices=[720, 1080], default=1080,
                        help="Output resolution height (default: %(default)s)")
    parser.add_argument("-cq", type=int, default=None,
                        help="Constant quality value for AV1_NVENC (auto if not specified)")

    args = parser.parse_args()

    orig = args.origination
    processing = args.processing
    completed = args.completed
    res_choice = str(args.resolution)

    # Auto-detect filter & default CQ
    filter_flags = "lanczos"
    if "\\tv\\" in orig.lower() or "/tv/" in orig.lower():
        filter_flags = "bicubic"
        cq_default = 28 if res_choice == "1080" else 32
    else:
        cq_default = 32 if res_choice == "1080" else 34

    cq = args.cq if args.cq is not None else cq_default

    print("\n=== Using These Settings ===")
    print(f"Origination: {orig}")
    print(f"Processing:  {processing}")
    print(f"Completed:   {completed if completed else '[Return to original folder]'}")
    print(f"Resolution:  {res_choice}")
    print(f"Filter:      {filter_flags}")
    print(f"CQ:          {cq}")
    print("=============================\n")

    os.makedirs(processing, exist_ok=True)
    if completed:
        os.makedirs(completed, exist_ok=True)

    suffix = " -EHX"

    for root, dirs, files in os.walk(orig):
        for f in files:
            if not f.lower().endswith((".mkv", ".mp4")):
                continue
            if "ehx" in f.lower() or "megusta" in f.lower():
                print(f"Skipping {f} (contains 'EHX' or 'MeGusta')")
                continue

            print("=" * 60)
            print(f"Processing: {f}")
            src = Path(root) / f
            tmp = Path(processing) / f
            shutil.copy2(src, tmp)

            # Detect audio info
            _, channels, abr = get_audio_info(str(tmp))
            lang = get_language_tag(str(tmp))
            lang_metadata = []
            if not lang:
                lang_metadata = ["-metadata:s:a:0", "language=eng"]

            width, height = ("1920", "1080") if res_choice == "1080" else ("1280", "720")
            out_file = Path(processing) / f"{Path(f).stem}{suffix}.mkv"

            ffmpeg_cmd = [
                "ffmpeg", "-y", "-i", str(tmp),
                "-vf", f"scale={width}:{height}:flags={filter_flags}:force_original_aspect_ratio=decrease",
                "-map", "0:v", "-map", "0:a", "-map", "0:s?",
                "-c:v", "av1_nvenc", "-preset", "p1", "-cq", str(cq), "-pix_fmt", "p010le",
                "-c:a", "aac", "-b:a", str(abr), "-ac", str(channels),
                *lang_metadata,
                "-metadata:s:a:0", f"bit_rate={abr}",
                "-c:s", "copy",
                str(out_file)
            ]

            subprocess.run(ffmpeg_cmd)

            target = Path(completed) / out_file.name if completed else Path(root) / out_file.name
            shutil.move(out_file, target)
            print(f"Moved file to {target}")

            if target.exists():
                print("Conversion confirmed. Deleting originals...")
                os.remove(src)
                os.remove(tmp)
            else:
                print("ERROR: Converted file not found. Skipping deletion of originals.")

if __name__ == "__main__":
    main()