#!/usr/bin/env python3
import os
import pathlib
import subprocess
import shutil
import sys
import argparse
import tempfile
import logging
import time
import psutil
from tqdm import tqdm

video_extensions: list[str] = ["mp4", "mkv", "avi", "wmv", "mov", "m4v", "ts", "flv", "mpg"]

total_files: int = 0

def is_system_busy() -> bool:
    """Returns true if the cpu is doing something else. maybe."""
    # First call all processes, so we can recheck them later.
    # See: https://psutil.readthedocs.io/en/latest/#psutil.Process.cpu_percent
    for proc in psutil.process_iter(['cpu_percent', 'pid']):
        pass
    total_cpu = 0
    our_pids = [os.getpid()]
    try:
        # add child processes
        our_pids.extend(p.pid for p in psutil.Process().children(recursive=True))
    except:
        pass

    for proc in psutil.process_iter(['cpu_percent', 'pid']):
        if proc.info['pid'] not in our_pids:
            total_cpu += proc.info['cpu_percent']

    cpu_count = psutil.cpu_count()
    if cpu_count is None:
        cpu_count = 1
    return total_cpu > 50 * cpu_count  # TODO: Parameter

def acquire_file_lock(filepath: pathlib.Path) -> bool:
    lock_file = filepath.with_suffix(filepath.suffix + '.lock')
    try:
        # Atomic file creation
        lock_fd = os.open(str(lock_file), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
        with os.fdopen(lock_fd, 'w') as f:
            f.write(f"{os.getpid()}\n{time.time()}\n")
        return True
    except FileExistsError:
        try:
            lock_age = time.time() - lock_file.stat().st_mtime
            if lock_age > 86400:  # 24 hours
                lock_file.unlink()
                return acquire_file_lock(filepath)
        except FileNotFoundError:
            return acquire_file_lock(filepath)
        return False

def release_file_lock(filepath: pathlib.Path) -> None:
    lock_file = filepath.with_suffix(filepath.suffix + '.lock')
    try:
        lock_file.unlink()
    except FileNotFoundError:
        pass

def setup_logging(verbose: bool) -> None:
    """Configure logging to both file and stdout."""
    log_level = logging.DEBUG if verbose else logging.INFO

    # Use local log directory
    xdg_state_home = os.environ.get('XDG_STATE_HOME', str(pathlib.Path.home() / '.local' / 'state'))
    log_dir = pathlib.Path(xdg_state_home) / 'vcmp'
    log_dir.mkdir(parents=True, exist_ok=True)

    log_file = log_dir / 'vcmp.log'

    # Create console handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(log_level)
    console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))

    # Create file handler
    file_handler = logging.FileHandler(log_file)
    file_handler.setLevel(logging.DEBUG)  # Always log everything to file
    file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))

    # Set up root logger
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.DEBUG)
    root_logger.addHandler(console_handler)
    root_logger.addHandler(file_handler)

def ffmpeg_command(input: pathlib.Path, output: pathlib.Path, nice: int = 20) -> list[str]:
    return [
        "nice", "-n", str(nice),
        "ffmpeg",
        "-progress", "pipe:1",  # Send progress info to stdout
        "-i", str(input),
        "-vf", "fps='if(gte(source_fps,100),source_fps/4,if(gte(source_fps,50),source_fps/2,source_fps))',scale=-2:'min(1080,ih)'",
        "-c:v", "libsvtav1",
        "-crf", "40",
        "-c:a", "aac",
        "-b:a", "64k",
        str(output)
    ]

def clean_directory(directory_path: pathlib.Path) -> None:
    """Yeet all files and directories from a directory."""
    for item in directory_path.iterdir():
        if item.is_dir():
            shutil.rmtree(item)  # Yeet the directory
        else:
            item.unlink()  # Yeet the file

def is_video(path: pathlib.Path) -> bool:
    return path.suffix.lower()[1:] in video_extensions

