#!/usr/bin/env python3
"""
Usage:
  python3 count_files (<number>) (<directory>)
  python3 count_files (<number>) (<directory>) --undo
"""

import os
from pathlib import Path
import sys
import re
from termcolor import colored


VIDEO_EXTENSIONS = {
    ".mkv", ".mp4", ".avi", ".mov", ".wmv", ".flv", ".webm", ".m4v"
}


def main():
    args = sys.argv[1:]

    undo = "--undo" in args
    if undo:
        args.remove("--undo")

    if len(args) == 0:
        number = 0
        directory = Path(os.getcwd())
    elif len(args) == 1:
        try: 
            number = int(args[0])
            directory = Path(os.getcwd())
        except:
            number = 0
            directory = Path(args[0])
    elif len(args) == 2:
        try:
            number = int(args[0])
            directory = Path(args[1])
        except:
            number = int(args[1])
            directory = Path(args[0])
    else:
        print("Usage: recount_videos.py (<directory>) [--undo]")
        sys.exit(1)

    if not directory.exists():
        print(colored("Directory does not exist", "red"))
        sys.exit(1)

    import subprocess
    subprocess.run(["python", "/home/honney/.bin/tracker.py", "add", "count_files"])

    files = sorted(
        [f for f in directory.iterdir()
         if f.is_file() and f.suffix.lower() in VIDEO_EXTENSIONS],
        key=lambda f: f.name
    )

    if not files:
        print(colored("No video files found", "yellow"))
        return

    if undo:
        # detect prefix like "00", "001", etc.
        pattern = re.compile(r"^\d+")

        for f in files:
            if not pattern.match(f.name):
                continue

            new_name = Path(directory / pattern.sub("", f.name, count=1))

            if new_name.exists():
                print(colored(f"Conflict exists: {new_name.name}", "red"))
                continue

            print(
                colored("Restored: ", "green") +
                colored(f.name, "red") +
                colored(" -> ", "white") +
                colored(new_name.name, "green")
            )

            f.rename(new_name)

    else:
        max_len = max(len(str(len(files))), 2)

        for i, f in enumerate(files):
            if number:
                new_name = Path(f"{f.parent}/{i+number:0{max_len}d}{f.name}")
            else:
                new_name = Path(f"{f.parent}/{i+1:0{max_len}d}{f.name}")

            if new_name.exists():
                print(colored(f"Exists, skipping: {new_name.name}", "red"))
                continue

            print(
                colored("Renamed: ", "green") +
                colored(f.name, "red") +
                colored(" -> ", "white") +
                colored(new_name.name, "green")
            )

            f.rename(new_name)


if __name__ == "__main__":
    main()