Organize Extracted EQ 1 Equipment

organize_equipment.py

#!/usr/bin/env python3
"""
EverQuest Equipment GLB Organizer
Categorizes and renames 471 GLB equipment files based on EQSage items.json metadata.
"""

import json
import os
import shutil
import re
import argparse
from pathlib import Path
from collections import defaultdict, Counter
from typing import Dict, List, Set, Tuple, Optional


class SlotCategorizer:
	"""Maps slot IDs to equipment categories by analyzing item names."""

	WEAPON_KEYWORDS = {
		'sword': 'Swords',
		'axe': 'Axes',
		'mace': 'Maces',
		'hammer': 'Hammers',
		'dagger': 'Daggers',
		'bow': 'Bows',
		'staff': 'Staves',
		'stave': 'Staves',
		'spear': 'Spears',
		'club': 'Clubs',
		'arrow': 'Ammunition',
		'throwing': 'Ammunition',
	}

	ARMOR_KEYWORDS = {
		'helm': 'Head',
		'cap': 'Head',
		'crown': 'Head',
		'coif': 'Head',
		'mask': 'Head',
		'breastplate': 'Chest',
		'tunic': 'Chest',
		'robe': 'Chest',
		'chestplate': 'Chest',
		'greaves': 'Legs',
		'leggings': 'Legs',
		'pants': 'Legs',
		'boots': 'Feet',
		'shoes': 'Feet',
		'sandals': 'Feet',
		'gauntlets': 'Hands',
		'gloves': 'Hands',
		'vambraces': 'Arms',
		'sleeves': 'Arms',
		'bracer': 'Arms',
		'belt': 'Waist',
		'girdle': 'Waist',
		'sash': 'Waist',
		'cloak': 'Back',
		'cape': 'Back',
		'mantle': 'Back',
		'shield': 'Shields',
		'buckler': 'Shields',
	}

	ACCESSORY_KEYWORDS = {
		'necklace': 'Neck',
		'amulet': 'Neck',
		'choker': 'Neck',
		'ring': 'Fingers',
		'band': 'Fingers',
		'earring': 'Ears',
		'stud': 'Ears',
	}

	def __init__(self, items_data: Dict):
		self.items_data = items_data
		self.slot_to_category: Dict[str, Tuple[str, str]] = {}
		self._build_slot_mapping()

	def _build_slot_mapping(self):
		"""Analyze all items to build slot -> (category, subcategory) mapping."""
		slot_item_names: Dict[str, List[str]] = defaultdict(list)

		# Collect all item names per slot
		for it_key, it_data in self.items_data.items():
			if not it_key.startswith('IT'):
				continue

			# it_data is a dict of slot_id: [item_names]
			for slot_id, items_list in it_data.items():
				if isinstance(items_list, list):
					for item_name in items_list:
						if isinstance(item_name, str):
							slot_item_names[slot_id].append(item_name.lower())

		# Categorize each slot based on item names
		for slot_id, names in slot_item_names.items():
			category, subcategory = self._categorize_slot(names)
			if category:
				self.slot_to_category[slot_id] = (category, subcategory)

	def _categorize_slot(self, item_names: List[str]) -> Tuple[Optional[str], Optional[str]]:
		"""Determine category/subcategory for a slot based on item names."""
		votes: Counter = Counter()

		for name in item_names:
			# Check weapons first
			for keyword, subcat in self.WEAPON_KEYWORDS.items():
				if keyword in name:
					votes[('Weapons', subcat)] += 1
					break
			else:
				# Check armor
				for keyword, subcat in self.ARMOR_KEYWORDS.items():
					if keyword in name:
						votes[('Armor', subcat)] += 1
						break
				else:
					# Check accessories
					for keyword, subcat in self.ACCESSORY_KEYWORDS.items():
						if keyword in name:
							votes[('Accessories', subcat)] += 1
							break

		if votes:
			(category, subcategory), _ = votes.most_common(1)[0]
			return category, subcategory

		return None, None

	def get_category(self, slots: List[str]) -> Tuple[Optional[str], Optional[str]]:
		"""Get category/subcategory for an IT based on its slots."""
		votes: Counter = Counter()

		for slot_id in slots:
			if slot_id in self.slot_to_category:
				votes[self.slot_to_category[slot_id]] += 1

		if votes:
			(category, subcategory), _ = votes.most_common(1)[0]
			return category, subcategory

		return None, None


