import sys

import subprocess

import argparse

def install_package(package):

	subprocess.check_call([sys.executable, "-m", "pip", "install", package])

# Check and install required packages

required_packages = ['plexapi', 'requests', 'fuzzywuzzy', 'colorama', 'python-Levenshtein']

for package in required_packages:

	try:

		__import__(package.replace('-', '_'))  # Replace hyphen with underscore for import

	except ImportError:

		print(f"Installing {package}...")

		install_package(package)

import requests

from plexapi.server import PlexServer

import re

import csv

import time

import os

import configparser

from datetime import datetime

from fuzzywuzzy import fuzz

from fuzzywuzzy import process

from colorama import init, Fore, Style

CONFIG_FILE = 'plex_collection_config.ini'

# Initialize colorama

init()

### Plex server details - will be set via args or user input ###

PLEX_URL = ''

PLEX_TOKEN = ''

MOVIE_LIBRARY_NAME = ''

CSV_FILE_PATH = ''

IMDB_COLLECTION_NAME = ''

LOG_FILE_PATH = ''

REPORT_FILE_PATH = ''

# Fuzzy matching thresholds

TITLE_YEAR_MATCH_RATIO = 85  # Minimum match ratio for title + year

TITLE_ONLY_MATCH_RATIO = 90  # Minimum match ratio for title only

def load_config():

	"""Load saved configuration"""

	config = configparser.ConfigParser()

	if os.path.exists(CONFIG_FILE):

		config.read(CONFIG_FILE)

	return config

def save_config(url, token, library):

	"""Save configuration to file"""

	config = configparser.ConfigParser()

	config['DEFAULT'] = {

		'PLEX_URL': url,

		'PLEX_TOKEN': token,

		'MOVIE_LIBRARY_NAME': library

	}

	with open(CONFIG_FILE, 'w') as configfile:

		config.write(configfile)

def display_instructions():

	"""Display instructions for users who need guidance"""

	print(f"\n{Fore.CYAN}=== Instructions ==={Style.RESET_ALL}")

	print(f"{Fore.WHITE}1. Go to an IMDB list and export the list as a .CSV (like https://www.imdb.com/list/ls569744078/)")

	print(f"2. Rename that .CSV file from the IMDB list to something more memorable and friendly")

	print(f"3. Get your Plex token via these instructions: https://support.plex.tv/articles/204059436-finding-an-authentication-token-x-plex-token/{Style.RESET_ALL}\n")

def get_existing_collections():

	"""Get list of existing collections in the movie library"""

	try:

		plex = PlexServer(PLEX_URL, PLEX_TOKEN)

		movie_library = plex.library.section(MOVIE_LIBRARY_NAME)

		# Get all collections from the movie library

		collections = movie_library.collections()

		return [collection.title for collection in collections]

	except Exception as e:

		log_message(f"Error retrieving existing collections: {str(e)}", message_type="ERROR", print_to_console=True)

		return []

def parse_arguments():

	"""Parse command line arguments"""

	parser = argparse.ArgumentParser(description="Plex Collection Creator")

	parser.add_argument('-u', '--url', help='Plex server URL (default: http://127.0.0.1:32400)')

	parser.add_argument('-t', '--token', help='Plex authentication token')

	parser.add_argument('-c', '--csv', help='Path to the CSV file containing the movie list')

	parser.add_argument('-n', '--name', help='Name for the collection to be created or updated')

	parser.add_argument('-l', '--library', help='Plex movie library name (default: Movies)')

	parser.add_argument('-i', '--instructions', action='store_true', help='Display instructions and exit')

	return parser.parse_args()

def initialize_logging():

	"""Create log file with timestamp"""

	global LOG_FILE_PATH, REPORT_FILE_PATH

	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

	LOG_FILE_PATH = f'plex_collection_log_{IMDB_COLLECTION_NAME}_{timestamp}.txt'

	REPORT_FILE_PATH = f'plex_collection_report_{IMDB_COLLECTION_NAME}_{timestamp}.csv'

	with open(LOG_FILE_PATH, 'w', encoding='utf-8') as log_file:

		log_file.write(f"Plex Collection Creation Log - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")

		log_file.write(f"Collection Name: {IMDB_COLLECTION_NAME}\n")

		log_file.write(f"CSV File: {CSV_FILE_PATH}\n")

		log_file.write(f"Plex URL: {PLEX_URL}\n")

		log_file.write(f"Plex Library: {MOVIE_LIBRARY_NAME}\n")

		log_file.write("="*50 + "\n\n")

