import os
import re
import subprocess
import json
import requests
from pathlib import Path
from urllib.parse import quote
from rich import print

# === CONFIG ===
MOVIES_DIR = Path(r"PATH:/TO/VIDEOS") << Replace this with your folder path
VIDEO_EXTS = {".mp4", ".mkv", ".avi", ".mov", ".webm"}
OMDB_API_KEY = "API_KEY"  # << Replace this with your OMDb API key
YEAR_REGEX = re.compile(r"\((\d{4})\)")

# === FUNCTIONS ===

def get_video_duration(file_path):
    cmd = [
        "ffprobe", "-v", "error", "-show_entries",
        "format=duration", "-of", "json", str(file_path)
    ]
    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    try:
        duration_sec = float(json.loads(result.stdout)["format"]["duration"])
        return round(duration_sec / 60, 1)
    except:
        return None

def clean_title(file_stem):
    title = re.sub(r"[.\-_]", " ", file_stem)
    title = re.sub(r"\{.*?\}", "", title)
    title = re.sub(r"\[.*?\]", "", title)
    title = re.sub(r"\b(1080p|720p|4k|bluray|webrip|hdrip|x264|x265|h264|h265|aac|dvdrip|hevc)\b", "", title, flags=re.I)
    title = re.sub(r"\s{2,}", " ", title).strip()
    return title

def search_omdb(title, year=None):
    clean = re.sub(r"\(\d{4}\)", "", title)
    clean = re.sub(r"\bHDTV|WEBRip|BluRay|1080p|720p|x264|x265|HEVC|AAC\b", "", clean, flags=re.I)
    clean = re.sub(r"[\[\]\{\}\.\-_]", " ", clean)
    clean = re.sub(r"\s{2,}", " ", clean).strip()

    params = {
        "apikey": OMDB_API_KEY,
        "t": clean,
        "type": "movie",
    }
    if year:
        params["y"] = year

    resp = requests.get("http://www.omdbapi.com/", params=params)
    data = resp.json()

    if data.get("Response") == "True":
        runtime = data.get("Runtime", "")
        if "min" in runtime:
            try:
                return int(runtime.replace("min", "").strip())
            except:
                return None
    else:
        print(f"[dim]OMDb Error: {data.get('Error')}[/dim]")

    return None

def scan_and_compare():
    flagged = []

    for root, dirs, files in os.walk(MOVIES_DIR):
        for file in files:
            if Path(file).suffix.lower() in VIDEO_EXTS:
                if "{edition-" in file:
                    continue  # Skip already-tagged files

                full_path = Path(root) / file
                local_duration = get_video_duration(full_path)
                if not local_duration:
                    continue

                file_stem = Path(file).stem
                folder_name = Path(root).name

                title_raw = folder_name or file_stem
                title = clean_title(title_raw)

                match = YEAR_REGEX.search(folder_name) or YEAR_REGEX.search(file_stem)
                year = match.group(1) if match else None

                print(f"\n[cyan]Checking:[/cyan] {file_stem}")
                print(f" → Title Guess: [green]{title}[/green], Year: {year}")
                print(f" → Local runtime: {local_duration} min")

                imdb_runtime = search_omdb(title, year)

                if imdb_runtime:
                    diff = round(local_duration - imdb_runtime, 1)
                    if diff >= 5:
                        print(f"[bold yellow] → Potential edition tag:[/bold yellow] {{edition-Extended Cut}} (IMDb: {imdb_runtime} min, diff: +{diff} min)")
                        flagged.append((full_path, title))
                    elif diff <= -5:
                        print(f"[bold] → Note:[/bold] Local runtime is [red]{abs(diff)} min shorter[/red] than IMDb ({imdb_runtime} min)")
                        flagged.append((full_path, title))
                    else:
                        print(f" → Matches IMDb runtime ({imdb_runtime} min)")
                else:
                    print("[red] × IMDb runtime not found[/red]")
                    flagged.append((full_path, title))

    return flagged

def interactive_mode(flagged_movies):
    if not flagged_movies:
        print("\n[green]All movies match expected runtimes. No action needed.[/green]")
        return

    print("\n[bold underline]Flagged Movies:[/bold underline]")
    for idx, (path, title) in enumerate(flagged_movies, 1):
        print(f"[{idx}] {title}")

    while True:
        try:
            selection = int(input("\nEnter number of movie to investigate or 0 to cancel: "))
            if selection == 0:
                break

            selected_path, selected_title = flagged_movies[selection - 1]
            query = f"{selected_title} extended cut runtime"
            search_url = f"https://www.google.com/search?q={quote(query)}"
            print(f"{search_url}")

            tag = input("Enter edition tag (e.g., Director's Cut, Extended, Final): ").strip()
            if not tag:
                print("No tag entered. Skipping.")
                continue

            tag_str = f" {{edition-{tag}}}"
            folder = selected_path.parent
            stem = selected_path.stem
            base_name = re.sub(r"\{edition-.*?\}", "", stem).strip()
            ext = selected_path.suffix

            new_name = f"{base_name}{tag_str}{ext}"
            new_path = folder / new_name

            print(f"\n[cyan]Renaming all files in folder starting with:[/cyan] {base_name}")
            for f in folder.iterdir():
                if f.is_file() and f.stem.startswith(base_name):
                    new_f_name = f"{re.sub(r'\\{edition-.*?\\}', '', f.stem).strip()}{tag_str}{f.suffix}"
                    new_f_path = folder / new_f_name
                    print(f"→ {f.name} → {new_f_path.name}")
                    f.rename(new_f_path)

            print("[green]Rename complete.[/green]")

        except (ValueError, IndexError):
            print("[red]Invalid selection.[/red]")

# === MODIFIED RUN ===
if __name__ == "__main__":
    flagged = scan_and_compare()
    interactive_mode(flagged)