class FilenameSanitizer:
	"""Converts item names to filesystem-safe filenames."""

	INVALID_CHARS = re.compile(r'[<>:"/\\|?*]')
	WHITESPACE = re.compile(r'\s+')

	@staticmethod
	def sanitize(name: str, max_length: int = 100) -> str:
		"""Convert item name to lowercase filesystem-safe filename."""
		# Remove invalid characters
		name = FilenameSanitizer.INVALID_CHARS.sub('', name)
		# Replace whitespace with underscores
		name = FilenameSanitizer.WHITESPACE.sub('_', name)
		# Lowercase
		name = name.lower()
		# Remove leading/trailing underscores
		name = name.strip('_')
		# Truncate if too long
		if len(name) > max_length:
			name = name[:max_length].rstrip('_')

		return name


class ItemNameSelector:
	"""Selects best representative name from multiple items per IT."""

	SKIP_PREFIXES = ['summoned:', 'class ', 'test']

	@staticmethod
	def select_best_name(item_names: List[str]) -> str:
		"""Choose the best representative name from a list."""
		if not item_names:
			return ''

		# Filter out generic names
		filtered = [
			name for name in item_names
			if not any(name.lower().startswith(prefix) for prefix in ItemNameSelector.SKIP_PREFIXES)
		]

		if not filtered:
			return item_names[0]

		# Return shortest non-generic name
		return min(filtered, key=len)


class FileOrganizer:
	"""Orchestrates the file organization process."""

	def __init__(self, source_dir: Path, target_dir: Path, items_json_path: Path, dry_run: bool = False):
		self.source_dir = source_dir
		self.target_dir = target_dir
		self.items_json_path = items_json_path
		self.dry_run = dry_run

		self.items_data: Dict = {}
		self.categorizer: Optional[SlotCategorizer] = None
		self.file_mapping: List[Dict] = []
		self.used_filenames: Dict[str, int] = defaultdict(int)

	def load_items_json(self):
		"""Load and parse items.json."""
		print(f"Loading items.json from {self.items_json_path}...")
		with open(self.items_json_path, 'r', encoding='utf-8') as f:
			self.items_data = json.load(f)
		print(f"Loaded {len(self.items_data)} entries from items.json")

	def initialize_categorizer(self):
		"""Initialize the slot categorizer."""
		print("Building slot categorization mapping...")
		self.categorizer = SlotCategorizer(self.items_data)
		print(f"Mapped {len(self.categorizer.slot_to_category)} slots to categories")

	def scan_source_files(self) -> List[Path]:
		"""Scan source directory for GLB files."""
		glb_files = list(self.source_dir.glob('*.glb'))
		print(f"Found {len(glb_files)} GLB files in {self.source_dir}")
		return glb_files

	def process_file(self, file_path: Path) -> Dict:
		"""Process a single GLB file and determine its mapping."""
		filename = file_path.name
		it_match = re.match(r'it(\d+)\.glb', filename, re.IGNORECASE)

		# Special files (non-IT names)
		if not it_match:
			return {
				'original': filename,
				'it_number': None,
				'category': 'Special',
				'subcategory': None,
				'new_filename': filename,
				'item_name': filename.replace('.glb', ''),
				'all_item_names': '',
				'slots': ''
			}

		it_number = it_match.group(1)
		it_key = f"IT{it_number}"

		# Check if IT exists in items.json
		if it_key not in self.items_data:
			return {
				'original': filename,
				'it_number': it_number,
				'category': 'Uncategorized',
				'subcategory': None,
				'new_filename': filename,
				'item_name': f'it{it_number}',
				'all_item_names': '',
				'slots': ''
			}

		# Get IT data
		it_data = self.items_data[it_key]

		# Collect all item names and slot IDs
		# it_data is a dict of slot_id: [item_names]
		all_item_names = []
		slot_ids = []
		for slot_id, items_list in it_data.items():
			slot_ids.append(slot_id)
			if isinstance(items_list, list):
				all_item_names.extend(items_list)

		# Categorize
		category, subcategory = self.categorizer.get_category(slot_ids)

		if not category:
			category = 'Uncategorized'
			subcategory = None

		# Select best name
		best_name = ItemNameSelector.select_best_name(all_item_names)
		if not best_name:
			best_name = f'item_it{it_number}'

		# Sanitize filename
		sanitized = FilenameSanitizer.sanitize(best_name)
		if not sanitized:
			sanitized = f'item_it{it_number}'

		# Handle collisions
		new_filename = f"{sanitized}.glb"
		full_path = f"{category}/{subcategory or ''}/{new_filename}".strip('/')

		if full_path in self.used_filenames:
			new_filename = f"{sanitized}_it{it_number}.glb"

		self.used_filenames[f"{category}/{subcategory or ''}/{new_filename}".strip('/')] += 1

		return {
			'original': filename,
			'it_number': it_number,
			'category': category,
			'subcategory': subcategory,
			'new_filename': new_filename,
			'item_name': best_name,
			'all_item_names': '; '.join(all_item_names),
			'slots': ', '.join(slot_ids)
		}

	def create_directory_structure(self):
		"""Create target directory structure."""
		categories = {
			'Weapons': ['Swords', 'Axes', 'Maces', 'Hammers', 'Daggers', 'Bows', 'Staves', 'Spears', 'Clubs', 'Ammunition'],
			'Armor': ['Head', 'Chest', 'Legs', 'Feet', 'Hands', 'Arms', 'Waist', 'Back', 'Shields'],
			'Accessories': ['Neck', 'Fingers', 'Ears'],
			'Uncategorized': [],
			'Special': []
		}

		for category, subcategories in categories.items():
			cat_path = self.target_dir / category
			if not self.dry_run:
				cat_path.mkdir(parents=True, exist_ok=True)

			for subcat in subcategories:
				subcat_path = cat_path / subcat
				if not self.dry_run:
					subcat_path.mkdir(parents=True, exist_ok=True)

	def copy_files(self):
		"""Copy files to organized structure."""
		print(f"\n{'[DRY RUN] ' if self.dry_run else ''}Copying files...")

		for mapping in self.file_mapping:
			source_path = self.source_dir / mapping['original']

			# Build target path
			if mapping['subcategory']:
				target_path = self.target_dir / mapping['category'] / mapping['subcategory'] / mapping['new_filename']
			else:
				target_path = self.target_dir / mapping['category'] / mapping['new_filename']

			if self.dry_run:
				print(f"  {mapping['original']} -> {target_path.relative_to(self.target_dir)}")
			else:
				target_path.parent.mkdir(parents=True, exist_ok=True)
				shutil.copy2(source_path, target_path)

	def organize(self):
		"""Main organization workflow."""
		# Load data
		self.load_items_json()
		self.initialize_categorizer()

		# Scan and process files
		glb_files = self.scan_source_files()

		print("\nProcessing files...")
		for file_path in glb_files:
			mapping = self.process_file(file_path)
			self.file_mapping.append(mapping)

		# Create directory structure
		print(f"\n{'[DRY RUN] ' if self.dry_run else ''}Creating directory structure...")
		self.create_directory_structure()

		# Copy files
		self.copy_files()

		print(f"\n{'[DRY RUN] ' if self.dry_run else ''}Organization complete!")