def log_message(message, print_to_console=True, message_type="INFO"):

	"""Log messages to file and optionally print to console with colors"""

	timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

	log_entry = f"[{timestamp}] {message}\n"

	with open(LOG_FILE_PATH, 'a', encoding='utf-8') as log_file:

		log_file.write(log_entry)

	if print_to_console:

		if message_type == "SUCCESS":

			print(f"{Fore.GREEN}{message}{Style.RESET_ALL}")

		elif message_type == "ERROR":

			print(f"{Fore.RED}{message}{Style.RESET_ALL}")

		elif message_type == "WARNING":

			print(f"{Fore.YELLOW}{message}{Style.RESET_ALL}")

		else:

			print(message)

def add_collection(library_key, rating_key, movie_title):

	"""Add movie to collection and return success status"""

	headers = {"X-Plex-Token": PLEX_TOKEN}

	params = {

		"type": 1,

		"id": rating_key,

		"collection[0].tag.tag": IMDB_COLLECTION_NAME,

		"collection.locked": 1

	}

	url = f"{PLEX_URL}/library/sections/{library_key}/all"

	try:

		response = requests.put(url, headers=headers, params=params)

		return response.status_code == 200

	except Exception as e:

		log_message(f"Error adding collection for {movie_title}: {str(e)}", message_type="ERROR")

		return False

def find_imdb_id(imdb_url):

	"""Extract id from imdb_url"""

	match = re.search(r'(tt\d+)', imdb_url)

	return match.group(0) if match else None

def clean_title(title):

	"""Clean title for better matching"""

	if not title:

		return ""

	return title.lower().replace("the ", "").strip()

def find_movie_by_title_year(movies, title, year, original_title=None):

	"""Find movie by title and year with fuzzy matching"""

	# Try exact match first

	for movie in movies:

		if (movie.title.lower() == title.lower() or 

			(original_title and movie.title.lower() == original_title.lower())):

			if str(movie.year) == str(year):

				return movie

	# Try fuzzy matching

	best_match = None

	highest_score = 0

	for movie in movies:

		# Compare with main title

		title_score = fuzz.ratio(clean_title(movie.title), clean_title(title))

		# Compare with original title if available

		original_score = 0

		if original_title:

			original_score = fuzz.ratio(clean_title(movie.title), clean_title(original_title))

		current_score = max(title_score, original_score)

		# Only consider if year matches

		if str(movie.year) == str(year) and current_score > highest_score:

			highest_score = current_score

			best_match = movie

	if best_match and highest_score >= TITLE_YEAR_MATCH_RATIO:

		log_message(f"Fuzzy matched '{title}' ({year}) to '{best_match.title}' ({best_match.year}) with score {highest_score}", message_type="WARNING")

		return best_match

	return None

def find_movie_by_title(movies, title, original_title=None):

	"""Find movie by title only with fuzzy matching"""

	best_match = None

	highest_score = 0

	for movie in movies:

		title_score = fuzz.ratio(clean_title(movie.title), clean_title(title))

		original_score = 0

		if original_title:

			original_score = fuzz.ratio(clean_title(movie.title), clean_title(original_title))

		current_score = max(title_score, original_score)

		if current_score > highest_score:

			highest_score = current_score

			best_match = movie

	if best_match and highest_score >= TITLE_ONLY_MATCH_RATIO:

		log_message(f"Title-only matched '{title}' to '{best_match.title}' with score {highest_score}", message_type="WARNING")

		return best_match

	return None

def find_existing_collection_items(plex, collection_name):

	"""Find all movies that are already in the specified collection"""

	existing_items = {}

	# Search all movie sections

	try:

		movie_library = plex.library.section(MOVIE_LIBRARY_NAME)

		collection_items = movie_library.search(collection=collection_name)

		# Create a map of both IMDB IDs and rating keys

		for item in collection_items:

			# Add by rating key

			existing_items[item.ratingKey] = item

			# Add by IMDB ID if available

			if 'imdb://' in item.guid:

				imdb_id = item.guid.split('imdb://')[1].split('?')[0]

				existing_items[imdb_id] = item

		log_message(f"Found {len(collection_items)} movies already in collection '{collection_name}'")

		return existing_items

	except Exception as e:

		log_message(f"Error finding existing collection items: {str(e)}", message_type="ERROR")

		return {}

