#!/usr/bin/python3
#
#   qBittorrent-api:
#   https://qbittorrent-api.readthedocs.io/en/v2023.4.47/introduction.html
#
#
#
#   Make sure to set the folling variables. 
#       <USERNAME>
#       <SERVER IP ADDRESS> 
#       <PORT>
#
#
import re
import sys
import logging
import qbittorrentapi

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Checking for correct number of command-line arguments
if len(sys.argv) != 3:
    logging.error('[ACAT] - Incorrect number of arguments. Exiting script.')
    sys.exit(1)

# Extract torrent name and hash from command line arguments
torrent_name = sys.argv[1]
torrent_hash = sys.argv[2]

logging.warning('[ACAT] - -------------')
logging.warning('[ACAT] - -------------')
logging.warning('[ACAT] - *** STARTING ASSIGN CATEGORY SCRIPT ***')
logging.warning('[ACAT] - -------------')
logging.warning(f'[ACAT] - raw torrent name: {torrent_name}')
logging.warning(f'[ACAT] - raw torrent hash: {torrent_hash}')
logging.warning('[ACAT] - -------------')

# qbittorrent webui login crendentials.
qbitt_username = '<USERNAME>'
qbitt_password = '<PASSWORD>'

# qbittorrent webui address.
#url = 'http://192.168.1.55:8124'
#url = 'http://localhost:8124'
url = 'http://<SERVER IP ADDRESS>:<PORT>'

# The categories to use for matching torrents. Change as needed. 
category_1080_shows = "1080+ | Shows"
category_1080_movies = "1080+ | Movies"

try:
    # Connect to the qBittorrent API
    qb = qbittorrentapi.Client(host=url)
    qb.auth_log_in(username=qbitt_username, password=qbitt_password)

    # Define the regular expression pattern for 'SxxExx' or 'sxxexx' format.
    pattern = r"[Ss]\d{2}[Ee]\d{2}"

    # Check for 'SxxExx' or 'sxxexx' episode pattern and '1080p' or '2160p'
    if re.search(pattern, torrent_name) and ("1080p" in torrent_name or "2160p" in torrent_name):
        qb.torrents_set_category(category=category_1080_shows, hashes=torrent_hash)
        qb.torrents_set_auto_management(hashes=torrent_hash, enable=True)
        logging.warning(f'[ACAT] - Torrent "{torrent_name}" category set to: "{category_1080_shows}".')
        logging.warning('[ACAT] - -------------')

    # Then check for movies criteria.
    elif ("1080p" in torrent_name or "2160p" in torrent_name) and "YTS.MX" in torrent_name:
        qb.torrents_set_category(category=category_1080_movies, hashes=torrent_hash)
        qb.torrents_set_auto_management(hashes=torrent_hash, enable=True) 
        logging.warning(f'[ACAT] - Torrent "{torrent_name}" category set to: "{category_1080_movies}".')
        logging.warning('[ACAT] - -------------')

    # Torrent is not a Hi-def movie or show
    else:
        logging.warning(f'[ACAT] - No category change for Torrent "{torrent_name}".')
        logging.warning('[ACAT] - -------------')

    # Disconnect from the qBittorrent API
    qb.auth_log_out()

except Exception as e:
    logging.error(f'[ACAT] - Error occurred: {e}')
    logging.warning('[ACAT] - -------------')

logging.warning('[ACAT] - *** ENDING ASSIGN CATEGORY SCRIPT ***')
logging.warning('[ACAT] - -------------')
logging.warning('[ACAT] - -------------')