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

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

total_files: int = 0

ignore_busy = False

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 rotating file handler
    file_handler = RotatingFileHandler(
        log_file,
        backupCount=100
    )
    file_handler.setLevel(logging.DEBUG)  # Always log everything to file
    file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
    file_handler.doRollover()

    # 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", "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 not ignore_busy and 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, 'w') 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)

    global ignore_busy
    ignore_busy = bool(args.ignore_busy)

    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()