class ManifestGenerator:
	"""Generates documentation files."""

	def __init__(self, target_dir: Path, file_mapping: List[Dict]):
		self.target_dir = target_dir
		self.file_mapping = file_mapping

	def generate_manifest_csv(self):
		"""Generate MANIFEST.csv."""
		csv_path = self.target_dir / 'MANIFEST.csv'

		with open(csv_path, 'w', encoding='utf-8') as f:
			f.write('original_filename,it_number,category,subcategory,new_filename,primary_item_name,all_item_names,slots\n')

			for mapping in sorted(self.file_mapping, key=lambda x: x['original']):
				f.write(f'"{mapping["original"]}",'
				        f'"{mapping["it_number"] or ""}",'
				        f'"{mapping["category"]}",'
				        f'"{mapping["subcategory"] or ""}",'
				        f'"{mapping["new_filename"]}",'
				        f'"{mapping["item_name"]}",'
				        f'"{mapping["all_item_names"]}",'
				        f'"{mapping["slots"]}"\n')

		print(f"Generated {csv_path}")

	def generate_stats(self):
		"""Generate STATS.txt."""
		stats_path = self.target_dir / 'STATS.txt'

		# Count by category
		category_counts: Counter = Counter()
		subcategory_counts: Dict[str, Counter] = defaultdict(Counter)

		for mapping in self.file_mapping:
			category_counts[mapping['category']] += 1
			if mapping['subcategory']:
				subcategory_counts[mapping['category']][mapping['subcategory']] += 1

		with open(stats_path, 'w', encoding='utf-8') as f:
			f.write('EverQuest Equipment GLB Organization Statistics\n')
			f.write('=' * 60 + '\n\n')
			f.write(f'Total Files: {len(self.file_mapping)}\n\n')

			for category in sorted(category_counts.keys()):
				count = category_counts[category]
				f.write(f'{category}: {count} files\n')

				if category in subcategory_counts:
					for subcat, subcount in sorted(subcategory_counts[category].items()):
						f.write(f'  {subcat}: {subcount}\n')
				f.write('\n')

			# List uncategorized files
			uncategorized = [m for m in self.file_mapping if m['category'] == 'Uncategorized']
			if uncategorized:
				f.write('\nUncategorized Files:\n')
				f.write('-' * 60 + '\n')
				for mapping in sorted(uncategorized, key=lambda x: x['original']):
					f.write(f'  {mapping["original"]}\n')

		print(f"Generated {stats_path}")

	def generate_readme(self):
		"""Generate README.md."""
		readme_path = self.target_dir / 'README.md'

		with open(readme_path, 'w', encoding='utf-8') as f:
			f.write('# EverQuest Equipment GLB Files\n\n')
			f.write('This directory contains 471 GLB equipment models from EverQuest, organized by category.\n\n')
			f.write('## Organization\n\n')
			f.write('Files are organized into the following structure:\n\n')
			f.write('- **Weapons/**: Swords, Axes, Maces, Hammers, Daggers, Bows, Staves, Spears, Clubs, Ammunition\n')
			f.write('- **Armor/**: Head, Chest, Legs, Feet, Hands, Arms, Waist, Back, Shields\n')
			f.write('- **Accessories/**: Neck, Fingers, Ears\n')
			f.write('- **Uncategorized/**: Files without items.json entries\n')
			f.write('- **Special/**: Non-IT named files (bbboard, cube, ladders, etc.)\n\n')
			f.write('## Files\n\n')
			f.write('- **MANIFEST.csv**: Complete mapping of original to new filenames\n')
			f.write('- **STATS.txt**: Organization statistics and counts\n')
			f.write('- **README.md**: This file\n\n')
			f.write('## Examples\n\n')

			# Show some examples
			examples = [
				m for m in self.file_mapping
				if m['category'] not in ['Uncategorized', 'Special'] and m['it_number']
			][:5]

			for mapping in examples:
				subcategory_part = f"/{mapping['subcategory']}" if mapping['subcategory'] else ""
				f.write(f"- `{mapping['original']}` → `{mapping['category']}{subcategory_part}/{mapping['new_filename']}`\n")

			f.write('\n## Source\n\n')
			f.write('Models extracted using LanternExtractor from EverQuest game files.\n')
			f.write('Metadata from EQSage items.json database.\n')

		print(f"Generated {readme_path}")

	def generate_all(self):
		"""Generate all documentation files."""
		self.generate_manifest_csv()
		self.generate_stats()
		self.generate_readme()


