#!/usr/bin/env python3
"""
Smart Video Encoder - AV1 with HEVC Fallback
Runs on Windows with native ab-av1.exe and ffmpeg
Refactored to use vmaf_common.py
"""

import os
import sys
import subprocess
import json
import shutil
import time
import argparse
import signal
import platform
import threading
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

import vmaf_common as common

# --- Configuration ---
TARGET_VMAF_MIN = common.DEFAULT_CONFIG["target_vmaf"]
MIN_SAVINGS_PERCENT = common.DEFAULT_CONFIG["min_savings"]
MAX_JOBS = common.DEFAULT_CONFIG["cpu_jobs"]

AV1_QUALITY = common.DEFAULT_CONFIG["av1_crf"]
HEVC_QUALITY = common.DEFAULT_CONFIG["hevc_crf"]

SAMPLE_DURATION = 60
SAMPLE_START = 300
MAX_SAMPLE_SIZE_MB = 150

TEMP_DIR = common.get_temp_dir()

# Tools
FFMPEG_BIN = "ffmpeg"
AB_AV1_BIN = "ab-av1"

# Global state
_shutdown_requested = False
_lock_files = {}

def signal_handler(signum, frame):
    global _shutdown_requested
    _shutdown_requested = True
    print("\n\n[WARNING] Shutdown requested. Finishing current tasks...")

if platform.system() != "Windows":
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

def run_command_streaming(cmd, description=""):
    """Run command and stream output in real-time"""
    print(f"    ▶ Running: {description}")

    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True,
        bufsize=1
    )

    output_lines = []
    if process.stdout:
        for line in process.stdout:
            if _shutdown_requested:
                process.terminate()
                break
            line = line.rstrip()
            print(f"      {line}")
            output_lines.append(line)

    process.wait()
    return process.returncode, "\n".join(output_lines)


def test_av1_sample(filepath, output_path):
    """Test AV1 encoding on a sample segment"""
    print(f"  🧪 Testing AV1 with {SAMPLE_DURATION}s sample...")

    cmd = [
        "ffmpeg",
        "-hide_banner",
        "-loglevel", "warning",
        "-ss", str(SAMPLE_START),
        "-i", str(filepath),
        "-t", str(SAMPLE_DURATION),
        "-c:v", "libsvtav1",
        "-crf", str(AV1_QUALITY),
        "-preset", "6",
        "-c:a", "copy",
        "-c:s", "copy",
        "-y",
        str(output_path)
    ]

    returncode, output = run_command_streaming(cmd, f"AV1 sample test")
    return returncode == 0 and output_path.exists()


def encode_av1_full(filepath, output_path):
    """Full AV1 encode using ab-av1"""
    print(f"  🎬 Full AV1 encode...")

    cmd = [
        "ab-av1", "encode",
        "-i", str(filepath),
        "-o", str(output_path),
        "--crf", str(AV1_QUALITY),
        "--preset", "6",
        "--acodec", "copy"
    ]

    return run_command_streaming(cmd, f"AV1 full encode")[0] == 0


def encode_hevc_full(filepath, output_path):
    """Full HEVC encode using ffmpeg"""
    print(f"  🎬 Full HEVC encode...")

    cmd = [
        "ffmpeg",
        "-hide_banner",
        "-loglevel", "warning",
        "-i", str(filepath),
        "-c:v", "libx265",
        "-crf", str(HEVC_QUALITY),
        "-preset", "medium",
        "-c:a", "copy",
        "-c:s", "copy",
        "-y",
        str(output_path)
    ]

    return run_command_streaming(cmd, f"HEVC full encode")[0] == 0


