import os
import re
import shutil
from pymediainfo import MediaInfo # You might need to install this: pip install pymediainfo

# Define the network locations - KEEP THESE UNCHANGED as per user's original request
LOCATION1 = r"\\truenas\isolation\videos\redmilf_refresh"
LOCATION2 = r"\\truenas\isolation\videos\redmilf"
FINAL_LOCATION = r"\\truenas\isolation\videos\redmilf_final" # This is the final destination for ALL copied files

# Ensure the final destination directory exists
os.makedirs(FINAL_LOCATION, exist_ok=True)

def bytes_to_mb(bytes_value):
    """Converts bytes to megabytes, rounded to two decimal places."""
    return round(bytes_value / (1024 * 1024), 2)

def extract_did_number(filename):
    """
    Extracts the DIDXXX or DIDXXXX number from a filename.
    Returns the number as a string if found, otherwise None.
    """
    # Regex to match DID followed by 3 or 4 digits
    match = re.search(r"DID(\d{3,4})", filename, re.IGNORECASE)
    if match:
        return match.group(1) # Return the matched number string (e.g., "001" or "1234")
    return None

def get_file_info(filepath):
    """
    Gets file size and attempts to get codec/bitrate using pymediainfo.
    Returns a dictionary with 'size_bytes', 'size_mb', 'codec', 'bitrate'.
    If mediainfo fails, 'codec' and 'bitrate' will be None.
    """
    try:
        size_bytes = os.path.getsize(filepath)
        size_mb = bytes_to_mb(size_bytes)
        media_info = MediaInfo.parse(filepath)
        codec = None
        bitrate = 0 # Use 0 as a default for comparison

        for track in media_info.tracks:
            if track.track_type == 'Video':
                if track.format:
                    codec = track.format
                if track.bit_rate:
                    bitrate = track.bit_rate
                break # Assuming one video track is sufficient

        return {"size_bytes": size_bytes, "size_mb": size_mb, "codec": codec, "bitrate": bitrate}
    except Exception as e:
        print(f"Warning: Could not get media info for {filepath}: {e}")
        # Fallback if mediainfo fails
        size_bytes = os.path.getsize(filepath)
        size_mb = bytes_to_mb(size_bytes)
        return {"size_bytes": size_bytes, "size_mb": size_mb, "codec": None, "bitrate": 0}

def main():
    files_in_loc1 = {}
    files_in_loc2 = {}

    print("Welcome to the DID Video Comparison and Management Script!")
    print("\nDo you want to enable 'Auto-Run' mode?")
    print("In 'Auto-Run' mode, the script will automatically:")
    print("  1. Keep the file with the larger size.")
    print(f"  2. Use the filename from '{os.path.basename(LOCATION2)}' for the final copy.")
    print("If you decline, you will be prompted for each file manually.")

    auto_run_choice = input("Enter 'yes' to enable Auto-Run, or 'no' to run manually (yes/no): ").strip().lower()
    auto_run = (auto_run_choice == 'yes')

    if auto_run:
        print("\nAuto-Run mode enabled. Processing files automatically.")
    else:
        print("\nManual mode selected. You will be prompted for each file.")

    # Objective 1: Identify files with the same naming convention
    # Scan Location 1
    for root, _, files in os.walk(LOCATION1):
        for file in files:
            did_num = extract_did_number(file) # Changed function call
            if did_num:
                if did_num not in files_in_loc1:
                    files_in_loc1[did_num] = []
                files_in_loc1[did_num].append(os.path.join(root, file))

    # Scan Location 2
    for root, _, files in os.walk(LOCATION2):
        for file in files:
            did_num = extract_did_number(file) # Changed function call
            if did_num:
                if did_num not in files_in_loc2:
                    files_in_loc2[did_num] = []
                files_in_loc2[did_num].append(os.path.join(root, file))

    # Process matched files
    matched_did_numbers = set(files_in_loc1.keys()).intersection(set(files_in_loc2.keys()))

    if not matched_did_numbers:
        print(f"No matching DID files found in both '{LOCATION1}' and '{LOCATION2}'.")
        return

    for did_num in sorted(list(matched_did_numbers)):
        # Assuming only one file per DIDXXX/XXXX number per location for simplicity.
        path1 = files_in_loc1[did_num][0]
        path2 = files_in_loc2[did_num][0]

        file1_info = get_file_info(path1)
        file2_info = get_file_info(path2)

        print(f"\n--- Processing DID{did_num} ---")
        print(f"  File 1 (from '{os.path.basename(os.path.dirname(path1))}'): {os.path.basename(path1)} (Size: {file1_info['size_mb']} MB, Codec: {file1_info['codec']}, Bitrate: {file1_info['bitrate']} bps)")
        print(f"  File 2 (from '{os.path.basename(os.path.dirname(path2))}'): {os.path.basename(path2)} (Size: {file2_info['size_mb']} MB, Codec: {file2_info['codec']}, Bitrate: {file2_info['bitrate']} bps)")


        file_to_keep_path = None
        final_filename = None

        if auto_run:
            # Auto-run logic: keep larger size, use filename from Location 2
            if file1_info['size_bytes'] >= file2_info['size_bytes']:
                file_to_keep_path = path1
            else:
                file_to_keep_path = path2
            final_filename = os.path.basename(path2) # Always use filename from Location 2
            print(f"  [Auto] Keeping: {os.path.basename(file_to_keep_path)}")
            print(f"  [Auto] Using filename: {final_filename}")

        else: # Manual mode
            # Original objective 2: Compare size, quality, and codec and list the better file
            recommended_better_path = None
            # Prioritize bitrate for recommendation
            if file1_info['bitrate'] > file2_info['bitrate']:
                recommended_better_path = path1
            elif file2_info['bitrate'] > file1_info['bitrate']:
                recommended_better_path = path2
            else:
                # If bitrates are similar or unavailable, use size as a proxy for recommendation
                if file1_info['size_bytes'] >= file2_info['size_bytes']:
                    recommended_better_path = path1
                else:
                    recommended_better_path = path2
            print(f"  Recommended to keep: {os.path.basename(recommended_better_path)}")


            # Original objective 3: Prompt the user which one they would like to keep
            while True:
                choice = input("Which file would you like to keep? (1/2): ").strip()
                if choice == '1':
                    file_to_keep_path = path1
                    break
                elif choice == '2':
                    file_to_keep_path = path2
                    break
                else:
                    print("Invalid choice. Please enter '1' or '2'.")

            # Original objective 4: Prompt the user for each file, which one's file name will be used for the final copy
            while True:
                name_choice = input(f"Which file's name would you like to use for the final copy of DID{did_num}? (1/2): ").strip() # Changed prompt
                if name_choice == '1':
                    final_filename = os.path.basename(path1)
                    break
                elif name_choice == '2':
                    final_filename = os.path.basename(path2)
                    break
                else:
                    print("Invalid choice. Please enter '1' or '2'.")

        # Objective 5: Copy the selected file to the new location with the chosen name
        destination_path = os.path.join(FINAL_LOCATION, final_filename)
        try:
            print(f"Copying '{os.path.basename(file_to_keep_path)}' to '{destination_path}'...")
            shutil.copy2(file_to_keep_path, destination_path) # copy2 preserves metadata
            print("Copy successful!")
        except Exception as e:
            print(f"Error copying file {file_to_keep_path} to {destination_path}: {e}")

    print("\nComparison and copying process complete.")

if __name__ == "__main__":
    main()