def generate_report(report_data):

	"""Generate CSV report of processed movies"""

	with open(REPORT_FILE_PATH, 'w', newline='', encoding='utf-8') as csvfile:

		fieldnames = ['Title', 'Year', 'Original Title', 'IMDB_ID', 'Status', 'Match_Method', 'Plex_Title', 'Plex_Year', 'Notes']

		writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

		writer.writeheader()

		for item in report_data:

			writer.writerow({

				'Title': item['title'],

				'Year': item['year'],

				'Original Title': item.get('original_title', ''),

				'IMDB_ID': item['imdb_id'],

				'Status': item['status'],

				'Match_Method': item.get('match_method', ''),

				'Plex_Title': item.get('plex_title', ''),

				'Plex_Year': item.get('plex_year', ''),

				'Notes': item.get('notes', '')

			})

def verify_plex_connection():

	"""Verify connection to Plex server"""

	try:

		plex = PlexServer(PLEX_URL, PLEX_TOKEN)

		log_message(f"Connected to Plex server at {PLEX_URL}")

		return plex

	except Exception as e:

		log_message(f"Failed to connect to Plex server: {str(e)}", message_type="ERROR")

		print(f"\n{Fore.RED}Unable to connect to Plex server. Please verify your URL and token.{Style.RESET_ALL}")

		display_instructions()

		return None

def run_imdb_list():

	initialize_logging()

	report_data = []

	# Verify Plex connection

	plex = verify_plex_connection()

	if not plex:

		return []

	# Read local CSV file

	log_message(f"Reading IMDB list from: {CSV_FILE_PATH}")

	try:

		with open(CSV_FILE_PATH, mode='r', encoding='utf-8') as csv_file:

			reader = csv.DictReader(csv_file)

			rows = list(reader)

			url_column = 'URL' if 'URL' in reader.fieldnames else 'url'

			log_message(f"Found {len(rows)} movies in the CSV file")

	except Exception as e:

		log_message(f"Error reading CSV file: {str(e)}", message_type="ERROR")

		return []

	# Get all movies from Plex library

	log_message(f"Retrieving movies from Plex library: {MOVIE_LIBRARY_NAME}")

	try:

		movie_library = plex.library.section(MOVIE_LIBRARY_NAME)

		all_movies = movie_library.all()

		log_message(f"Found {len(all_movies)} movies in Plex library")

	except Exception as e:

		log_message(f"Error accessing Plex library: {str(e)}", message_type="ERROR")

		print(f"\n{Fore.RED}Error: Library '{MOVIE_LIBRARY_NAME}' not found. Please check your library name.{Style.RESET_ALL}")

		return []

	# Find existing collection items

	existing_collection_items = find_existing_collection_items(plex, IMDB_COLLECTION_NAME)

	# Create IMDB ID to movie mapping

	imdb_map = {}

	for movie in all_movies:

		if 'imdb://' in movie.guid:

			imdb_id = movie.guid.split('imdb://')[1].split('?')[0]

			imdb_map[imdb_id] = movie

	# Process each movie from CSV in order

	log_message(f"\nStarting collection creation/update for {IMDB_COLLECTION_NAME}")

	added_count = 0

	skipped_count = 0

	missing_count = 0

	# First pass - find all matches and create ordered list

	ordered_movies = []

	for row in rows:

		imdb_url = row.get(url_column, '')

		imdb_id = find_imdb_id(imdb_url)

		title = row.get('Title', 'Unknown')

		original_title = row.get('Original Title', '')

		year = row.get('Year', '')

		report_item = {

			'title': title,

			'year': year,

			'original_title': original_title,

			'imdb_id': imdb_id or '',

			'status': '',

			'plex_title': '',

			'plex_year': '',

			'notes': ''

		}

		if not imdb_id:

			status = "ERROR: No IMDB ID found"

			report_item.update({

				'status': status,

				'notes': f'Invalid URL: {imdb_url}'

			})

			log_message(f"{status} - {title} ({year})", message_type="ERROR")

			missing_count += 1

			report_data.append(report_item)

			continue

		# Try to find movie by IMDB ID first

		movie = imdb_map.get(imdb_id)

		match_method = "IMDB ID" if movie else ""

		# If not found by IMDB ID, try title + year

		if not movie:

			movie = find_movie_by_title_year(all_movies, title, year, original_title)

			if movie:

				match_method = "Title + Year Fuzzy Match"

		# If still not found, try title only

		if not movie:

			movie = find_movie_by_title(all_movies, title, original_title)

			if movie:

				match_method = "Title Only Fuzzy Match"

				# Verify year isn't too different

				if year and abs(int(movie.year) - int(year)) > 2:

					log_message(f"Rejected title-only match for {title} due to year mismatch (Plex: {movie.year}, CSV: {year})", message_type="WARNING")

					movie = None

		if not movie:

			status = "MISSING: Not found in Plex library"

			report_item.update({

				'status': status,

				'notes': 'Could not match by IMDB ID, title+year, or title only'

			})

			log_message(f"{status} - {title} ({year})", message_type="ERROR")

			missing_count += 1

			report_data.append(report_item)

			continue

		# Check if the movie is already in the collection

		already_in_collection = False

		if movie.ratingKey in existing_collection_items:

			already_in_collection = True

		elif imdb_id in existing_collection_items:

			already_in_collection = True

		if already_in_collection:

			status = "SKIPPED:"

			report_item.update({

				'status': status,

				'match_method': match_method,

				'plex_title': movie.title,

				'plex_year': movie.year

			})

			log_message(f"{status} {movie.title} ({movie.year})", message_type="INFO")

			skipped_count += 1

			report_data.append(report_item)

			continue

		# Add to our ordered list for movies not in collection

		ordered_movies.append({

			'movie': movie,

			'report_item': report_item,

			'match_method': match_method,

			'original_position': len(ordered_movies)  # Maintain original order

		})

	# Second pass - add to collection in order

	for item in ordered_movies:

		movie = item['movie']

		report_item = item['report_item']

		match_method = item['match_method']

		# Try to add to collection

		success = add_collection(movie.librarySectionID, movie.ratingKey, movie.title)

		if success:

			status = "ADDED:"

			report_item.update({

				'status': status,

				'match_method': match_method,

				'plex_title': movie.title,

				'plex_year': movie.year

			})

			message_type = "SUCCESS" if match_method == "IMDB ID" else "WARNING"

			log_message(f"{status} {movie.title} ({movie.year}) - (matched by {match_method})", message_type=message_type)

			added_count += 1

		else:

			status = "ERROR:"

			report_item.update({

				'status': status,

				'match_method': match_method,

				'plex_title': movie.title,

				'plex_year': movie.year,

				'notes': 'Plex API error'

			})

			log_message(f"{status} {movie.title} ({movie.year})", message_type="ERROR")

			missing_count += 1

		report_data.append(report_item)

	# Generate reports

	log_message(f"\nProcessing complete: {Fore.GREEN}{added_count} added{Style.RESET_ALL}, {Fore.YELLOW}{skipped_count} already in collection{Style.RESET_ALL}, {Fore.RED}{missing_count} not found/failed{Style.RESET_ALL}.")

	generate_report(report_data)

	log_message(f"Detailed report saved to: {REPORT_FILE_PATH}")

	log_message(f"Full log saved to: {LOG_FILE_PATH}")

	return report_data

