import os
import sys
import shutil
import json
import hashlib
import platform
import subprocess
import time
from pathlib import Path
from datetime import datetime

# --- Defaults ---
DEFAULT_CONFIG = {
    "tv_dir": r"Z:\tv",
    "content_dir": r"Z:\content",
    "target_vmaf": 93.0,
    "min_savings": 12.0,
    "gpu_jobs": 2,
    "cpu_jobs": 1,
    "av1_crf": 34,
    "hevc_crf": 28,
    "temp_dir_windows": r"C:\Users\bnair\Videos\encodes",
    "temp_dir_linux": "~/Videos/encodes"
}

# --- Paths ---
def get_base_paths(args=None):
    r"""
    Determine root paths for locks and logs.
    Priority:
    1. Shared Network Drive (parent of tv_dir) -> Z:\.vmaf_locks
    2. Local Fallback (current dir / logs)
    """
    tv_dir = Path(args.tv_dir) if args and hasattr(args, 'tv_dir') else Path(DEFAULT_CONFIG["tv_dir"])
    
    # Logic: Locks MUST be at the common root of content to be shared.
    # We attempt to find the common parent if both are supplied, otherwise default to tv_dir parent.
    
    potential_roots = []
    if args:
        if hasattr(args, 'tv_dir') and args.tv_dir: potential_roots.append(Path(args.tv_dir))
        if hasattr(args, 'content_dir') and args.content_dir: potential_roots.append(Path(args.content_dir))
    
    if not potential_roots:
        potential_roots.append(Path(DEFAULT_CONFIG["tv_dir"]))

    # Find common path
    try:
        if len(potential_roots) > 1:
            # If they are on different drives, common_path might fail or return empty on Windows
            # In that case, we fall back to the first one (TV dir parent)
            common_root = os.path.commonpath(potential_roots)
            shared_root = Path(common_root)
            # If common root is the root drive itself (e.g. Z:\), that's fine.
            # If it's a subdir (Z:\Media), also fine.
            # But if it's too deep (e.g. just Z:\), we want to make sure we don't dump .vmaf_locks in root unless intended.
            # Typically, we want the parent of the library folders.
            # If os.path.commonpath returns Z:\Media, and dirs are Z:\Media\TV and Z:\Media\Content
            # Then locks go to Z:\Media\.vmaf_locks. Perfect.
        else:
            shared_root = potential_roots[0].parent
    except:
        # Fallback for different drives
        shared_root = potential_roots[0].parent
    
    # Defaults
    lock_dir = Path("locks").resolve()
    log_dir = Path("logs").resolve()
    
    # Network Logic
    if shared_root.exists():
        network_lock = shared_root / ".vmaf_locks"
        # We prefer local logs to avoid network I/O issues during logging, 
        # but locks MUST be shared.
        try:
            network_lock.mkdir(parents=True, exist_ok=True)
            lock_dir = network_lock
        except Exception as e:
            print(f"[Warning] Could not create network locks at {network_lock}: {e}")
            print(f"          Falling back to local locks: {lock_dir}")

    # Ensure existence
    lock_dir.mkdir(parents=True, exist_ok=True)
    log_dir.mkdir(parents=True, exist_ok=True)
    
    return lock_dir, log_dir

def get_temp_dir(args=None):
    # Check args override
    if args and hasattr(args, 'temp_dir') and args.temp_dir:
        path = Path(args.temp_dir)
    elif platform.system() == "Windows":
        path = Path(DEFAULT_CONFIG["temp_dir_windows"])
    else:
        path = Path(DEFAULT_CONFIG["temp_dir_linux"]).expanduser()
        
    path.mkdir(parents=True, exist_ok=True)
    return path

# --- Logging ---
def log_event(log_dir, filename, data):
    log_path = Path(log_dir) / filename
    data["timestamp"] = datetime.now().isoformat()
    try:
        with open(log_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(data) + "\n")
    except Exception as e:
        print(f"[ERROR] Failed to write log: {e}")

# --- Global Processed Log ---
# Prevents re-encoding files that were already processed in previous runs
_processed_cache = None

def load_processed_log(log_dir):
    """Load set of already-processed file paths from processed.jsonl"""
    global _processed_cache
    if _processed_cache is not None:
        return _processed_cache
    
    _processed_cache = set()
    log_path = Path(log_dir) / "processed.jsonl"
    
    if log_path.exists():
        try:
            with open(log_path, "r", encoding="utf-8") as f:
                for line in f:
                    try:
                        entry = json.loads(line.strip())
                        if "file" in entry:
                            _processed_cache.add(entry["file"])
                    except:
                        continue
        except Exception as e:
            print(f"[Warning] Could not load processed log: {e}")
    
    return _processed_cache

def is_already_processed(log_dir, filepath):
    """Check if file was already processed in a previous run"""
    processed = load_processed_log(log_dir)
    return str(filepath) in processed

