profilarr/import.py

import json
import requests
import os

# ANSI escape sequences for colors
class Colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

# Load configuration for main app
with open('config.json', 'r') as config_file:
    config = json.load(config_file)

def print_success(message):
    print(Colors.OKGREEN + message + Colors.ENDC)

def print_error(message):
    print(Colors.FAIL + message + Colors.ENDC)

def print_connection_error():
    print(Colors.FAIL + "Failed to connect to the service! Please check if it's running and accessible." + Colors.ENDC)

def get_user_choice():
    choice = input("Enter the app you want to import to (radarr/sonarr): ").lower()
    while choice not in ["radarr", "sonarr"]:
        print_error("Invalid input. Please enter either 'radarr' or 'sonarr'.")
        choice = input("Enter the source (radarr/sonarr): ").lower()
    return choice

def get_import_choice():
    print()
    print(Colors.HEADER + "Choose what to import:" + Colors.ENDC)
    print("1. Custom Formats")
    print("2. Quality Profiles")
    choice = input("Enter your choice (1/2): ").strip()
    while choice not in ["1", "2"]:
        print_error("Invalid input. Please enter 1 or 2.")
        choice = input("Enter your choice (1/2): ").strip()
    return choice

def get_app_config(source):
    return config['master'][source]

def select_file(directory):
    files = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]
    print()
    print(Colors.OKBLUE + "Available files:" + Colors.ENDC)
    for i, file in enumerate(files, 1):
        print(f"{i}. {file}")
    choice = int(input("Select a file to import: "))
    return files[choice - 1]

def import_custom_formats(source_config):
    headers = {"X-Api-Key": source_config['api_key']}
    get_url = f"{source_config['base_url']}/api/v3/customformat"
    try:
        response = requests.get(get_url, headers=headers)
        if response.status_code == 200:
            existing_formats = response.json()
            existing_names_to_id = {format['name']: format['id'] for format in existing_formats}

            selected_file = select_file('./custom_formats')
            added_count, updated_count = 0, 0

            with open(os.path.join('./custom_formats', selected_file), 'r') as import_file:
                import_formats = json.load(import_file)

            print()

            for format in import_formats:
                format_name = format['name']
                if format_name in existing_names_to_id:
                    format_id = existing_names_to_id[format_name]
                    put_url = f"{source_config['base_url']}/api/v3/customformat/{format_id}"
                    response = requests.put(put_url, json=format, headers=headers)
                    if response.status_code in [200, 201, 202]:
                        print(Colors.WARNING + f"Updating custom format '{format_name}': " + Colors.ENDC, end='')
                        print_success("SUCCESS")
                        updated_count += 1
                    else:
                        print_error(f"Updating custom format '{format_name}': FAIL")
                        print(response.content.decode())

                else:
                    post_url = f"{source_config['base_url']}/api/v3/customformat"
                    response = requests.post(post_url, json=format, headers=headers)
                    if response.status_code in [200, 201]:
                        print(Colors.OKBLUE + f"Adding custom format '{format_name}': " + Colors.ENDC, end='')
                        print_success("SUCCESS")
                        added_count += 1
                    else:
                        print_error(f"Adding custom format '{format_name}': FAIL")
                        print(response.content.decode())

            print()
            print_success(f"Successfully added {added_count} custom formats, updated {updated_count} custom formats.")

        else:
            print_error(f"Failed to retrieve existing custom formats from {get_url}! (HTTP {response.status_code})")
            print(response.content.decode())

    except requests.exceptions.ConnectionError:
        print_connection_error()

def import_quality_profiles(source_config):
    headers = {"X-Api-Key": source_config['api_key']}
    try:
        cf_import_sync(source_config)

        profile_dir = './profiles'
        profiles = [f for f in os.listdir(profile_dir) if f.endswith('.json')]

        print()
        print(Colors.HEADER + "Available Profiles:" + Colors.ENDC)
        for i, profile in enumerate(profiles, 1):
            print(f"{i}. {profile}")

        print()
        selection = input("Please enter the number of the profile you want to import: ")
        try:
            selected_file = profiles[int(selection) - 1]
        except (ValueError, IndexError):
            print_error("Invalid selection, please enter a valid number.")
            return

        with open(os.path.join(profile_dir, selected_file), 'r') as file:
            try:
                quality_profiles = json.load(file)
            except json.JSONDecodeError as e:
                print_error(f"Error loading selected profile: {e}")
                return

            for profile in quality_profiles:
                existing_format_names = set()
                if 'formatItems' in profile:
                    for format_item in profile['formatItems']:
                        format_name = format_item.get('name')
                        if format_name:
                            existing_format_names.add(format_name)
                            if format_name in source_config['custom_formats']:
                                format_item['format'] = source_config['custom_formats'][format_name]

                for format_name, format_id in source_config['custom_formats'].items():
                    if format_name not in existing_format_names:
                        profile.setdefault('formatItems', []).append({
                            "format": format_id,
                            "name": format_name,
                            "score": 0
                        })

                post_url = f"{source_config['base_url']}/api/v3/qualityprofile"
                response = requests.post(post_url, json=profile, headers=headers)

                if response.status_code in [200, 201]:
                    print_success(f"Successfully added Quality Profile {profile['name']}")
                elif response.status_code == 409:
                    print_error(f"Failed to add Quality Profile {profile['name']} due to a naming conflict. Quality profile names must be unique. (HTTP {response.status_code})")
                else:
                    try:
                        errors = response.json()
                        message = errors.get("message", "No Message Provided")
                        print_error(f"Failed to add Quality Profile {profile['name']}! (HTTP {response.status_code})")
                        print(message)
                    except json.JSONDecodeError:
                        print_error("Failed to parse error message:")
                        print(response.text)

    except requests.exceptions.ConnectionError:
        print_connection_error()

def cf_import_sync(source_config):
    headers = {"X-Api-Key": source_config['api_key']}
    custom_format_url = f"{source_config['base_url']}/api/v3/customformat"
    try:
        response = requests.get(custom_format_url, headers=headers)
        if response.status_code == 200:
            data = response.json()
            source_config['custom_formats'] = {format['name']: format['id'] for format in data}
        elif response.status_code == 401:
            print_error("Authentication error: Invalid API key. Terminating program.")
            exit(1)
        else:
            print_error(f"Failed to retrieve custom formats! (HTTP {response.status_code})")
            print(response.content.decode())
            exit(1)

    except requests.exceptions.ConnectionError:
        print_connection_error()
        exit(1)


if __name__ == "__main__":
    user_choice = get_user_choice()
    source_config = get_app_config(user_choice)
    import_choice = get_import_choice()

    if import_choice == "1":
        import_custom_formats(source_config)
    elif import_choice == "2":
        import_quality_profiles(source_config)