import os
import subprocess
import time
from pymongo import MongoClient, errors
from concurrent.futures import ThreadPoolExecutor

client = MongoClient("mongodb://root:example@207.244.199.30:27017/")
db = client["niggaflix"]
collection = db["movies"]

collection.create_index("tmdbid", unique=True)

def upload_movie(
    tmdbid: str,
    status: str = "completed",
    progress: int = 0,
    timeLeft: str = "Unknown",
    createdAt:str = time.time(),
    size: int = 0,
    quality: str = "default"
):
   
    movie = {
        "tmdbid": tmdbid,
        "status": status,
        "progress": progress,
        "timeLeft": timeLeft,
        "createdAt": createdAt,
        "size": size,
        "quality": quality
    }

    try:
        result = collection.insert_one(movie)
        print(f"Movie uploaded with id: {result.inserted_id} ({tmdbid})")
    except errors.DuplicateKeyError:
        print(f"Movie with tmdbid {tmdbid} already exists!")

def get_segment_resolution(ts_url: str) -> str:
    cmd = [
        "ffprobe",
        "-v", "error",
        "-select_streams", "v:0",
        "-show_entries", "stream=width,height",
        "-of", "csv=p=0",
        ts_url
    ]
    
    try:
        result = subprocess.run(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True
        )
        line = result.stdout.strip().splitlines()[0]
        parts = line.split(',')
        if len(parts) >= 2:
            width, height = parts[:2]
            return f"{height}p"
        else:
            return "Unknown"
    except Exception:
        return "Unknown"

def get_folder_size(folder_path):
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(folder_path):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            if os.path.isfile(fp):
                total_size += os.path.getsize(fp)
    return total_size

def process_movie(tmdbid, base_dir):
    movie_path = os.path.join(base_dir, tmdbid)
    ts_url = os.path.join(movie_path, "segment_000.ts")
    print(ts_url)

    if not os.path.exists(ts_url):
        print(f"No segment_000.ts for {tmdbid}, skipping.")
        return

    size = get_folder_size(movie_path)
    quality = get_segment_resolution(ts_url)
    upload_movie(
        tmdbid=tmdbid,
        status="completed",
        progress=100,
        timeLeft="0",
        size=size,
        quality=quality
    )

if __name__ == "__main__":
    base_dir = "movies"
    tmdbIds = os.listdir(base_dir)

    with ThreadPoolExecutor(max_workers=20) as executor:
        for tmdbid in tmdbIds:
            executor.submit(process_movie, tmdbid, base_dir)