import os
import subprocess
import logging
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
import psutil

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Constants
SIZE_THRESHOLD = 5 * 1024 * 1024 * 1024  # 5GB in bytes
HANDBRAKE_PATH = r"HandBrakeCLI.exe"
PRESET_FILE = "xiaowan.json"
PRESET_NAME = "1080P x264 xiaowan"

# New constant for CPU cores per process
CORES_PER_PROCESS = 4  # Adjust this value as needed

def set_cpu_affinity():
    # Get the current process
    process = psutil.Process()
    # Get the number of CPU cores
    num_cores = psutil.cpu_count()
    # Calculate which cores to use
    cores_to_use = list(range(CORES_PER_PROCESS))
    # Set the CPU affinity
    process.cpu_affinity(cores_to_use)

def compress_video(input_path, output_path):
    try:
        command = [
            HANDBRAKE_PATH,
            "--preset-import-file", PRESET_FILE,
            "-Z", PRESET_NAME,
            "-i", str(input_path),
            "-o", str(output_path)
        ]
        # Run the subprocess with CPU affinity set
        process = subprocess.Popen(command)
        # Set CPU affinity for the subprocess
        psutil.Process(process.pid).cpu_affinity(list(range(CORES_PER_PROCESS)))
        # Wait for the process to complete
        process.wait()
        if process.returncode != 0:
            raise subprocess.CalledProcessError(process.returncode, command)
        logging.info(f"Successfully compressed: {input_path}")
    except subprocess.CalledProcessError as e:
        logging.error(f"Error compressing {input_path}: {e}")

def process_folder(input_folder, output_folder):
    input_folder = Path(input_folder).resolve()
    output_folder = Path(output_folder).resolve()
    
    # Calculate the number of workers based on available cores and cores per process
    max_workers = max(1, psutil.cpu_count() // CORES_PER_PROCESS)
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = []
        for root, _, files in os.walk(input_folder):
            for file in files:
                input_path = Path(root) / file
                relative_path = input_path.relative_to(input_folder)
                output_path = output_folder / relative_path

                if input_path.suffix.lower() in ['.mp4', '.avi', '.mkv', '.mov']:
                    if input_path.stat().st_size > SIZE_THRESHOLD:
                        output_path.parent.mkdir(parents=True, exist_ok=True)
                        futures.append(executor.submit(compress_video, input_path, output_path))
                    else:
                        logging.info(f"Skipping {input_path} (size < 5GB)")
                else:
                    logging.info(f"Skipping non-video file: {input_path}")
        
        for future in as_completed(futures):
            try:
                future.result()
            except Exception as e:
                logging.error(f"An error occurred during compression: {e}")

def main():
    input_folder = input("Enter the input folder path: ")
    output_folder = input("Enter the output folder path: ")

    if not os.path.exists(input_folder):
        logging.error(f"Input folder does not exist: {input_folder}")
        return

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
        logging.info(f"Created output folder: {output_folder}")

    process_folder(input_folder, output_folder)

if __name__ == "__main__":
    try:
        # Set CPU affinity for the main process
        set_cpu_affinity()
        main()
    except Exception as e:
        logging.exception(f"An unexpected error occurred: {e}")