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

# Define the network locations
LOCATION1 = r"\\truenas\isolation\videos\redmilf_refresh"
LOCATION2 = r"\\truenas\isolation\videos\redmilf"
FINAL_LOCATION = r"\\truenas\isolation\videos\redmilf_final"

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

def extract_milf_number(filename):
    """
    Extracts the MILFXXX number from a filename.
    Returns the number as a string if found, otherwise None.
    """
    match = re.search(r"MILF(\d{3})", filename, re.IGNORECASE)
    if match:
        return match.group(1)
    return None

def get_file_info(filepath):
    """
    Gets file size and attempts to get codec/bitrate using pymediainfo.
    Returns a dictionary with 'size', 'codec', 'bitrate'.
    If mediainfo fails, 'codec' and 'bitrate' will be None.
    """
    try:
        size = os.path.getsize(filepath)
        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": size, "codec": codec, "bitrate": bitrate}
    except Exception as e:
        print(f"Warning: Could not get media info for {filepath}: {e}")
        return {"size": os.path.getsize(filepath), "codec": None, "bitrate": 0}

def determine_better_file(file1_info, file2_info, path1, path2):
    """
    Compares two file infos to determine which one is 'better'.
    Prioritizes bitrate (quality), then size if bitrate is similar or unavailable.
    """
    print(f"\nComparing:")
    print(f"  File 1: {os.path.basename(path1)} (Size: {file1_info['size']} bytes, Codec: {file1_info['codec']}, Bitrate: {file1_info['bitrate']} bps)")
    print(f"  File 2: {os.path.basename(path2)} (Size: {file2_info['size']} bytes, Codec: {file2_info['codec']}, Bitrate: {file2_info['bitrate']} bps)")

    # Prioritize bitrate for quality
    if file1_info['bitrate'] > file2_info['bitrate']:
        return path1
    elif file2_info['bitrate'] > file1_info['bitrate']:
        return path2
    else:
        # If bitrates are similar or unavailable, use size as a proxy
        if file1_info['size'] >= file2_info['size']: # Greater or equal size usually means better quality
            return path1
        else:
            return path2

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

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

    # Scan Location 2
    for root, _, files in os.walk(LOCATION2):
        for file in files:
            milf_num = extract_milf_number(file)
            if milf_num:
                if milf_num not in files_in_loc2:
                    files_in_loc2[milf_num] = []
                files_in_loc2[milf_num].append(os.path.join(root, file))

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

    if not matched_milf_numbers:
        print("No matching MILF files found in both locations.")
        return

    for milf_num in sorted(list(matched_milf_numbers)):
        # Assuming only one file per MILFXXX number per location for simplicity.
        # If multiple files with the same MILFXXX number exist in one location,
        # this script will compare the first one it finds.
        # For a more robust solution, you might need to handle this by comparing
        # all permutations or selecting the "best" within each location first.
        path1 = files_in_loc1[milf_num][0]
        path2 = files_in_loc2[milf_num][0]

        # Objective 2: Compare size, quality, and codec and list the better file
        file1_info = get_file_info(path1)
        file2_info = get_file_info(path2)

        recommended_better_path = determine_better_file(file1_info, file2_info, path1, path2)

        # Objective 3: Prompt the user which one they would like to keep
        print(f"\nFor MILF{milf_num}:")
        print(f"  Option 1 (from '{os.path.basename(os.path.dirname(path1))}'): {os.path.basename(path1)}")
        print(f"  Option 2 (from '{os.path.basename(os.path.dirname(path2))}'): {os.path.basename(path2)}")
        print(f"  Recommended: {os.path.basename(recommended_better_path)}")

        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'.")

        # 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 MILF{milf_num}? (1/2): ").strip()
            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()