def get_video_duration(input_file: pathlib.Path) -> float:
    cmd = [
        "ffprobe",
        "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        str(input_file)
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return float(result.stdout)
    except (subprocess.CalledProcessError, ValueError):
        logging.error(f"Error getting duration for {input_file}")
        return 1.0

def run_progress(process, duration, file_index=0) -> bool:
    current_time = None
    last_update = 0
    progress = 0
    UPDATE_INTERVAL = 10 # TODO: Parameter
    start_time = time.time()
    is_interactive = sys.stdout.isatty()
    pbar = None
    last_time = 0
    UPDATE_INTERVAL = 1
    pbar = tqdm(total=duration, unit='s', bar_format='{l_bar}{bar}| {n:.2f}{unit}/{total:.2f}{unit} [{elapsed}<{remaining}, ' '{rate_fmt}{postfix}')
    last_time = 0
    while True:
        line = process.stdout.readline() if process.stdout else None

        # Check if we've reached the end of output (no more lines to read)
        # and if the process has finished (poll() returns exit code instead of None)
        if not line and process.poll() is not None:
            if pbar:
                pbar.close()
            break

        if line and line.startswith('out_time_ms='):
            try:
                if time.time() - last_update > UPDATE_INTERVAL:
                    if is_system_busy():
                        if pbar:
                            pbar.close()
                        return False
                    if is_interactive and pbar:
                        time_str = line.split('=')[1].strip()
                        current_time = float(time_str) / 1000000
                        current_time = min(current_time, duration)
                        pbar.update(current_time - last_time)
                        last_time = current_time
                    last_update = time.time()
            except:
                pass
    return True

def run(filepath: pathlib.Path, cmd: list[str], temp_path: pathlib.Path, file_index: int) -> None:
    duration = get_video_duration(filepath)
    xdg_state_home = os.environ.get('XDG_STATE_HOME', str(pathlib.Path.home() / '.local' / 'state'))
    ffmpeg_log_path = pathlib.Path(xdg_state_home) / 'vcmp' / 'ffmpeg.log'
    ffmpeg_log_path.parent.mkdir(parents=True, exist_ok=True)

    with open(ffmpeg_log_path, 'a') as log_file:
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=log_file,
            universal_newlines=True
        )
        if not run_progress(process, duration, file_index):
            process.kill()
            release_file_lock(filepath)
            if temp_path.exists():
                temp_path.unlink()
            sys.exit(0)
        sys.stdout.write("\n")
        sys.stdout.flush()

        # Wait for the process to complete and get the return code
        process.wait()

        if process.returncode != 0:
            logging.error(f"FFmpeg exited with code {process.returncode}")
            if temp_path.exists():
                temp_path.unlink()
            release_file_lock(filepath)
            raise subprocess.CalledProcessError(process.returncode, cmd)

def process_single_file(
        filepath: pathlib.Path,
        target_path: pathlib.Path,
        temp_dir: pathlib.Path,
        remove_source: bool,
        nice: int,
        file_index: int
) -> bool:  # Return success/failure status
    if not acquire_file_lock(filepath):
        logging.info(f"Skipping {filepath}: already being processed")
        return False

    temp_path = temp_dir / filepath.with_suffix(".mp4").name
    clean_directory(temp_dir) # Should be empty if mktemp was used, but clear it regardless.
    logging.info(f"Processing ({file_index}/{total_files}): {filepath} to {target_path} ")
    logging.debug(f"Using temporary file: {temp_path}")

    cmd = ffmpeg_command(filepath, temp_path, nice)
    logging.debug(f"Executing: {' '.join(cmd)}")
    try:
        run(filepath, cmd, temp_path, file_index)
    except KeyboardInterrupt:
        logging.warning("Process interrupted by user")
        if temp_path.exists():
            temp_path.unlink()
        release_file_lock(filepath)
        raise  # Re-raise KeyboardInterrupt to exit the program
    except subprocess.CalledProcessError as e:
        logging.error(f"FFmpeg error processing {filepath}: {e}")
        if temp_path.exists():
            temp_path.unlink()
        release_file_lock(filepath)
        return False
    else:
        # File successfully converted to av1, move to destination
        target_path.parent.mkdir(parents=True, exist_ok=True)
        success = False

        try:
            # Use cp to copy file since it handles permissions better
            subprocess.run(["cp", str(temp_path), str(target_path)], check=True)
            temp_path.unlink()
            success = True
        except subprocess.CalledProcessError as e:
            logging.error(f"Failed to copy file: {e}")
            if temp_path.exists():
                temp_path.unlink()

        if success:
            logging.info(f"Successfully compressed {filepath}")
            if remove_source:
                logging.info(f"Removing source file {filepath}")
                filepath.unlink()
            release_file_lock(filepath)
            return True

        release_file_lock(filepath)
        return False