def main():
	parser = argparse.ArgumentParser(description='Organize EverQuest equipment GLB files')
	parser.add_argument('--source', type=str, required=True, help='Source directory containing GLB files')
	parser.add_argument('--target', type=str, default='Equipment', help='Target directory for organized files')
	parser.add_argument('--items-json', type=str, required=True, help='Path to items.json')
	parser.add_argument('--dry-run', action='store_true', help='Preview actions without executing')

	args = parser.parse_args()

	source_dir = Path(args.source).resolve()
	target_dir = Path(args.target).resolve()
	items_json_path = Path(args.items_json).resolve()

	# Validate paths
	if not source_dir.exists():
		print(f"Error: Source directory does not exist: {source_dir}")
		return 1

	if not items_json_path.exists():
		print(f"Error: items.json not found: {items_json_path}")
		return 1

	if not args.dry_run and target_dir.exists():
		response = input(f"Target directory {target_dir} already exists. Continue? [y/N]: ")
		if response.lower() != 'y':
			print("Aborted.")
			return 0

	# Run organization
	organizer = FileOrganizer(source_dir, target_dir, items_json_path, args.dry_run)
	organizer.organize()

	# Generate documentation
	if not args.dry_run:
		print("\nGenerating documentation...")
		manifest_gen = ManifestGenerator(target_dir, organizer.file_mapping)
		manifest_gen.generate_all()

		print(f"\nAll done! Organized files are in: {target_dir}")
	else:
		print("\nDry run complete. Use without --dry-run to execute.")

	return 0


if __name__ == '__main__':
	exit(main())