def process_file(filepath, log_dir, log_category, lock_dir):
    """Process a single video file with AV1→HEVC fallback"""
    filepath = Path(filepath)
    filename = filepath.name

    # Check lock (using shared logic now)
    lock_file = common.acquire_lock(lock_dir, filepath)
    if not lock_file:
        print(f"  🔒 Skipping (locked): {filename}")
        return True

    try:
        print(f"\n{'='*60}")
        print(f"📁 Processing: {filename}")
        print(f"{'='*60}")

        # Get initial metadata
        metadata_before = common.get_video_info(filepath)
        if not metadata_before:
            print(f"❌ Could not read metadata, skipping")
            return False

        print(f"  📊 Original:")
        print(f"     Codec: {metadata_before['codec']}")
        print(f"     Size: {metadata_before['size'] / (1024**3):.2f} GB")
        print(f"     Bitrate: {metadata_before['bitrate']} kbps")
        print(f"     Duration: {metadata_before['duration'] / 60:.1f} min")

        # Skip if already AV1 or HEVC
        if metadata_before['codec'] in ['av1', 'hevc']:
            print(f"  ℹ️  Already optimized ({metadata_before['codec']}), skipping")
            # We don't log skips to main log to avoid clutter, but could if needed
            return True

        input_size = metadata_before['size']

        # --- PHASE 1: AV1 SAMPLE TEST ---
        sample_output = TEMP_DIR / f"{filepath.stem}.sample.mkv"
        use_av1 = True

        try:
            av1_test_passed = test_av1_sample(filepath, sample_output)

            if av1_test_passed:
                sample_size = sample_output.stat().st_size
                sample_size_mb = sample_size / (1024 * 1024)

                print(f"  📏 Sample size: {sample_size_mb:.1f} MB")

                # Extrapolate to full file size
                duration_ratio = metadata_before['duration'] / SAMPLE_DURATION
                estimated_full_size = sample_size * duration_ratio
                estimated_full_gb = estimated_full_size / (1024**3)
                input_gb = input_size / (1024**3)

                print(f"  📈 Estimated full size: {estimated_full_gb:.2f} GB")
                print(f"  📉 Original size: {input_gb:.2f} GB")

                if sample_size_mb > MAX_SAMPLE_SIZE_MB:
                    print(f"  ❌ AV1 REJECTED: Sample too large ({sample_size_mb:.1f} MB > {MAX_SAMPLE_SIZE_MB} MB)")
                    use_av1 = False
                elif estimated_full_size >= input_size:
                    print(f"  ❌ AV1 REJECTED: Estimated size ({estimated_full_gb:.2f} GB) >= original ({input_gb:.2f} GB)")
                    use_av1 = False
                else:
                    estimated_savings = (1 - estimated_full_size / input_size) * 100
                    print(f"  ✅ AV1 PASS: Estimated savings {estimated_savings:.1f}%")
            else:
                print(f"  ❌ AV1 sample test failed")
                use_av1 = False

            if sample_output.exists():
                sample_output.unlink()

        except Exception as e:
            print(f"  ❌ AV1 test error: {e}")
            use_av1 = False

        # --- PHASE 2: ENCODE ---
        temp_output = TEMP_DIR / f"{filepath.stem}.temp.mkv"
        final_status = "rejected"
        final_codec = None
        final_size = input_size
        final_savings = 0.0

        if use_av1:
            # Try AV1 full encode
            if encode_av1_full(filepath, temp_output):
                metadata_after = common.get_video_info(temp_output)
                if metadata_after:
                    final_size = metadata_after['size']
                    final_savings = (1 - final_size / input_size) * 100

                    if final_size < input_size:
                        final_status = "success"
                        final_codec = "av1"
                        print(f"  ✅ AV1 SUCCESS: Saved {final_savings:.1f}%")
                    else:
                        print(f"  ❌ AV1 FAILED: Final size >= original")
                        if temp_output.exists(): temp_output.unlink()

                        # Fall back to HEVC
                        print(f"  🔄 Trying HEVC fallback...")
                        if encode_hevc_full(filepath, temp_output):
                            metadata_after = common.get_video_info(temp_output)
                            if metadata_after:
                                final_size = metadata_after['size']
                                final_savings = (1 - final_size / input_size) * 100
                                if final_size < input_size:
                                    final_status = "success"
                                    final_codec = "hevc"
                                    print(f"  ✅ HEVC SUCCESS: Saved {final_savings:.1f}%")
                                else:
                                    print(f"  ❌ HEVC FAILED: Also larger than original")
                                    if temp_output.exists(): temp_output.unlink()
            else:
                print(f"  ❌ AV1 encode failed, trying HEVC...")
                if encode_hevc_full(filepath, temp_output):
                    metadata_after = common.get_video_info(temp_output)
                    if metadata_after:
                        final_size = metadata_after['size']
                        final_savings = (1 - final_size / input_size) * 100
                        if final_size < input_size:
                            final_status = "success"
                            final_codec = "hevc"
                            print(f"  ✅ HEVC SUCCESS: Saved {final_savings:.1f}%")
                        else:
                            print(f"  ❌ HEVC FAILED: Larger than original")
                            if temp_output.exists(): temp_output.unlink()
        else:
            # AV1 test failed, try HEVC directly
            print(f"  🔄 Trying HEVC directly...")
            if encode_hevc_full(filepath, temp_output):
                metadata_after = common.get_video_info(temp_output)
                if metadata_after:
                    final_size = metadata_after['size']
                    final_savings = (1 - final_size / input_size) * 100
                    if final_size < input_size:
                        final_status = "success"
                        final_codec = "hevc"
                        print(f"  ✅ HEVC SUCCESS: Saved {final_savings:.1f}%")
                    else:
                        print(f"  ❌ HEVC FAILED: Larger than original")
                        if temp_output.exists(): temp_output.unlink()

        # --- PHASE 3: FINALIZE ---
        if final_status == "success":
            # Replace original file (Safe Upload)
            if filepath.suffix:
                backup_path = filepath.with_suffix(f"{filepath.suffix}.backup")
            else:
                backup_path = Path(str(filepath) + ".backup")
                
            shutil.move(str(filepath), str(backup_path))
            shutil.move(str(temp_output), str(filepath))
            
            # Verify Integrity
            if Path(filepath).stat().st_size == final_size:
                backup_path.unlink()
                metadata_after = common.get_video_info(filepath)
                
                common.log_event(log_dir, f"{log_category}.jsonl", {
                    "file": str(filepath),
                    "status": "success",
                    "codec": final_codec,
                    "input_size": input_size,
                    "output_size": final_size,
                    "savings_percent": final_savings,
                    "metadata_before": metadata_before,
                    "metadata_after": metadata_after
                })
                
                print(f"\n  ✅ SUCCESS: Optimized with {final_codec.upper() if final_codec else 'unknown'}")
                print(f"     Savings: {final_savings:.1f}% ({input_size / (1024**3):.2f} GB → {final_size / (1024**3):.2f} GB)")
            else:
                print("  ❌ Critical Error: Copied file size mismatch! Restoring backup.")
                shutil.move(str(backup_path), str(filepath))

        else:
            common.log_event(log_dir, "rejected.jsonl", {
                "file": str(filepath),
                "status": "rejected",
                "reason": "both_codecs_larger_than_original",
                "input_size": input_size,
                "metadata": metadata_before
            })

            print(f"\n  ❌ REJECTED: Both AV1 and HEVC larger than original")
            print(f"     Keeping original file ({input_size / (1024**3):.2f} GB)")

        return True

    except Exception as e:
        print(f"❌ Error processing {filename}: {e}")
        return False

    finally:
        # Release shared lock
        if lock_file and lock_file.exists():
            lock_file.unlink()

