#!/usr/bin/env python3 """ Team Management Service for PetBot Handles team swapping, individual team editing, and team selection hub functionality. """ import asyncio import json from typing import Dict, List, Optional class TeamManagementService: """Service for managing player teams and team swapping operations.""" def __init__(self, database, pin_service): self.database = database self.pin_service = pin_service async def get_team_overview(self, player_id: int) -> Dict: """Get overview of all teams for a player.""" try: # Get active team active_team = await self.database.get_active_team(player_id) # Get saved team configurations team_configs = await self.database.get_player_team_configurations(player_id) # Structure the data teams = { "active": { "pets": active_team, "count": len(active_team), "is_active": True } } # Add saved configurations for i in range(1, 4): config = next((c for c in team_configs if c.get("slot") == i), None) if config: # team_data is already parsed by get_player_team_configurations team_data = config["team_data"] if config["team_data"] else {} teams[f"slot_{i}"] = { "name": config.get("name", f"Team {i}"), "pets": team_data, "count": len(team_data), "last_updated": config.get("updated_at"), "is_active": False } else: teams[f"slot_{i}"] = { "name": f"Team {i}", "pets": {}, "count": 0, "last_updated": None, "is_active": False } return {"success": True, "teams": teams} except Exception as e: return {"success": False, "error": f"Failed to get team overview: {str(e)}"} async def request_team_swap(self, player_id: int, nickname: str, source_slot: int) -> Dict: """Request to swap a saved team configuration to active team.""" try: # Validate slot if source_slot < 1 or source_slot > 3: return {"success": False, "error": "Invalid team slot. Must be 1, 2, or 3"} # Get the source team configuration config = await self.database.load_team_configuration(player_id, source_slot) if not config: return {"success": False, "error": f"No team configuration found in slot {source_slot}"} # Parse team data team_data = json.loads(config["team_data"]) if not team_data: return {"success": False, "error": f"Team {source_slot} is empty"} # Request PIN verification for team swap pin_request = await self.pin_service.request_verification( player_id=player_id, nickname=nickname, action_type="team_swap", action_data={ "source_slot": source_slot, "team_data": team_data, "config_name": config["config_name"] }, expiration_minutes=10 ) if pin_request["success"]: return { "success": True, "message": f"PIN sent to confirm swapping {config['config_name']} to active team", "source_slot": source_slot, "team_name": config["config_name"], "expires_in_minutes": pin_request["expires_in_minutes"] } else: return pin_request except Exception as e: return {"success": False, "error": f"Failed to request team swap: {str(e)}"} async def verify_team_swap(self, player_id: int, pin_code: str) -> Dict: """Verify PIN and execute team swap.""" try: # Define team swap callback async def apply_team_swap_callback(player_id, action_data): """Apply the team swap operation.""" source_slot = action_data["source_slot"] team_data = action_data["team_data"] config_name = action_data["config_name"] # Get current active team before swapping current_active = await self.database.get_active_team(player_id) # Apply the saved team as active team result = await self.database.apply_team_configuration(player_id, source_slot) if result["success"]: return { "success": True, "message": f"Successfully applied {config_name} as active team", "source_slot": source_slot, "pets_applied": len(team_data), "previous_active_count": len(current_active) } else: raise Exception(f"Failed to apply team configuration: {result.get('error', 'Unknown error')}") # Verify PIN and execute swap result = await self.pin_service.verify_and_execute( player_id=player_id, pin_code=pin_code, action_type="team_swap", action_callback=apply_team_swap_callback ) return result except Exception as e: return {"success": False, "error": f"Team swap verification failed: {str(e)}"} async def get_individual_team_data(self, player_id: int, team_identifier: str) -> Dict: """Get data for editing an individual team.""" try: if team_identifier == "active": # Get active team active_pets = await self.database.get_active_team(player_id) return { "success": True, "team_type": "active", "team_name": "Active Team", "team_data": active_pets, "is_active_team": True } else: # Get saved team configuration try: slot = int(team_identifier) if slot < 1 or slot > 3: return {"success": False, "error": "Invalid team slot"} except ValueError: return {"success": False, "error": "Invalid team identifier"} config = await self.database.load_team_configuration(player_id, slot) if config: # Parse team_data - it should be a JSON string containing list of pets try: team_pets = json.loads(config["team_data"]) if config["team_data"] else [] # Ensure team_pets is a list (new format) if isinstance(team_pets, list): pets_data = team_pets else: # Handle old format - convert dict to list pets_data = [] if isinstance(team_pets, dict): for position, pet_info in team_pets.items(): if pet_info and 'pet_id' in pet_info: # This is old format, we'll need to get full pet data pet_data = await self._get_full_pet_data(player_id, pet_info['pet_id']) if pet_data: pet_data['team_order'] = int(position) pets_data.append(pet_data) return { "success": True, "team_type": "saved", "team_slot": slot, "team_name": config["config_name"], "pets": pets_data, # Use 'pets' key expected by webserver "is_active_team": False, "last_updated": config.get("updated_at") } except json.JSONDecodeError: return { "success": True, "team_type": "saved", "team_slot": slot, "team_name": config["config_name"], "pets": [], "is_active_team": False, "last_updated": config.get("updated_at") } else: return { "success": True, "team_type": "saved", "team_slot": slot, "team_name": f"Team {slot}", "pets": [], # Use 'pets' key expected by webserver "is_active_team": False, "last_updated": None } except Exception as e: return {"success": False, "error": f"Failed to get team data: {str(e)}"} async def _get_full_pet_data(self, player_id: int, pet_id: int) -> Optional[Dict]: """Helper method to get full pet data for backward compatibility.""" try: import aiosqlite async with aiosqlite.connect(self.database.db_path) as db: db.row_factory = aiosqlite.Row cursor = await db.execute(""" SELECT p.id, p.nickname, p.level, p.hp, p.max_hp, p.attack, p.defense, p.speed, p.happiness, ps.name as species_name, ps.type1, ps.type2 FROM pets p JOIN pet_species ps ON p.species_id = ps.id WHERE p.id = ? AND p.player_id = ? """, (pet_id, player_id)) row = await cursor.fetchone() return dict(row) if row else None except Exception as e: print(f"Error getting full pet data: {e}") return None async def save_individual_team( self, player_id: int, nickname: str, team_identifier: str, team_data: Dict ) -> Dict: """Save changes to an individual team.""" try: if team_identifier == "active": # Save to active team action_type = "active_team_change" action_data = { "team_type": "active", "team_data": team_data } else: # Save to team configuration slot try: slot = int(team_identifier) if slot < 1 or slot > 3: return {"success": False, "error": "Invalid team slot"} except ValueError: return {"success": False, "error": "Invalid team identifier"} action_type = f"team_{slot}_change" action_data = { "team_type": "saved", "team_slot": slot, "team_data": team_data } # Request PIN verification pin_request = await self.pin_service.request_verification( player_id=player_id, nickname=nickname, action_type=action_type, action_data=action_data, expiration_minutes=10 ) return pin_request except Exception as e: return {"success": False, "error": f"Failed to save individual team: {str(e)}"} async def verify_individual_team_save(self, player_id: int, pin_code: str, team_identifier: str) -> Dict: """Verify PIN and save individual team changes.""" try: if team_identifier == "active": action_type = "active_team_change" else: try: slot = int(team_identifier) action_type = f"team_{slot}_change" except ValueError: return {"success": False, "error": "Invalid team identifier"} # Define save callback async def apply_individual_team_save_callback(player_id, action_data): """Apply individual team save.""" team_type = action_data["team_type"] team_data = action_data["team_data"] if team_type == "active": # Apply to active team changes_applied = await self._apply_to_active_team(player_id, team_data) return { "success": True, "message": "Active team updated successfully", "changes_applied": changes_applied, "team_type": "active" } else: # Save to configuration slot slot = action_data["team_slot"] changes_applied = await self._save_team_configuration(player_id, slot, team_data) return { "success": True, "message": f"Team {slot} configuration saved successfully", "changes_applied": changes_applied, "team_slot": slot, "team_type": "saved" } # Verify PIN and execute save result = await self.pin_service.verify_and_execute( player_id=player_id, pin_code=pin_code, action_type=action_type, action_callback=apply_individual_team_save_callback ) return result except Exception as e: return {"success": False, "error": f"Individual team save verification failed: {str(e)}"} async def _apply_to_active_team(self, player_id: int, team_data: Dict) -> int: """Apply team changes to active pets.""" changes_count = 0 # Deactivate all pets await self.database.execute(""" UPDATE pets SET is_active = FALSE, team_order = NULL WHERE player_id = ? """, (player_id,)) # Activate selected pets for pet_id, position in team_data.items(): if position: await self.database.execute(""" UPDATE pets SET is_active = TRUE, team_order = ? WHERE id = ? AND player_id = ? """, (position, int(pet_id), player_id)) changes_count += 1 return changes_count async def _save_team_configuration(self, player_id: int, slot: int, team_data: Dict) -> int: """Save team as configuration.""" pets_list = [] changes_count = 0 for pet_id, position in team_data.items(): if position: pet_info = await self.database.get_pet_by_id(pet_id) if pet_info and pet_info["player_id"] == player_id: # Create full pet data object in new format pet_dict = { 'id': pet_info['id'], 'nickname': pet_info['nickname'] or pet_info.get('species_name', 'Unknown'), 'level': pet_info['level'], 'hp': pet_info.get('hp', 0), 'max_hp': pet_info.get('max_hp', 0), 'attack': pet_info.get('attack', 0), 'defense': pet_info.get('defense', 0), 'speed': pet_info.get('speed', 0), 'happiness': pet_info.get('happiness', 0), 'species_name': pet_info.get('species_name', 'Unknown'), 'type1': pet_info.get('type1'), 'type2': pet_info.get('type2'), 'team_order': int(position) } pets_list.append(pet_dict) changes_count += 1 # Save configuration in new list format success = await self.database.save_team_configuration( player_id, slot, f'Team {slot}', json.dumps(pets_list) ) return changes_count if success else 0