def test_plex_connection():

	"""Test connection to Plex server and validate library name"""

	try:

		plex = PlexServer(PLEX_URL, PLEX_TOKEN)

		print(f"{Fore.GREEN}✓ Successfully connected to Plex server{Style.RESET_ALL}")

		# Validate library name

		try:

			library = plex.library.section(MOVIE_LIBRARY_NAME)

			print(f"{Fore.GREEN}✓ Successfully found library: {MOVIE_LIBRARY_NAME} (contains {len(library.all())} items){Style.RESET_ALL}")

			return True

		except Exception as e:

			print(f"{Fore.RED}✗ Library '{MOVIE_LIBRARY_NAME}' not found{Style.RESET_ALL}")

			print(f"{Fore.CYAN}Available libraries:{Style.RESET_ALL}")

			for section in plex.library.sections():

				print(f"  - {section.title}")

			return False

	except Exception as e:

		print(f"{Fore.RED}✗ Failed to connect to Plex server: {str(e)}{Style.RESET_ALL}")

		return False

def get_user_input():

	"""Get all required parameters from user if not provided via args"""

	global CSV_FILE_PATH, IMDB_COLLECTION_NAME, PLEX_URL, PLEX_TOKEN, MOVIE_LIBRARY_NAME

	# Parse command line arguments

	args = parse_arguments()

	# Check if the user just wants instructions

	if args.instructions:

		display_instructions()

		sys.exit(0)

	# Load saved config

	config = load_config()

	# Get Plex URL

	if args.url:

		PLEX_URL = args.url

	else:

		default_url = config.get('DEFAULT', 'PLEX_URL', fallback='http://127.0.0.1:32400')

		PLEX_URL = input(f"Enter your Plex server URL [Press Enter for {default_url}]: ").strip()

		if not PLEX_URL:

			PLEX_URL = default_url

	# Make sure URL has proper format

	if not PLEX_URL.startswith("http"):

		PLEX_URL = "http://" + PLEX_URL

	# Get Plex token

	if args.token:

		PLEX_TOKEN = args.token

	else:

		default_token = config.get('DEFAULT', 'PLEX_TOKEN', fallback='')

		PLEX_TOKEN = input(f"Enter your Plex authentication token [Press Enter to use saved token]: ").strip()

		if not PLEX_TOKEN:

			PLEX_TOKEN = default_token

		if not PLEX_TOKEN:

			print(f"{Fore.RED}Plex token is required{Style.RESET_ALL}")

			display_instructions()

			sys.exit(1)

	# Get Plex library name

	if args.library:

		MOVIE_LIBRARY_NAME = args.library

	else:

		default_library = config.get('DEFAULT', 'MOVIE_LIBRARY_NAME', fallback='Movies')

		MOVIE_LIBRARY_NAME = input(f"Enter your Plex movie library name [Press Enter for {default_library}]: ").strip()

		if not MOVIE_LIBRARY_NAME:

			MOVIE_LIBRARY_NAME = default_library

	# Save the current configuration (only if not using command line args)

	if not (args.url or args.token or args.library):

		save_config(PLEX_URL, PLEX_TOKEN, MOVIE_LIBRARY_NAME)

	# Test Plex connection

	print(f"\n{Fore.CYAN}Testing Plex connection...{Style.RESET_ALL}")

	if not test_plex_connection():

		retry = input(f"\n{Fore.YELLOW}Would you like to retry with different Plex settings? (y/n): {Style.RESET_ALL}").lower()

		if retry == 'y':

			PLEX_URL = ""

			PLEX_TOKEN = ""

			MOVIE_LIBRARY_NAME = ""

			get_user_input()

			return

		else:

			print(f"{Fore.RED}Exiting due to Plex connection issues.{Style.RESET_ALL}")

			sys.exit(1)

	# Get CSV file path

	if args.csv:

		CSV_FILE_PATH = args.csv

	else:

		while not (CSV_FILE_PATH and os.path.isfile(CSV_FILE_PATH)):

			path = input("\nEnter the path to your IMDB CSV file: ")

			if os.path.isfile(path):

				CSV_FILE_PATH = path

			else:

				print(f"{Fore.RED}File not found: {path}{Style.RESET_ALL}")

				print(f"{Fore.YELLOW}Remember to export a list from IMDB as a CSV file first.{Style.RESET_ALL}")

	# Prompt for new or existing collection

	if args.name:

		IMDB_COLLECTION_NAME = args.name

	else:

		# Get existing collections to offer as options

		existing_collections = get_existing_collections()

		if existing_collections:

			update_existing = input(f"\n{Fore.CYAN}Would you like to update an existing collection? (y/n): {Style.RESET_ALL}").lower() == 'y'

			if update_existing:

				print(f"\n{Fore.CYAN}Available collections:{Style.RESET_ALL}")

				for i, collection in enumerate(existing_collections, 1):

					print(f"{i}. {collection}")

				while True:

					choice = input(f"\nEnter the number of the collection to update (or 'n' for new collection): ")

					if choice.lower() == 'n':

						break

					try:

						index = int(choice) - 1

						if 0 <= index < len(existing_collections):

							IMDB_COLLECTION_NAME = existing_collections[index]

							print(f"{Fore.GREEN}Selected collection: {IMDB_COLLECTION_NAME}{Style.RESET_ALL}")

							return

						else:

							print(f"{Fore.RED}Invalid selection. Please try again.{Style.RESET_ALL}")

					except ValueError:

						print(f"{Fore.RED}Please enter a valid number or 'n'.{Style.RESET_ALL}")

		# Ask for new collection name

		while not IMDB_COLLECTION_NAME:

			name = input("\nEnter the name for your new collection: ")

			if name.strip():

				IMDB_COLLECTION_NAME = name.strip()

			else:

				print(f"{Fore.RED}Collection name cannot be empty{Style.RESET_ALL}")

if __name__ == "__main__":

	print(f"{Fore.CYAN}=== Plex Collection Creator ==={Style.RESET_ALL}")

	# Get user input or command line arguments

	get_user_input()

	# Run the script

	run_imdb_list()

	print(f"\n{Fore.CYAN}Processing complete. Check the log and report files for details.{Style.RESET_ALL}")