database_manager.py

"""
DynamoDB database manager for MCP MUD Game
"""
import boto3
import json
from typing import Dict, List, Optional, Any
from datetime import datetime
from botocore.exceptions import ClientError
from config import DYNAMODB_REGION, DYNAMODB_ENDPOINT, TABLES


class DatabaseManager:
    def __init__(self):
        """Initialize DynamoDB connection"""
        self.dynamodb = boto3.resource(
            'dynamodb',
            region_name=DYNAMODB_REGION,
            endpoint_url=DYNAMODB_ENDPOINT
        )
        self.tables = {}
        self._initialize_tables()
    
    def _initialize_tables(self):
        """Initialize DynamoDB tables"""
        # Players table
        self._create_table_if_not_exists(
            TABLES["players"],
            [
                {'AttributeName': 'player_id', 'KeyType': 'HASH'}
            ],
            [
                {'AttributeName': 'player_id', 'AttributeType': 'S'}
            ]
        )
        
        # Rooms table
        self._create_table_if_not_exists(
            TABLES["rooms"],
            [
                {'AttributeName': 'room_id', 'KeyType': 'HASH'}
            ],
            [
                {'AttributeName': 'room_id', 'AttributeType': 'S'}
            ]
        )
        
        # Game state table
        self._create_table_if_not_exists(
            TABLES["game_state"],
            [
                {'AttributeName': 'game_id', 'KeyType': 'HASH'}
            ],
            [
                {'AttributeName': 'game_id', 'AttributeType': 'S'}
            ]
        )
        
        # Action queue table
        self._create_table_if_not_exists(
            TABLES["action_queue"],
            [
                {'AttributeName': 'action_id', 'KeyType': 'HASH'}
            ],
            [
                {'AttributeName': 'action_id', 'AttributeType': 'S'},
                {'AttributeName': 'timestamp', 'AttributeType': 'N'}
            ],
            global_secondary_indexes=[
                {
                    'IndexName': 'timestamp-index',
                    'KeySchema': [
                        {'AttributeName': 'timestamp', 'KeyType': 'HASH'}
                    ],
                    'Projection': {'ProjectionType': 'ALL'},
                    'ProvisionedThroughput': {'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5}
                }
            ]
        )
        
        # Store table references
        for table_name in TABLES.values():
            self.tables[table_name] = self.dynamodb.Table(table_name)
    
    def _create_table_if_not_exists(self, table_name: str, key_schema: List[Dict], 
                                  attribute_definitions: List[Dict], 
                                  global_secondary_indexes: Optional[List[Dict]] = None):
        """Create DynamoDB table if it doesn't exist"""
        try:
            table = self.dynamodb.Table(table_name)
            table.load()
            print(f"Table {table_name} already exists")
        except ClientError as e:
            if e.response['Error']['Code'] == 'ResourceNotFoundException':
                print(f"Creating table {table_name}")
                
                table_config = {
                    'TableName': table_name,
                    'KeySchema': key_schema,
                    'AttributeDefinitions': attribute_definitions,
                    'ProvisionedThroughput': {
                        'ReadCapacityUnits': 5,
                        'WriteCapacityUnits': 5
                    }
                }
                
                if global_secondary_indexes:
                    table_config['GlobalSecondaryIndexes'] = global_secondary_indexes
                
                table = self.dynamodb.create_table(**table_config)
                table.wait_until_exists()
                print(f"Table {table_name} created successfully")
            else:
                raise e
    
    # Player operations
    def create_player(self, player_data: Dict[str, Any]) -> bool:
        """Create a new player character"""
        try:
            player_data['created_at'] = datetime.utcnow().isoformat()
            player_data['updated_at'] = datetime.utcnow().isoformat()
            
            self.tables[TABLES["players"]].put_item(
                Item=player_data,
                ConditionExpression='attribute_not_exists(player_id)'
            )
            return True
        except ClientError as e:
            if e.response['Error']['Code'] == 'ConditionalCheckFailedException':
                print(f"Player {player_data['player_id']} already exists")
                return False
            raise e
    
    def get_player(self, player_id: str) -> Optional[Dict[str, Any]]:
        """Get player data by ID"""
        try:
            response = self.tables[TABLES["players"]].get_item(
                Key={'player_id': player_id}
            )
            return response.get('Item')
        except ClientError as e:
            print(f"Error getting player {player_id}: {e}")
            return None
    
    def update_player(self, player_id: str, updates: Dict[str, Any]) -> bool:
        """Update player data"""
        try:
            updates['updated_at'] = datetime.utcnow().isoformat()
            
            update_expression = "SET "
            expression_attribute_values = {}
            
            for key, value in updates.items():
                update_expression += f"{key} = :{key}, "
                expression_attribute_values[f":{key}"] = value
            
            update_expression = update_expression.rstrip(", ")
            
            self.tables[TABLES["players"]].update_item(
                Key={'player_id': player_id},
                UpdateExpression=update_expression,
                ExpressionAttributeValues=expression_attribute_values
            )
            return True
        except ClientError as e:
            print(f"Error updating player {player_id}: {e}")
            return False
    
    # Room operations
    def create_room(self, room_data: Dict[str, Any]) -> bool:
        """Create a new room"""
        try:
            room_data['created_at'] = datetime.utcnow().isoformat()
            room_data['updated_at'] = datetime.utcnow().isoformat()
            
            self.tables[TABLES["rooms"]].put_item(Item=room_data)
            return True
        except ClientError as e:
            print(f"Error creating room: {e}")
            return False
    
    def get_room(self, room_id: str) -> Optional[Dict[str, Any]]:
        """Get room data by ID"""
        try:
            response = self.tables[TABLES["rooms"]].get_item(
                Key={'room_id': room_id}
            )
            return response.get('Item')
        except ClientError as e:
            print(f"Error getting room {room_id}: {e}")
            return None
    
    def update_room(self, room_id: str, updates: Dict[str, Any]) -> bool:
        """Update room data"""
        try:
            updates['updated_at'] = datetime.utcnow().isoformat()
            
            update_expression = "SET "
            expression_attribute_values = {}
            
            for key, value in updates.items():
                update_expression += f"{key} = :{key}, "
                expression_attribute_values[f":{key}"] = value
            
            update_expression = update_expression.rstrip(", ")
            
            self.tables[TABLES["rooms"]].update_item(
                Key={'room_id': room_id},
                UpdateExpression=update_expression,
                ExpressionAttributeValues=expression_attribute_values
            )
            return True
        except ClientError as e:
            print(f"Error updating room {room_id}: {e}")
            return False
    
    def get_all_rooms(self) -> List[Dict[str, Any]]:
        """Get all rooms from the database"""
        try:
            response = self.tables[TABLES["rooms"]].scan()
            return response.get('Items', [])
        except ClientError as e:
            print(f"Error getting all rooms: {e}")
            return []
    
    # Game state operations
    def get_game_state(self, game_id: str = "main") -> Optional[Dict[str, Any]]:
        """Get current game state"""
        try:
            response = self.tables[TABLES["game_state"]].get_item(
                Key={'game_id': game_id}
            )
            return response.get('Item')
        except ClientError as e:
            print(f"Error getting game state: {e}")
            return None
    
    def update_game_state(self, updates: Dict[str, Any], game_id: str = "main") -> bool:
        """Update game state"""
        try:
            updates['updated_at'] = datetime.utcnow().isoformat()
            
            update_expression = "SET "
            expression_attribute_values = {}
            
            for key, value in updates.items():
                update_expression += f"{key} = :{key}, "
                expression_attribute_values[f":{key}"] = value
            
            update_expression = update_expression.rstrip(", ")
            
            self.tables[TABLES["game_state"]].update_item(
                Key={'game_id': game_id},
                UpdateExpression=update_expression,
                ExpressionAttributeValues=expression_attribute_values
            )
            return True
        except ClientError as e:
            print(f"Error updating game state: {e}")
            return False
    
    # Action queue operations
    def add_action(self, action_data: Dict[str, Any]) -> bool:
        """Add action to queue"""
        try:
            action_data['timestamp'] = int(datetime.utcnow().timestamp() * 1000)  # milliseconds
            action_data['created_at'] = datetime.utcnow().isoformat()
            
            self.tables[TABLES["action_queue"]].put_item(Item=action_data)
            return True
        except ClientError as e:
            print(f"Error adding action to queue: {e}")
            return False
    
    def get_pending_actions(self, limit: int = 10) -> List[Dict[str, Any]]:
        """Get pending actions from queue"""
        try:
            response = self.tables[TABLES["action_queue"]].scan(
                FilterExpression='attribute_exists(action_id) AND #status = :status',
                ExpressionAttributeNames={'#status': 'status'},
                ExpressionAttributeValues={':status': 'pending'},
                Limit=limit
            )
            return response.get('Items', [])
        except ClientError as e:
            print(f"Error getting pending actions: {e}")
            return []