diff --git a/src/database.py b/src/database.py index 1b89401..52eb5bd 100644 --- a/src/database.py +++ b/src/database.py @@ -120,6 +120,46 @@ class Database: except: pass # Column already exists + # Add team_order column if it doesn't exist + try: + await db.execute("ALTER TABLE pets ADD COLUMN team_order INTEGER DEFAULT NULL") + await db.commit() + print("Added team_order column to pets table") + except: + pass # Column already exists + + # Migrate existing active pets to have team_order values + try: + # Find active pets without team_order + cursor = await db.execute(""" + SELECT id, player_id FROM pets + WHERE is_active = TRUE AND team_order IS NULL + ORDER BY player_id, id + """) + pets_to_migrate = await cursor.fetchall() + + if pets_to_migrate: + print(f"Migrating {len(pets_to_migrate)} active pets to have team_order values...") + + # Group pets by player + from collections import defaultdict + pets_by_player = defaultdict(list) + for pet in pets_to_migrate: + pets_by_player[pet[1]].append(pet[0]) + + # Assign team_order values for each player + for player_id, pet_ids in pets_by_player.items(): + for i, pet_id in enumerate(pet_ids[:6]): # Max 6 pets per team + await db.execute(""" + UPDATE pets SET team_order = ? WHERE id = ? + """, (i + 1, pet_id)) + + await db.commit() + print("Migration completed successfully") + except Exception as e: + print(f"Migration warning: {e}") + pass # Don't fail if migration has issues + await db.execute(""" CREATE TABLE IF NOT EXISTS location_spawns ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -408,6 +448,9 @@ class Database: if active_only: query += " AND p.is_active = TRUE" + # Order by team position for active pets, then by id for storage pets + query += " ORDER BY CASE WHEN p.is_active THEN COALESCE(p.team_order, 999) ELSE 999 END ASC, p.id ASC" + cursor = await db.execute(query, params) rows = await cursor.fetchall() return [dict(row) for row in rows] @@ -638,11 +681,16 @@ class Database: if not pet: return {"success": False, "error": f"No inactive pet found named '{pet_identifier}'"} - # Activate the pet - await db.execute("UPDATE pets SET is_active = TRUE WHERE id = ?", (pet["id"],)) + # Get next available team slot + next_slot = await self.get_next_available_team_slot(player_id) + if next_slot is None: + return {"success": False, "error": "Team is full (maximum 6 pets)"} + + # Activate the pet and assign team position + await db.execute("UPDATE pets SET is_active = TRUE, team_order = ? WHERE id = ?", (next_slot, pet["id"])) await db.commit() - return {"success": True, "pet": dict(pet)} + return {"success": True, "pet": dict(pet), "team_position": next_slot} async def deactivate_pet(self, player_id: int, pet_identifier: str) -> Dict: """Deactivate a pet by name or species name. Returns result dict.""" @@ -670,58 +718,122 @@ class Database: if active_count["count"] <= 1: return {"success": False, "error": "You must have at least one active pet!"} - # Deactivate the pet - await db.execute("UPDATE pets SET is_active = FALSE WHERE id = ?", (pet["id"],)) + # Deactivate the pet and clear team order + await db.execute("UPDATE pets SET is_active = FALSE, team_order = NULL WHERE id = ?", (pet["id"],)) await db.commit() return {"success": True, "pet": dict(pet)} - async def swap_pets(self, player_id: int, pet1_identifier: str, pet2_identifier: str) -> Dict: - """Swap the active status of two pets. Returns result dict.""" + # Team Order Methods + async def get_next_available_team_slot(self, player_id: int) -> int: + """Get the next available team slot (1-6)""" async with aiosqlite.connect(self.db_path) as db: - db.row_factory = aiosqlite.Row - - # Find both pets cursor = await db.execute(""" - SELECT p.*, ps.name as species_name - FROM pets p - JOIN pet_species ps ON p.species_id = ps.id - WHERE p.player_id = ? - AND (p.nickname = ? OR ps.name = ?) - LIMIT 1 - """, (player_id, pet1_identifier, pet1_identifier)) - pet1 = await cursor.fetchone() + SELECT team_order FROM pets + WHERE player_id = ? AND is_active = TRUE AND team_order IS NOT NULL + ORDER BY team_order ASC + """, (player_id,)) + used_slots = [row[0] for row in await cursor.fetchall()] + # Find first available slot (1-6) + for slot in range(1, 7): + if slot not in used_slots: + return slot + return None # Team is full + + async def set_pet_team_order(self, player_id: int, pet_id: int, position: int) -> Dict: + """Set a pet's team order position (1-6)""" + if position < 1 or position > 6: + return {"success": False, "error": "Team position must be between 1-6"} + + async with aiosqlite.connect(self.db_path) as db: + # Check if pet belongs to player + cursor = await db.execute("SELECT * FROM pets WHERE id = ? AND player_id = ?", (pet_id, player_id)) + pet = await cursor.fetchone() + if not pet: + return {"success": False, "error": "Pet not found"} + + # Check if position is already taken cursor = await db.execute(""" - SELECT p.*, ps.name as species_name - FROM pets p - JOIN pet_species ps ON p.species_id = ps.id - WHERE p.player_id = ? - AND (p.nickname = ? OR ps.name = ?) - LIMIT 1 - """, (player_id, pet2_identifier, pet2_identifier)) - pet2 = await cursor.fetchone() + SELECT id FROM pets + WHERE player_id = ? AND team_order = ? AND is_active = TRUE AND id != ? + """, (player_id, position, pet_id)) + existing_pet = await cursor.fetchone() - if not pet1: - return {"success": False, "error": f"Pet '{pet1_identifier}' not found"} - if not pet2: - return {"success": False, "error": f"Pet '{pet2_identifier}' not found"} + if existing_pet: + return {"success": False, "error": f"Position {position} is already taken"} - if pet1["id"] == pet2["id"]: - return {"success": False, "error": "Cannot swap a pet with itself"} + # Update pet's team order and make it active + await db.execute(""" + UPDATE pets SET team_order = ?, is_active = TRUE + WHERE id = ? AND player_id = ? + """, (position, pet_id, player_id)) - # Swap their active status - await db.execute("UPDATE pets SET is_active = ? WHERE id = ?", (not pet1["is_active"], pet1["id"])) - await db.execute("UPDATE pets SET is_active = ? WHERE id = ?", (not pet2["is_active"], pet2["id"])) await db.commit() + return {"success": True, "position": position} + + async def reorder_team_positions(self, player_id: int, new_positions: List[Dict]) -> Dict: + """Reorder team positions based on new arrangement""" + async with aiosqlite.connect(self.db_path) as db: + try: + # Validate all positions are 1-6 and no duplicates + positions = [pos["position"] for pos in new_positions] + if len(set(positions)) != len(positions): + return {"success": False, "error": "Duplicate positions detected"} + + for pos_data in new_positions: + position = pos_data["position"] + pet_id = pos_data["pet_id"] + + if position < 1 or position > 6: + return {"success": False, "error": f"Invalid position {position}"} + + # Verify pet belongs to player + cursor = await db.execute("SELECT id FROM pets WHERE id = ? AND player_id = ?", (pet_id, player_id)) + if not await cursor.fetchone(): + return {"success": False, "error": f"Pet {pet_id} not found"} + + # Clear all team orders first + await db.execute("UPDATE pets SET team_order = NULL WHERE player_id = ?", (player_id,)) + + # Set new positions + for pos_data in new_positions: + await db.execute(""" + UPDATE pets SET team_order = ?, is_active = TRUE + WHERE id = ? AND player_id = ? + """, (pos_data["position"], pos_data["pet_id"], player_id)) + + await db.commit() + return {"success": True, "message": "Team order updated successfully"} + + except Exception as e: + await db.rollback() + return {"success": False, "error": str(e)} + + async def remove_from_team_position(self, player_id: int, pet_id: int) -> Dict: + """Remove a pet from team (set to inactive and clear team_order)""" + async with aiosqlite.connect(self.db_path) as db: + # Check if pet belongs to player + cursor = await db.execute("SELECT * FROM pets WHERE id = ? AND player_id = ?", (pet_id, player_id)) + pet = await cursor.fetchone() + if not pet: + return {"success": False, "error": "Pet not found"} - return { - "success": True, - "pet1": dict(pet1), - "pet2": dict(pet2), - "pet1_now": "active" if not pet1["is_active"] else "storage", - "pet2_now": "active" if not pet2["is_active"] else "storage" - } + # Check if this is the only active pet + cursor = await db.execute("SELECT COUNT(*) as count FROM pets WHERE player_id = ? AND is_active = TRUE", (player_id,)) + active_count = await cursor.fetchone() + + if active_count["count"] <= 1: + return {"success": False, "error": "Cannot deactivate your only active pet"} + + # Remove from team + await db.execute(""" + UPDATE pets SET is_active = FALSE, team_order = NULL + WHERE id = ? AND player_id = ? + """, (pet_id, player_id)) + + await db.commit() + return {"success": True, "message": "Pet removed from team"} # Item and Inventory Methods async def add_item_to_inventory(self, player_id: int, item_name: str, quantity: int = 1) -> bool: @@ -873,7 +985,7 @@ class Database: FROM pets p JOIN pet_species ps ON p.species_id = ps.id WHERE p.player_id = ? AND p.is_active = 1 - ORDER BY p.id ASC + ORDER BY p.team_order ASC, p.id ASC """, (player_id,)) rows = await cursor.fetchall() return [dict(row) for row in rows] @@ -1601,12 +1713,18 @@ class Database: # Begin transaction await db.execute("BEGIN TRANSACTION") - # Update pet active status based on new team - for pet_id, is_active in team_changes.items(): - await db.execute(""" - UPDATE pets SET is_active = ? - WHERE id = ? AND player_id = ? - """, (is_active, int(pet_id), player_id)) + # Update pet active status and team_order based on new team + for pet_id, position in team_changes.items(): + if position: # If position is a number (1-6), pet is active + await db.execute(""" + UPDATE pets SET is_active = TRUE, team_order = ? + WHERE id = ? AND player_id = ? + """, (position, int(pet_id), player_id)) + else: # If position is False, pet is inactive + await db.execute(""" + UPDATE pets SET is_active = FALSE, team_order = NULL + WHERE id = ? AND player_id = ? + """, (int(pet_id), player_id)) # Mark any pending change as verified await db.execute(""" @@ -1713,18 +1831,28 @@ class Database: # Get current pet states cursor = await db.execute(""" - SELECT id, is_active FROM pets WHERE player_id = ? + SELECT id, is_active, team_order FROM pets WHERE player_id = ? """, (player_id,)) - current_pets = {str(row["id"]): bool(row["is_active"]) for row in await cursor.fetchall()} + current_pets = {str(row["id"]): row["team_order"] if row["is_active"] else False for row in await cursor.fetchall()} # Apply proposed changes to current state new_state = current_pets.copy() - for pet_id, new_active_state in proposed_changes.items(): + for pet_id, new_position in proposed_changes.items(): if pet_id in new_state: - new_state[pet_id] = new_active_state + new_state[pet_id] = new_position - # Count active pets in new state - active_count = sum(1 for is_active in new_state.values() if is_active) + # Count active pets and validate positions + active_positions = [pos for pos in new_state.values() if pos] + active_count = len(active_positions) + + # Check for valid positions (1-6) + for pos in active_positions: + if not isinstance(pos, int) or pos < 1 or pos > 6: + return {"valid": False, "error": f"Invalid team position: {pos}"} + + # Check for duplicate positions + if len(active_positions) != len(set(active_positions)): + return {"valid": False, "error": "Duplicate team positions detected"} # Validate constraints if active_count < 1: