import subprocess
import json
import os
from collections import defaultdict
from pathlib import Path

# Common video file extensions
VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.m4v', '.mpg', '.mpeg'}

def get_video_info(file_path):
    """
    Extract video bitrate and resolution using ffprobe with modified bitrate calculation
    """
    try:
        # First command to get stream information
        result = subprocess.run([
            'ffprobe', 
            '-v', 'quiet', 
            '-print_format', 'json', 
            '-show_streams', 
            '-select_streams', 'v:0', 
            file_path
        ], capture_output=True, text=True)
        
        # Second command specifically for bitrate using -count_packets
        bitrate_result = subprocess.run([
            'ffprobe',
            '-v', 'quiet',
            '-select_streams', 'v:0',
            '-show_entries', 'format=duration,size',
            '-print_format', 'json',
            file_path
        ], capture_output=True, text=True)
        
        probe_data = json.loads(result.stdout)
        bitrate_data = json.loads(bitrate_result.stdout)
        
        video_stream = probe_data['streams'][0]
        width = int(video_stream.get('width', 0))
        height = int(video_stream.get('height', 0))
        
        # Calculate bitrate from file size and duration
        try:
            duration = float(bitrate_data['format']['duration'])
            size_bits = float(bitrate_data['format']['size']) * 8
            bitrate = (size_bits / duration) / 1000  # Convert to kbps
        except (KeyError, ZeroDivisionError):
            # Fallback method using direct bit_rate if available
            bitrate = float(video_stream.get('bit_rate', 0)) / 1000
            if bitrate == 0:
                # Second fallback: try format bitrate
                format_result = subprocess.run([
                    'ffprobe',
                    '-v', 'quiet',
                    '-print_format', 'json',
                    '-show_format',
                    file_path
                ], capture_output=True, text=True)
                format_data = json.loads(format_result.stdout)
                bitrate = float(format_data['format'].get('bit_rate', 0)) / 1000
        
        return {
            'bitrate': bitrate,
            'width': width,
            'height': height,
            'path': file_path,
            'filename': os.path.basename(file_path),
            'size_mb': os.path.getsize(file_path) / (1024 * 1024)  # Convert to MB
        }
    
    except Exception as e:
        print(f"\nError analyzing video {file_path}: {e}")
        return None

def categorize_bitrate(video_info):
    """
    Categorize bitrate based on resolution and bitrate
    """
    width = video_info['width']
    height = video_info['height']
    bitrate = video_info['bitrate']
    pixels = width * height
    
    # Resolution categories and their bitrate thresholds
    categories = {
        '4K': {
            'pixels': (3840 * 2160, float('inf')),
            'thresholds': {
                'Low': (0, 10000),
                'Medium': (10000, 20000),
                'High': (20000, float('inf'))
            }
        },
        '1080p': {
            'pixels': (1920 * 1080, 3840 * 2160),
            'thresholds': {
                'Low': (0, 5000),
                'Medium': (5000, 10000),
                'High': (10000, float('inf'))
            }
        },
        '720p': {
            'pixels': (1280 * 720, 1920 * 1080),
            'thresholds': {
                'Low': (0, 2500),
                'Medium': (2500, 5000),
                'High': (5000, float('inf'))
            }
        },
        'SD': {
            'pixels': (0, 1280 * 720),
            'thresholds': {
                'Low': (0, 1000),
                'Medium': (1000, 2500),
                'High': (2500, float('inf'))
            }
        }
    }
    
    # Find resolution category
    resolution_category = None
    for res_name, res_info in categories.items():
        min_pixels, max_pixels = res_info['pixels']
        if min_pixels < pixels <= max_pixels:
            resolution_category = res_name
            break
    
    if not resolution_category:
        return None
    
    # Find bitrate category
    for bitrate_category, (min_rate, max_rate) in categories[resolution_category]['thresholds'].items():
        if min_rate <= bitrate < max_rate:
            return f"{resolution_category} - {bitrate_category}"
    
    return None

def find_video_files(folder_path):
    """
    Recursively find all video files in the given folder
    """
    video_files = []
    for path in Path(folder_path).rglob('*'):
        if path.suffix.lower() in VIDEO_EXTENSIONS:
            video_files.append(str(path))
    return video_files

def analyze_folder(folder_path):
    """
    Analyze all videos in a folder and return categorized results
    """
    print(f"Scanning folder: {folder_path}")
    
    # Find all video files
    video_files = find_video_files(folder_path)
    total_files = len(video_files)
    print(f"Found {total_files} video files")
    
    # Analyze each video
    categorized_videos = defaultdict(list)
    for i, file_path in enumerate(video_files, 1):
        print(f"\rAnalyzing video {i}/{total_files}: {os.path.basename(file_path)}", end='')
        
        video_info = get_video_info(file_path)
        if video_info:
            category = categorize_bitrate(video_info)
            if category:
                categorized_videos[category].append(video_info)
    
    print("\nAnalysis complete!")
    return categorized_videos

def display_results(categorized_videos):
    """
    Display categorized videos sorted by bitrate
    """
    print("\nVideo Analysis Results:")
    print("=" * 100)
    
    # Sort categories in order: 4K, 1080p, 720p, SD, each with High, Medium, Low
    category_order = []
    for res in ['4K', '1080p', '720p', 'SD']:
        for quality in ['High', 'Medium', 'Low']:
            category_order.append(f"{res} - {quality}")
    
    # Display results for each category
    for category in category_order:
        if category in categorized_videos:
            videos = categorized_videos[category]
            # Sort videos by bitrate (highest to lowest)
            videos.sort(key=lambda x: x['bitrate'], reverse=True)
            
            print(f"\n{category} Quality Videos:")
            print("-" * 100)
            print(f"{'Filename':<50} {'Resolution':<15} {'Bitrate':>10} {'Size':>10}")
            print("-" * 100)
            
            for video in videos:
                filename = video['filename']
                if len(filename) > 47:
                    filename = filename[:44] + "..."
                print(f"{filename:<50} {video['width']}x{video['height']:<15} {video['bitrate']:>8.0f}k {video['size_mb']:>8.1f}MB")

def main():
    # Prompt for folder path
    folder_path = input("Enter the folder path to analyze: ").strip('"')
    
    # Validate folder exists
    if not os.path.isdir(folder_path):
        print("Folder does not exist. Please check the path and try again.")
        return
    
    # Analyze videos
    categorized_videos = analyze_folder(folder_path)
    
    # Display results
    display_results(categorized_videos)

if __name__ == "__main__":
    main()