import os, shutil, config
import ffmpeg
from tqdm import tqdm

def is_av1(filepath):
    try:
        probe = ffmpeg.probe(filepath)
        for stream in probe['streams']:
            if stream['codec_type'] == 'video' and 'codec_name' in stream:
                if stream['codec_name'] == 'av1':
                    return True
    except ffmpeg.Error as e:
        print(f"Error probing {filepath}: {e}")
    return False

def get_video_info(filepath):
    try:
        probe = ffmpeg.probe(filepath)
        format_info = probe['format']
        video_stream = next(
            (stream for stream in probe['streams'] if stream['codec_type'] == 'video'),
            None
        )
        if video_stream:
            # Convert from bits/sec to kbps
            bitrate_kbps = int(format_info['bit_rate']) // 1000  
            width = video_stream['width']
            height = video_stream['height']
            return bitrate_kbps, (width, height)
    except ffmpeg.Error as e:
        print(f"Error getting video info for {filepath}: {e}")
    return None, (None, None)

def get_files(folder):
    all_files = []
    for root, _, filenames in os.walk(folder):
        for filename in filenames:
            if filename.lower().endswith(('.mp4', '.mkv', '.avi', '.mov')):
                if not "encoded" in root:
                    all_files.append(os.path.join(root, filename))
    return sorted(all_files, key=os.path.getsize)

def get_target_bitrate(width, height):
    resolutions = {
        (854, 480): 1000,
        (1280, 720): 1500,
        (1920, 1080): 3000,
        (2560, 1440): 5000,
        (3840, 2160): 12000
    }
    
    for res, bitrate in resolutions.items():
        if width <= res[0] and height <= res[1]:
            return bitrate
    
    return 2500

def get_fps(filepath):
    try:
        probe = ffmpeg.probe(filepath)
        video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
        if video_stream and 'r_frame_rate' in video_stream:
            fps_str = video_stream['r_frame_rate']
            num, den = map(int, fps_str.split('/'))
            fps = num / den
            return fps
    except ffmpeg.Error as e:
        print(f"Error getting FPS for {filepath}: {e}")
    return None

def encode_video(filepath, output_path, target_bitrate):
    try:
        fps = get_fps(filepath)
        if fps is None:
            print(f"Could not determine FPS for {filepath}.")
            return False
        
        keyframe_interval = int(fps)  # Set the keyframe interval to match 1 second

        # Calculate 1.5x for max bitrate
        max_bitrate = int(1.5 * target_bitrate)

        print(f"  Encoding {filepath} to AV1 at {target_bitrate} kbps...")
        (
            ffmpeg
            .input(filepath)
            .output(
                output_path,
                vcodec='av1_nvenc',
                format='mp4',
                b=f"{target_bitrate}k",
                maxrate=f"{max_bitrate}k",
                bufsize=f"{max_bitrate}k",
                preset='p5',
                g=keyframe_interval
            )
            .run(
                overwrite_output=True,
            )
        )
        print(f"  Finished encoding {os.path.basename(filepath)} to AV1 at {target_bitrate} kbps "
              f"(maxrate={max_bitrate} kbps).")
    except ffmpeg.Error as e:
        print(f"  Error encoding {os.path.basename(filepath)} to AV1: {e}")

def check_and_replace_if_smaller(original_path, temp_output_path):
    if not os.path.exists(temp_output_path):
        print(f"[ERROR] Temp file {temp_output_path} not found. Skipping replacement...")
        return

    original_size = os.path.getsize(original_path)
    processed_size = os.path.getsize(temp_output_path)

    size_original_mb = original_size / (1024 * 1024)
    size_processed_mb = processed_size / (1024 * 1024)
    size_diff_perc = (1 - processed_size / original_size) * 100
    size_diff_mb = size_original_mb - size_processed_mb

    if size_original_mb < 25:
        shutil.move(temp_output_path, original_path)
        return True

    if processed_size >= original_size or size_diff_mb < 1:
        os.remove(temp_output_path)
        return False
    else:
        print(100*"=")
        print(f"  Re-encoded is smaller by {size_diff_perc:.2f}% ({size_diff_mb:.2f} MB). Replacing original.")
        print(f"  Original: {size_original_mb:.2f} MB \n  Re-encoded: {size_processed_mb:.2f} MB.")
        print(100*"=")
        shutil.move(temp_output_path, original_path)
        return True

def update_codec_db(video_id, codec):
    conn, cursor = config.get_local_db_connection()
    cursor.execute("UPDATE videos SET codec = %s WHERE id = %s", (codec, video_id))
    conn.commit()
    conn.close()

def reencode_videos_av1():
    # get videos
    conn, cursor = config.get_local_db_connection()
    cursor.execute("SELECT * FROM videos WHERE codec IS NULL AND status != 'missing' AND filepath IS NOT NULL ORDER BY size ASC;")
    videos = cursor.fetchall()

    for video in tqdm(videos, desc="Processing videos", unit="file"):
        input_path = video['filepath']

        if not os.path.exists(input_path):
            print(f"🚫 File not found: {input_path}")
            continue
        
        file_size_in_mb = os.path.getsize(input_path) / (1024 * 1024)
        print(f"\nProcessing {os.path.basename(input_path)} ({file_size_in_mb:.2f} MB)...")
        
        # 2) Get current bitrate & resolution
        current_bitrate, (width, height) = get_video_info(input_path)
        if not current_bitrate:
            print("Video's bitrate is not available. Skipping")
            continue

        target_bitrate = get_target_bitrate(width, height)
        
        # If current bitrate <= target, it's not worth it to re-encode
        if current_bitrate <= target_bitrate:
            target_bitrate = current_bitrate
        
        if is_av1(input_path):
            print("Video is already encoded in AV1. Skipping")
            update_codec_db(video['id'], 'av1')
            continue
        
        # 3) Re-encode
        output_path = os.path.join('.temp', os.path.basename(input_path))
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        encoded = encode_video(input_path, output_path, target_bitrate)

        if not encoded:
            print("Encoding failed. Skipping.")
            continue

        # 4) Compare file sizes and replace if smaller
        check_and_replace_if_smaller(input_path, output_path)
        update_codec_db(video['id'], 'av1')

if __name__ == "__main__":
    reencode_videos_av1()