import boto3
from typing import List, Dict, Any, Callable

class DynamoDBClient:
    def __init__(self, region_name: str = "us-east-2"):
        self.client = boto3.client("dynamodb", region_name=region_name)

    def get_client(self):
        return self.client

    def get_table_item(self, table_name: str, key: dict):
        response = self.client.get_item(TableName=table_name, Key=key)
        return response["Item"]

    def put_table_item(self, table_name: str, item: dict):
        self.client.put_item(TableName=table_name, Item=item)

    def write_to_dynamodb_batch(self, table_name: str, data: List[Dict], converter_function: Callable):
        """
        Write data to DynamoDB in batches of 25 items.
        
        Args:
            table_name: The DynamoDB table name
            data: List of raw data items to be converted and written
            converter_function: Function to convert each data item to DynamoDB format
        """
        if not data:
            print("No data provided for batch write")
            return
            
        batch_size = 25  # DynamoDB batch write limit
        total_items = 0
        
        # Convert data and filter out None items
        converted_items = []
        for item in data:
            converted_item = converter_function(item)
            if converted_item:  # Skip None items
                converted_items.append(converted_item)
        
        print(f"Converting {len(data)} items, {len(converted_items)} valid items to write")
        
        # Process in batches
        for i in range(0, len(converted_items), batch_size):
            batch = converted_items[i:i + batch_size]
            
            # Prepare batch write request
            request_items = {
                table_name: [
                    {"PutRequest": {"Item": item}}
                    for item in batch
                ]
            }
            
            try:
                response = self.client.batch_write_item(RequestItems=request_items)
                
                # Handle unprocessed items with retry logic
                retry_count = 0
                max_retries = 3
                
                while response.get("UnprocessedItems") and retry_count < max_retries:
                    retry_count += 1
                    unprocessed_count = sum(len(items) for items in response["UnprocessedItems"].values())
                    print(f"Retrying {unprocessed_count} unprocessed items (attempt {retry_count}/{max_retries})...")
                    
                    # Retry unprocessed items
                    response = self.client.batch_write_item(RequestItems=response["UnprocessedItems"])
                
                # Check if there are still unprocessed items after max retries
                if response.get("UnprocessedItems"):
                    unprocessed_count = sum(len(items) for items in response["UnprocessedItems"].values())
                    print(f"Warning: {unprocessed_count} items remain unprocessed after {max_retries} retries")
                
                total_items += len(batch)
                print(f"Successfully wrote batch of {len(batch)} items to DynamoDB table '{table_name}'")
                
            except Exception as e:
                print(f"Error writing batch to DynamoDB table '{table_name}': {e}")
                raise
        
        print(f"Total items successfully processed: {total_items}")
        return total_items