def process_directory(
        *,
        source_dir: pathlib.Path,
        destination_dir: pathlib.Path,
        temp_dir: pathlib.Path,
        remove_source: bool = False,
        nice: int = 20,
        count_only: bool = False
) -> int:
    """Process all video files in source directory recursively."""
    count = 0
    successful = 0
    for dirpath, _dirnames, filenames in os.walk(source_dir):
        for filename in filenames:
            filepath = pathlib.Path(dirpath) / filename

            # File might have been deleted since starting script
            if not filepath.exists() or not is_video(filepath):
                logging.debug(f"Skipping {filepath}: not a video file or doesn't exist")
                continue
            if ".cmp." in filepath.name:
                logging.debug(f"Skipping {filepath}: already compressed")
                continue

            # Determine destination path, remove source if already exists and requested
            target_path = destination_dir / filepath.relative_to(source_dir)
            target_path = target_path.with_suffix(".cmp.mp4")
            if target_path.exists():
                if remove_source:
                    logging.info(f"Removing source file {filepath} as compressed version exists")
                    filepath.unlink()
                logging.debug(f"Skipping {filepath}: target already exists")
                continue

            count += 1
            if count_only:
                continue

            if process_single_file(filepath, target_path, temp_dir, remove_source, nice, count):
                successful += 1

    if not count_only:
        logging.info(f"Processed {count} files, {successful} successful")
    return count


def main():
    parser = argparse.ArgumentParser(
        description="""
        Video compression utility that converts videos to AV1 format with reduced quality/size.
        Converts to 1080p max resolution, reduces high framerates, and uses low bitrate audio.
        The compressed output files will have '.cmp' added to the filename.
        """
    )
    parser.add_argument('-s', '--source-dir', type=pathlib.Path, default=pathlib.Path.cwd(),
                      help='Source directory containing videos to compress (default: current directory)')
    parser.add_argument('-d', '--destination-dir', type=pathlib.Path,
                      help='Destination directory for compressed videos (default: same as source)')
    parser.add_argument('-t', '--temp-dir', type=pathlib.Path, default=pathlib.Path(tempfile.mkdtemp()),
                      help='Temporary directory for processing (default: system temp directory)')
    parser.add_argument('-r', '--remove-source', action='store_true',
                      help='Remove source files after successful compression')
    parser.add_argument('-v', '--verbose', action='store_true',
                      help='Enable verbose logging output')
    parser.add_argument('-n', '--nice', type=lambda x: int(x) if 0 <= int(x) <= 39 else exec('raise ValueError("Nice value must be between 0 and 39")'),
                      default=20,
                      help='Nice value for process priority (default: 20). Valid values are 0 (highest priority) to 39 (lowest priority). Lower values use more system resources.')
    parser.add_argument('-i', '--ignore-busy', action='store_true',
                      help='Ignore system busy check and run anyway')
    args = parser.parse_args()

    setup_logging(bool(args.verbose))

    source_dir = args.source_dir
    destination_dir = args.destination_dir if args.destination_dir else args.source_dir
    temp_dir = args.temp_dir

    if not args.ignore_busy and is_system_busy():
        logging.error("System is busy, try again later or use --ignore-busy flag to override")
        sys.exit(1)

    logging.info(f"Starting video compression from {source_dir} to {destination_dir}")

    global total_files
    total_files = process_directory(
        source_dir=source_dir,
        destination_dir=destination_dir,
        temp_dir=temp_dir,
        count_only=True
    )

    process_directory(
        source_dir=source_dir,
        destination_dir=destination_dir,
        temp_dir=temp_dir,
        remove_source=args.remove_source,
        nice=args.nice
    )
if __name__ == "__main__":
    main()