def mark_processed(log_dir, filepath, codec, vmaf, savings):
    """Mark a file as processed (prevents re-encoding in future runs)"""
    global _processed_cache
    log_path = Path(log_dir) / "processed.jsonl"
    
    entry = {
        "file": str(filepath),
        "codec": codec,
        "vmaf": vmaf,
        "savings": savings,
        "timestamp": datetime.now().isoformat()
    }
    
    try:
        with open(log_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(entry) + "\n")
        
        # Update cache
        if _processed_cache is not None:
            _processed_cache.add(str(filepath))
    except Exception as e:
        print(f"[Warning] Could not write to processed log: {e}")

# --- Dependencies ---
def check_dependencies(required_tools=None):
    if required_tools is None:
        required_tools = ["ffmpeg", "ffprobe", "ab-av1"]
        
    missing = []
    # Determine bin path relative to this file
    bin_path = Path(__file__).parent.parent / "bin"
    
    for tool in required_tools:
        if tool == "ab-av1":
            # Check bundled first in bin/
            bundled = bin_path / "ab-av1.exe"
            if bundled.exists(): continue
            # Check PATH
            if shutil.which("ab-av1"): continue
            missing.append("ab-av1")
        else:
            if not shutil.which(tool):
                missing.append(tool)
    
    if missing:
        print(f"[!] CRITICAL: Missing tools: {', '.join(missing)}")
        print(f"    Checked bundled path: {bin_path}")
        sys.exit(1)

