strm_generator.py

import os
import re
import sqlite3
import requests
import logging
from dotenv import load_dotenv
from tmdbv3api import TMDb, Search

# =====================================================================
# CONFIGURATION
# =====================================================================
# Load environment variables
load_dotenv()
TORBOX_API_TOKEN = os.getenv("TORBOX_API_TOKEN", "")
TMDB_API_KEY = os.getenv("TMDB_API_KEY", "")

# Destination folders under Plex_STRM
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
LOG_DIR = os.path.join(BASE_DIR, "logs")
os.makedirs(LOG_DIR, exist_ok=True)

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(os.path.join(LOG_DIR, "strm_generator.log"), encoding="utf-8"),
        logging.StreamHandler()
    ]
)

DEST_BASE = os.path.join(BASE_DIR, "Plex_STRM")

DEST_MOVIES = os.path.join(DEST_BASE, "Movies")
DEST_TV = os.path.join(DEST_BASE, "TV")
DEST_ANIME_MOVIES = os.path.join(DEST_BASE, "Anime Movies")
DEST_ANIME_TV = os.path.join(DEST_BASE, "Anime TV")
DEST_UNCATEGORIZED = os.path.join(DEST_BASE, "Uncategorized")

DB_PATH = os.path.join(BASE_DIR, "media_cache.db")

# Safety Threshold: Maximum number of links allowed to be auto-deleted in a single cycle.
MAX_DELETIONS_PER_CYCLE = 50

# Setup TMDb Client
tmdb = TMDb()
if TMDB_API_KEY:
    tmdb.api_key = TMDB_API_KEY
tmdb.language = 'en'

# Local Parsing Regex Engines
TV_RE = re.compile(r"([sS]\d+|[sS]eason\s*\d+|[sS]ezon\s*\d+|\d+x\d+|Сезон\s*\d+|Сери[ия]\s*\d+|\[Tenrai-Sensei\])", re.IGNORECASE)
ANIME_GROUP_RE = re.compile(r"^\[(.*?)\]") 
YEAR_RE = re.compile(r"\b(19\d{2}|20\d{2})\b")