def scan_directory(directory, extensions=None):
    """Scan directory for video files"""
    if extensions is None:
        extensions = {".mkv", ".mp4", ".mov", ".avi", ".ts"}

    files = []
    for dirpath, dirnames, filenames in os.walk(directory):
        # Skip processed/system files
        dirnames[:] = [d for d in dirnames if not d.startswith("_") and d not in [".recycle", "@eaDir"]]

        for filename in filenames:
            filepath = Path(dirpath) / filename
            if filepath.suffix.lower() in extensions:
                if "_av1" in filepath.stem.lower() or "_hevc" in filepath.stem.lower():
                    continue
                files.append(filepath)

    return sorted(files)

def main():
    parser = argparse.ArgumentParser(description="Smart Video Encoder - AV1 with HEVC Fallback")
    parser.add_argument("--tv-dir", default=common.DEFAULT_CONFIG["tv_dir"], help="TV directory")
    parser.add_argument("--content-dir", default=common.DEFAULT_CONFIG["content_dir"], help="Content directory")
    parser.add_argument("--jobs", type=int, default=MAX_JOBS, help="Parallel jobs")
    parser.add_argument("--tv-only", action="store_true", help="Process TV only")
    parser.add_argument("--content-only", action="store_true", help="Process content only")
    args = parser.parse_args()

    print("="*60)
    print("🎬 Smart Video Encoder - AV1 with HEVC Fallback")
    print("="*60)

    # Setup
    common.check_dependencies()
    lock_dir, log_dir = common.get_base_paths(args)
    TEMP_DIR.mkdir(parents=True, exist_ok=True)

    print(f"\n📁 Configuration:")
    print(f"   TV Directory: {args.tv_dir}")
    print(f"   Content Directory: {args.content_dir}")
    print(f"   Parallel Jobs: {args.jobs}")
    print(f"   Locks: {lock_dir}")
    print(f"   Logs: {log_dir}")
    print()

    tasks = []
    if not args.content_only:
        tv_dir = Path(args.tv_dir)
        if tv_dir.exists():
            tv_files = scan_directory(tv_dir)
            print(f"📺 TV Files: {len(tv_files)}")
            for f in tv_files:
                tasks.append((f, log_dir, "tv_shows", lock_dir))

    if not args.tv_only:
        content_dir = Path(args.content_dir)
        if content_dir.exists():
            content_files = scan_directory(content_dir)
            print(f"📦 Content Files: {len(content_files)}")
            for f in content_files:
                tasks.append((f, log_dir, "content", lock_dir))

    if not tasks:
        print("❌ No files to process")
        return

    print(f"\n🚀 Processing {len(tasks)} files...")
    print("="*60)

    success_count = 0
    fail_count = 0

    with ThreadPoolExecutor(max_workers=args.jobs) as executor:
        futures = {}
        for item in tasks:
            # item = (filepath, log_dir, log_category, lock_dir)
            future = executor.submit(process_file, *item)
            futures[future] = item[0]

        for future in as_completed(futures):
            if _shutdown_requested:
                break

            filepath = futures[future]
            try:
                result = future.result()
                if result:
                    success_count += 1
                else:
                    fail_count += 1
            except Exception as e:
                print(f"❌ Error processing {filepath}: {e}")
                fail_count += 1

    print("\n" + "="*60)
    print("📊 Summary:")
    print(f"   Processed: {success_count + fail_count}")
    print(f"   ✅ Success: {success_count}")
    print(f"   ❌ Failed: {fail_count}")
    print("="*60)

if __name__ == "__main__":
    main()