# --- Metadata ---
def get_video_info(filepath):
    """Get video info using ffprobe (Standardized)"""
    try:
        cmd = [
            "ffprobe", "-v", "quiet",
            "-print_format", "json",
            "-show_format", "-show_streams",
            str(filepath)
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        data = json.loads(result.stdout)
        
        video = next((s for s in data["streams"] if s["codec_type"] == "video"), None)
        if not video: return None
            
        size = int(data["format"].get("size", 0))
        duration = float(data["format"].get("duration", 0))
        
        # Bitrate calc
        bitrate = int(data["format"].get("bitrate", 0))
        if bitrate == 0 and duration > 0:
            bitrate = int((size * 8) / duration)

        return {
            "codec": video.get("codec_name"),
            "width": int(video.get("width", 0)),
            "height": int(video.get("height", 0)),
            "duration": duration,
            "size": size,
            "bitrate": bitrate,
            "fps": video.get("r_frame_rate", "0/0")
        }
    except Exception:
        return None

# --- Locks ---
def acquire_lock(lock_dir, filepath):
    """
    Simple file-based lock.
    Returns lock_path if acquired, None if failed.
    """
    # Use hash of full path (not just filename) for uniqueness
    fhash = hashlib.md5(str(filepath).encode()).hexdigest()
    lock_file = lock_dir / f"{fhash}.lock"
    
    if lock_file.exists():
        # Check staleness (24h)
        try:
            if time.time() - lock_file.stat().st_mtime > 86400:
                lock_file.unlink()
            else:
                # Lock exists and is fresh - file is being processed elsewhere
                return None
        except:
            return None
            
    try:
        # Write hostname and timestamp for debugging
        import socket
        lock_info = {
            "file": str(filepath),
            "host": socket.gethostname(),
            "timestamp": time.time()
        }
        lock_file.write_text(json.dumps(lock_info))
        return lock_file
    except Exception as e:
        print(f"[Warning] Failed to acquire lock for {filepath.name}: {e}")
        return None

# --- Hardware Detection ---
def detect_hardware_encoder(args=None):
    """
    Detects available hardware encoders via ffmpeg (Cross-Platform).
    
    Supports --av1-encoder and --hevc-encoder flags:
        hw  = prefer hardware encoder (error if unavailable)
        sw  = force software encoder
        off = disable this codec entirely
    """
    
    # Get user preferences from args
    av1_pref = getattr(args, 'av1_encoder', 'hw') if args else 'hw'
    hevc_pref = getattr(args, 'hevc_encoder', 'hw') if args else 'hw'
    cpu_only = getattr(args, 'cpu_only', False) if args else False
    
    # If cpu_only, treat both as 'sw'
    if cpu_only:
        av1_pref = 'sw'
        hevc_pref = 'sw'
    
    try:
        # Run ffmpeg -encoders
        res = subprocess.run(["ffmpeg", "-hide_banner", "-encoders"], capture_output=True, text=True)
        out = res.stdout
        
        # Build available encoder maps
        hw_av1_encoders = []
        hw_hevc_encoders = []
        hw_type = "cpu"
        
        # 1. AMD (AMF) - Windows (Preferred)
        if "av1_amf" in out: hw_av1_encoders.append(("av1_amf", "amf"))
        if "hevc_amf" in out: hw_hevc_encoders.append(("hevc_amf", "amf"))
        
        # 2. NVIDIA (NVENC) - Windows/Linux
        if "av1_nvenc" in out: hw_av1_encoders.append(("av1_nvenc", "nvenc"))
        if "hevc_nvenc" in out: hw_hevc_encoders.append(("hevc_nvenc", "nvenc"))
        
        # 3. AMD (VAAPI) - Linux
        if "av1_vaapi" in out: hw_av1_encoders.append(("av1_vaapi", "vaapi"))
        if "hevc_vaapi" in out: hw_hevc_encoders.append(("hevc_vaapi", "vaapi"))
        
        # 4. Intel (QSV) - Windows/Linux
        if "av1_qsv" in out: hw_av1_encoders.append(("av1_qsv", "qsv"))
        if "hevc_qsv" in out: hw_hevc_encoders.append(("hevc_qsv", "qsv"))
        
        # 5. Apple Silicon (VideoToolbox) - macOS
        if "av1_videotoolbox" in out: hw_av1_encoders.append(("av1_videotoolbox", "videotoolbox"))
        if "hevc_videotoolbox" in out: hw_hevc_encoders.append(("hevc_videotoolbox", "videotoolbox"))
        
        # Software encoders
        has_libsvtav1 = "libsvtav1" in out
        has_libx265 = "libx265" in out
        
        # Resolve AV1 encoder
        av1_enc = None
        if av1_pref == 'off':
            av1_enc = None
        elif av1_pref == 'sw':
            if has_libsvtav1:
                av1_enc = "libsvtav1"
            else:
                print("[Warning] libsvtav1 not available, AV1 disabled")
        elif av1_pref == 'hw':
            if hw_av1_encoders:
                av1_enc, hw_type = hw_av1_encoders[0]  # First available (AMD priority)
            elif has_libsvtav1:
                av1_enc = "libsvtav1"
                print("[Info] No AV1 HW encoder, using libsvtav1 (CPU)")
            else:
                print("[Warning] No AV1 encoder available")
        
        # Resolve HEVC encoder
        hevc_enc = None
        if hevc_pref == 'off':
            hevc_enc = None
        elif hevc_pref == 'sw':
            if has_libx265:
                hevc_enc = "libx265"
            else:
                print("[Warning] libx265 not available, HEVC disabled")
        elif hevc_pref == 'hw':
            if hw_hevc_encoders:
                hevc_enc, hevc_hw = hw_hevc_encoders[0]
                # Update hw_type if we didn't get one from AV1
                if hw_type == "cpu":
                    hw_type = hevc_hw
            elif has_libx265:
                hevc_enc = "libx265"
                print("[Info] No HEVC HW encoder, using libx265 (CPU)")
            else:
                print("[Warning] No HEVC encoder available")
        
        # Determine final hw_type label
        if av1_enc and "lib" not in av1_enc:
            pass  # hw_type already set
        elif hevc_enc and "lib" not in hevc_enc:
            pass  # hw_type already set from HEVC
        else:
            hw_type = "cpu"
        
        return av1_enc, hevc_enc, hw_type
        
    except Exception as e:
        print(f"[Warning] HW Detection failed: {e}")
        return None, None, "error"

def get_encoder_args(codec, encoder, qp):
    """Returns correct ffmpeg args for specific HW vendor"""
    if not encoder: return []
    
    # Software (CPU)
    if encoder == "libsvtav1":
        # CRF 0-63 (Lower is better)
        return ["-c:v", "libsvtav1", "-crf", str(qp), "-preset", "6", "-g", "240"]
        
    if encoder == "libx265":
        return ["-c:v", "libx265", "-crf", str(qp), "-preset", "medium"]

    # NVIDIA NVENC
    if "nvenc" in encoder:
        # p6 = better quality, spatial-aq for better perception
        return ["-c:v", encoder, "-rc", "constqp", "-qp", str(qp), "-preset", "p6", "-spatial-aq", "1"]
        
    # AMD AMF
    if "amf" in encoder:
        return ["-c:v", encoder, "-usage", "transcoding", "-rc", "cqp", "-qp_i", str(qp), "-qp_p", str(qp), "-qp_b", str(qp), "-quality", "quality"]

    # Intel QSV
    if "qsv" in encoder:
        return ["-c:v", encoder, "-global_quality", str(qp), "-look_ahead", "1"]
        
    # Apple VideoToolbox
    if "videotoolbox" in encoder:
        # Map 0-51 QP to 100-0 Quality (approx)
        q = int(100 - (qp * 2))
        return ["-c:v", encoder, "-q:v", str(q)]
        
    # Linux VAAPI
    if "vaapi" in encoder:
        # Uses -qp normally? or -global_quality? Depends on driver.
        # Often needs: -vf format=nv12,hwupload
        # Safe bet for vaapi is usually CQP via -qp or -global_quality
        return ["-c:v", encoder, "-qp", str(qp)]
        
    return []

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, key=lambda x: x.stat().st_size, reverse=True)