# =====================================================================
# LOCAL DATABASE CACHING LAYER
# =====================================================================
def init_db():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS media_cache (
            raw_title TEXT PRIMARY KEY,
            media_type TEXT,
            is_anime INTEGER,
            tmdb_id INTEGER
        )
    ''')
    cursor.execute("PRAGMA table_info(media_cache)")
    columns = [col[1] for col in cursor.fetchall()]
    if 'tmdb_id' not in columns:
        cursor.execute("ALTER TABLE media_cache ADD COLUMN tmdb_id INTEGER")
    conn.commit()
    conn.close()

def get_cached_metadata(raw_title):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute("SELECT media_type, is_anime, tmdb_id FROM media_cache WHERE raw_title = ?", (raw_title,))
    result = cursor.fetchone()
    conn.close()
    return result if result else None

def save_to_cache(raw_title, media_type, is_anime, tmdb_id=None):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute(
        "INSERT OR REPLACE INTO media_cache (raw_title, media_type, is_anime, tmdb_id) VALUES (?, ?, ?, ?)",
        (raw_title, media_type, int(is_anime), tmdb_id)
    )
    conn.commit()
    conn.close()

# =====================================================================
# FILENAME CLEANING & API CALLS
# =====================================================================
import urllib.parse

def clean_title(raw_name):
    clean = urllib.parse.unquote(raw_name)
    clean = re.sub(r"^\[.*?\]", "", clean)
    clean = re.sub(r"\[.*?\]$", "", clean)
    clean = re.sub(r"\(.*?\)$", "", clean)
    clean = re.sub(r"[\._\-]", " ", clean)
    
    # Aggressively remove release group tags, resolutions, audio codecs, and edition tags
    junk_words = r"\b(2160p|1080p|720p|4k|UHD|HDR|HDR10|DV|Dolby Vision|REMUX|BLURAY|BDRip|WEB\-DL|WEBRip|WEB|x264|x265|H264|HEVC|DDP5|AAC|FLAC|TrueHD|Atmos|DTS-HD|DTS|MA|10\-Bit|Hi10|Dual\-Audio|v\d+|Super Duper Cut|Extended|Unrated|DC|Directors Cut|PROPER|REPACK|COMPLETE)\b"
    clean = re.compile(junk_words, re.IGNORECASE).split(clean)[0]
    clean = TV_RE.split(clean)[0]
    
    # Extract year if present
    year_match = YEAR_RE.search(clean)
    year = year_match.group(1) if year_match else None
    
    # Remove year from the clean title for TMDB searching
    if year:
        clean = re.sub(rf"\b{year}\b", "", clean)
        
    return clean.strip(), year

def identify_via_api(raw_name, item_files=None):
    if not TMDB_API_KEY:
        logging.warning("TMDB_API_KEY is not set. Falling back to regex only.")
        return fallback_identify(raw_name, item_files)

    guessed_tv = bool(TV_RE.search(raw_name))
    has_leading_bracket = bool(ANIME_GROUP_RE.search(raw_name))
    
    cleaned, extracted_year = clean_title(raw_name)
    has_year = bool(extracted_year) or bool(YEAR_RE.search(raw_name))
    
    if not cleaned:
        return "uncategorized", False, None

    try:
        # \w includes unicode word characters (like Russian/Cyrillic), \s is whitespace
        cleaned = re.sub(r'[^\w\s]', '', cleaned).strip()
        search = Search()
        
        # Determine media type using regex first, because search.multi often returns TV shows for Movie queries
        guessed_tv = bool(TV_RE.search(raw_name))
        
        # Try specific search endpoints first for maximum accuracy
        results = None
        if guessed_tv:
            results = search.tv_shows(term=cleaned, release_year=extracted_year)
            if results:
                # Add media_type attribute so downstream code works
                for r in results:
                    r.media_type = 'tv'
        else:
            results = search.movies(term=cleaned, year=extracted_year)
            if results:
                for r in results:
                    r.media_type = 'movie'
                    
        # Fallback to multi-search if specific search fails
        if not results:
            results = search.multi(term=cleaned)
        
        if results:
            for res in results:
                media_type = getattr(res, 'media_type', None)
                if media_type not in ['movie', 'tv']:
                    continue
                
                genres = getattr(res, 'genre_ids', [])
                is_animation = 16 in genres
                
                is_japanese = False
                if media_type == 'tv':
                    origin_country = getattr(res, 'origin_country', [])
                    is_japanese = 'JP' in origin_country
                elif media_type == 'movie':
                    prod_countries = getattr(res, 'production_countries', [])
                    is_japanese = any(c.get('iso_3166_1') == 'JP' for c in prod_countries if isinstance(c, dict))
                
                is_anime = is_animation and (is_japanese or has_leading_bracket)
                
                return media_type, is_anime, res.id

    except Exception as e:
        logging.error(f"API Exception routing metadata matching profiles for '{cleaned}': {e}")

    return fallback_identify(raw_name, item_files)

def fallback_identify(raw_name, item_files=None):
    guessed_tv = bool(TV_RE.search(raw_name))
    
    # If the folder name doesn't have a TV tag, check if any of the files inside do
    if not guessed_tv and item_files:
        for f in item_files:
            if TV_RE.search(f):
                guessed_tv = True
                break
                
    has_leading_bracket = bool(ANIME_GROUP_RE.search(raw_name))
    has_year = bool(YEAR_RE.search(raw_name))

    if guessed_tv:
        if has_leading_bracket:
            anime_keywords = ["dual-audio", "10-bit", "hi10p", "english dub", "eng dub", "subbed"]
            raw_lower = raw_name.lower()
            
            if any(keyword in raw_lower for keyword in anime_keywords):
                return "tv", True, None
                
            known_anime_groups = ["yameii", "judas", "tenrai-sensei", "subsplease", "erai-raws", "horriblesubs", "ember", "crucible", "varyg", "anime time", "scy", "datttwisted", "trix", "cleo", "yurasuka", "asw", "nii-sama"]
            group_match = ANIME_GROUP_RE.match(raw_name)
            if group_match:
                group_name = group_match.group(1).lower()
                if any(g in group_name for g in known_anime_groups):
                    return "tv", True, None
                    
        return "tv", False, None
    
    if has_leading_bracket:
        is_tracker_tag = bool(re.search(r'^\[(bitsearch|btsearch|tgx|rartv|yts|rarbg|ettv|eztv|kickass|1337x)\]', raw_name, re.IGNORECASE))
        if has_year:
            return "movie", not is_tracker_tag, None
        return "tv", not is_tracker_tag, None

    if has_year:
        return "movie", False, None

    return "uncategorized", False, None

# =====================================================================
# TORBOX API LOGIC
# =====================================================================
def api_request(endpoint, params=None):
    if not TORBOX_API_TOKEN:
        logging.error("TORBOX_API_TOKEN not set in .env")
        return None
        
    url = f"https://api.torbox.app/v1/api/{endpoint}"
    headers = {
        "Authorization": f"Bearer {TORBOX_API_TOKEN}",
        "Accept": "application/json"
    }
    try:
        response = requests.get(url, headers=headers, params=params)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        logging.error(f"Failed fetching {endpoint}: {e}")
        return None

def get_torbox_items():
    items = []
    
    # Fetch torrents
    res = api_request("torrents/mylist")
    if res and res.get("success") and "data" in res:
        for t in res["data"]:
            if t.get("download_state") == "completed" or t.get("download_finished"):
                items.append({
                    "id": t.get("id"),
                    "name": t.get("name"),
                    "files": t.get("files", []),
                    "type": "torrents"
                })

    # Fetch webdls
    res = api_request("webdl/mylist")
    if res and res.get("success") and "data" in res:
        for w in res["data"]:
            if w.get("download_state") == "completed" or w.get("download_finished"):
                items.append({
                    "id": w.get("id"),
                    "name": w.get("name"),
                    "files": w.get("files", []),
                    "type": "webdl"
                })
                
    return items

def get_requestdl_url(item_type, item_id, file_id, filename):
    # Extract extension (e.g., '.mkv', '.mp4') to trick Plex into direct-playing the stream
    ext = os.path.splitext(filename)[1].lower() if filename else ".mkv"
    if item_type == "torrents":
        return f"https://api.torbox.app/v1/api/torrents/requestdl?token={TORBOX_API_TOKEN}&torrent_id={item_id}&file_id={file_id}&redirect=true#{ext}"
    else:
        return f"https://api.torbox.app/v1/api/webdl/requestdl?token={TORBOX_API_TOKEN}&web_id={item_id}&file_id={file_id}&redirect=true#{ext}"

# =====================================================================
# SYNC ENGINE
# =====================================================================
def setup_dirs():
    all_paths = [DEST_MOVIES, DEST_TV, DEST_ANIME_MOVIES, DEST_ANIME_TV, DEST_UNCATEGORIZED]
    for path in all_paths:
        os.makedirs(path, exist_ok=True)

def is_video_file(filename):
    ext = os.path.splitext(filename)[1].lower()
    return ext in ['.mkv', '.mp4', '.avi', '.ts', '.webm', '.m2ts', '.wmv', '.mov']

def sync_media():
    if not TORBOX_API_TOKEN:
        logging.warning("TorBox API Token missing. Aborting sync. Please populate .env file.")
        return

    logging.info("Fetching active items from TorBox...")
    active_items = get_torbox_items()
    
    # 2. Empty Source Shield
    if not active_items:
        logging.warning("API returned 0 active items. Network drop suspected. Skipping sync cycle to protect links.")
        return

    logging.info(f"Found {len(active_items)} active items. Generating STRM files...")
    all_dest_folders = [DEST_MOVIES, DEST_TV, DEST_ANIME_MOVIES, DEST_ANIME_TV, DEST_UNCATEGORIZED]
    generated_strm_paths = set()
    
    # Cache manual overrides locally once for performance
    valid_override_folders = {
        DEST_MOVIES: ("movie", False),
        DEST_TV: ("tv", False),
        DEST_ANIME_MOVIES: ("movie", True),
        DEST_ANIME_TV: ("tv", True)
    }
    dir_cache = {}
    for folder in valid_override_folders:
        try:
            dir_cache[folder] = os.listdir(folder)
        except FileNotFoundError:
            dir_cache[folder] = []

    # --- CREATION PHASE ---
    for item in active_items:
        raw_name = item["name"]
        
        # If the torrent name is just a 40-character Info Hash OR a raw magnet link, try to grab a real name from the video files
        if re.match(r'^[a-fA-F0-9]{40}$', raw_name) or raw_name.lower().startswith("magnet"):
            video_files = [f for f in item.get("files", []) if is_video_file(f.get("name", ""))]
            if video_files:
                first_video_name = video_files[0].get("short_name", video_files[0].get("name", ""))
                better_name = os.path.splitext(first_video_name)[0]
                if better_name:
                    logging.info(f"Info Hash detected '{raw_name}'. Using filename '{better_name}' as the folder name instead.")
                    raw_name = better_name
                    
        # Ensure the folder name does not exceed filesystem limits (255 chars max)
        folder_name = raw_name
        
        # Sanitize folder name for Windows and other OS invalid characters
        folder_name = re.sub(r'[\\/*?:"<>|]', "", folder_name)
        
        # Strip Russian/Cyrillic characters ONLY if English characters are also present
        # This prevents Plex from getting confused by hybrid titles like "Жестокое лето Cruel Summer"
        if re.search(r'[А-Яа-яЁё]', folder_name) and re.search(r'[A-Za-z]', folder_name):
            folder_name = re.sub(r'[А-Яа-яЁё]+', '', folder_name)
            folder_name = re.sub(r'\s+', ' ', folder_name)
            folder_name = re.sub(r'^[\s\-/|]+', '', folder_name).strip()
            
        # Format the folder name to include the (Year) so Plex matches it flawlessly even without TMDB
        _, extracted_year = clean_title(raw_name)
        if extracted_year and not re.search(rf'\({extracted_year}\)', folder_name):
            # If the year is already in the folder name, let's format it with parentheses
            folder_name = re.sub(rf'\b{extracted_year}\b', f'({extracted_year})', folder_name)
            
        # Note: encoding to utf-8 to correctly handle multi-byte chars (like Cyrillic)
        if len(folder_name.encode('utf-8')) > 255:
            # truncate by characters, checking byte length
            while len(folder_name.encode('utf-8')) > 255:
                folder_name = folder_name[:-1]
            folder_name = folder_name.strip()
            
        cached = get_cached_metadata(raw_name)
        media_type, is_anime, tmdb_id = None, None, None
        # If it's cached but uncategorized, we ignore the cache so we can retry the API/Fallback engine!
        # Unless the user explicitly forced it to stay there with 'force_uncategorized'
        if cached and (cached[0] != "uncategorized" or cached[0] == "force_uncategorized"):
            media_type, is_anime, tmdb_id = get_cached_metadata(raw_name)
        
        # --- LEGACY RD SYNC OVERRIDE ---
        rd_override = None
        # Use an environment variable for the Plex Media Dir if a user wants to bridge legacy setups
        RD_MEDIA_DIR = os.getenv("PLEX_MEDIA_DIR", "")
        if RD_MEDIA_DIR and os.path.exists(RD_MEDIA_DIR):
            for cat in ["Anime Movies", "Anime TV", "Movies", "TV"]:
                if os.path.exists(os.path.join(RD_MEDIA_DIR, cat, raw_name)):
                    rd_override = cat
                    break
        
        file_names = [f.get("name", "") for f in item.get("files", [])]
        if media_type is None:
            api_media_type, api_is_anime, tmdb_id = identify_via_api(raw_name, file_names)
            
            if rd_override == "Anime TV":
                media_type, is_anime = "tv", True
            elif rd_override == "Anime Movies":
                media_type, is_anime = "movie", True
            elif rd_override == "Movies":
                media_type, is_anime = "movie", False
            elif rd_override == "TV":
                media_type, is_anime = "tv", False
            else:
                media_type, is_anime = api_media_type, api_is_anime
                
            save_to_cache(raw_name, media_type, is_anime, tmdb_id)
            logging.info(f"Analyzed '{raw_name}' -> Type: {media_type.upper()}, Anime: {bool(is_anime)}, TMDB: {tmdb_id} (RD Override: {rd_override})")
        # --- MANUAL OVERRIDE CHECK ---
        # If the user manually moved the folder to a valid category inside Plex_STRM, respect that location
        found_override = False
        for folder, (o_media_type, o_is_anime) in valid_override_folders.items():
            if found_override: break
            if folder in dir_cache:
                for entry in dir_cache[folder]:
                    # Match if the entry is the exact folder name OR starts with folder_name + " {tmdb-"
                    # This safely matches any manual TMDB ID changes made by the user
                    if entry == folder_name or entry.startswith(f"{folder_name} {{tmdb-"):
                        # Extract the tmdb_id from the folder if it was manually changed by the user!
                        match = re.search(r'\{tmdb-(\d+)\}', entry)
                        new_tmdb_id = match.group(1) if match else tmdb_id
                        
                        if media_type != o_media_type or is_anime != o_is_anime or tmdb_id != new_tmdb_id:
                            media_type = o_media_type
                            is_anime = o_is_anime
                            tmdb_id = new_tmdb_id
                            save_to_cache(raw_name, media_type, is_anime, tmdb_id)
                        found_override = True
                        break

        if media_type == "uncategorized" or media_type == "force_uncategorized":
            target_dir = DEST_UNCATEGORIZED
        elif is_anime:
            target_dir = DEST_ANIME_TV if media_type == "tv" else DEST_ANIME_MOVIES
        else:
            target_dir = DEST_TV if media_type == "tv" else DEST_MOVIES
            
        # Magic Match: Append TMDB ID to folder name to guarantee flawless Plex matching!
        if tmdb_id and not re.search(r'\{tmdb-\d+\}', folder_name):
            folder_name = f"{folder_name} {{tmdb-{tmdb_id}}}"
            
        base_item_dir = os.path.join(target_dir, folder_name)
        os.makedirs(base_item_dir, exist_ok=True)

        for f in item["files"]:
            filename = f.get("name", "")
            short_name = f.get("short_name", filename)
            file_id = f.get("id")
            
            # Sanitize short_name for Windows
            short_name = re.sub(r'[\\/*?:"<>|]', "", short_name)
            
            if not is_video_file(filename):
                continue
            
            # Use short_name to construct the STRM path preserving internal folder structure
            strm_rel_path = os.path.splitext(short_name)[0] + ".strm"
            target_strm_path = os.path.join(base_item_dir, strm_rel_path)
            
            os.makedirs(os.path.dirname(target_strm_path), exist_ok=True)
            generated_strm_paths.add(os.path.abspath(target_strm_path))
            
            if not os.path.exists(target_strm_path):
                strm_url = get_requestdl_url(item["type"], item["id"], file_id, filename)
                try:
                    with open(target_strm_path, "w", encoding="utf-8") as strm_file:
                        strm_file.write(strm_url)
                    logging.info(f"Created STRM: {strm_rel_path} -> {os.path.basename(target_dir)}")
                except Exception as e:
                    logging.error(f"Failed to create STRM for {short_name}: {e}")

    # --- SANITATION / CLEANUP PHASE ---
    dead_links = []
    
    for folder in all_dest_folders:
        if not os.path.exists(folder): continue
        for root, dirs, files in os.walk(folder):
            for file in files:
                if file.endswith(".strm"):
                    full_path = os.path.abspath(os.path.join(root, file))
                    if full_path not in generated_strm_paths:
                        dead_links.append(full_path)

    # 3. Partial Disconnect Shield
    if len(dead_links) > MAX_DELETIONS_PER_CYCLE:
        logging.warning(f"SHIELD ACTIVE: Detected {len(dead_links)} dead .strm links at once.")
        logging.warning("Aborting cleanup phase to prevent mass-deletion.")
    else:
        for link_path in dead_links:
            try:
                os.remove(link_path)
                logging.info(f"Removed dead link: {os.path.basename(link_path)}")
            except Exception as e:
                logging.error(f"Failed to remove dead link {os.path.basename(link_path)}: {e}")

    # Cleanup empty directories
    for base_dir in all_dest_folders:
        if not os.path.exists(base_dir): continue
        for root, dirs, files in os.walk(base_dir, topdown=False):
            for name in dirs:
                dir_path = os.path.join(root, name)
                contents = os.listdir(dir_path)
                # If the only thing left is a macOS .DS_Store file, delete it so the folder can be removed
                if contents == ['.DS_Store']:
                    os.remove(os.path.join(dir_path, '.DS_Store'))
                    contents = []
                
                if not contents:
                    try:
                        os.rmdir(dir_path)
                        logging.info(f"Removed empty directory: {dir_path}")
                    except Exception as e:
                        logging.error(f"Failed to remove directory {dir_path}: {e}")

if __name__ == "__main__":
    init_db()
    setup_dirs()
    logging.info("TorBox Bridge Online.")
    sync_media()
    logging.info("Sync complete.")