import json
import redis 
import os
import pickle
import io
import base64
from typing import Any, Optional

class DjangoModelStub:
    """Stub class to represent Django models when unpickling"""
    def __init__(self, *args, **kwargs):
        if args and isinstance(args[0], dict):
            for key, value in args[0].items():
                setattr(self, key, value)
        
        for key, value in kwargs.items():
            setattr(self, key, value)
            
    def __getstate__(self):
        return self.__dict__
        
    def __setstate__(self, state):
        for key, value in state.items():
            setattr(self, key, value)

class CustomUnpickler(pickle.Unpickler):
    """Custom unpickler that handles Django models without requiring Django"""
    def find_class(self, module, name):
        if module.startswith('django.') or module.startswith('studio_api.'):
            return DjangoModelStub
        return super().find_class(module, name)

def handler(event, context):
    """
    AWS Lambda handler for Redis operations.
    
    Expected event format:
    {
        "operation": "get"|"set_json"|"delete",
        "keys": ["key1", "key2", ...],
        "values": ["value1", "value2", ...],  # Only for 'set_json' operation
        "ttl": 3600,  # Optional, only for 'set_json' operation
    }
    """
    try:
        operation = event.get('operation')
        if not operation:
            return error_response(400, "Missing required parameter: 'operation'")
            
        keys = event.get('keys')
        if not keys:
            return error_response(400, "Missing required parameter: 'keys'")
        
        redis_client = get_redis_client()
        
        if operation == 'get':
            result = handle_get(redis_client, keys)
        elif operation == 'set_json':
            values = event.get('values')
            if not values:
                return error_response(400, "Missing required parameter: 'values' for 'set_json' operation")
            ttl = event.get('ttl')
            result = handle_set_json(redis_client, keys, values, ttl)
        elif operation == 'delete':
            result = handle_delete(redis_client, keys)
        else:
            return error_response(400, f"Invalid operation: {operation}")
        
        return success_response(result)
        
    except Exception as e:
        return error_response(500, f"Error processing Redis operation: {str(e)}")

def get_redis_client():
    redis_url = os.environ.get('REDIS_URL')
    
    if not redis_url:
        raise ValueError("REDIS_URL environment variable is not set")
    
    return redis.Redis.from_url(redis_url)

def handle_get(redis_client, keys):
    def deserialize(value):
        if value is None:
            return None
            
        try:
            unpickler = CustomUnpickler(io.BytesIO(value))
            return unpickler.load()
        except Exception:
            pass
        
        try:
            return json.loads(value.decode('utf-8'))
        except Exception:
            pass
        
        try:
            return value.decode('utf-8')
        except Exception:
            pass
        
        try:
            return base64.b64encode(value).decode('ascii')
        except Exception:
            return value
    if len(keys) == 1:
        value = redis_client.get(keys[0])
        return deserialize(value)
    else:
        # For multiple keys, use pipeline for efficiency
        pipe = redis_client.pipeline()
        for key in keys:
            pipe.get(key)
        values = pipe.execute()
        
        # Deserialize all values
        results = {}
        for i, key in enumerate(keys):
            results[key] = deserialize(values[i])
        return results

def handle_set_json(redis_client, keys, values, ttl=None):
    def serialize(value):
        if isinstance(value, (dict, list)):
            return json.dumps(value)
        return value
 
    if len(keys) == 1:
        success = redis_client.set(keys[0], serialize(values[0]), ex=ttl)
        return success
    else:
        # For multiple keys, use pipeline for efficiency
        pipe = redis_client.pipeline()
        for i, key in enumerate(keys):
            pipe.set(key, serialize(values[i]), ex=ttl)
        results = pipe.execute()
        return dict(zip(keys, results))

def handle_delete(redis_client, keys):
    if len(keys) == 1:
        deleted = redis_client.delete(keys[0])
        return deleted > 0
    else:
        # For multiple keys, use pipeline for efficiency
        pipe = redis_client.pipeline()
        for key in keys:
            pipe.delete(key)
        results = pipe.execute()
        return dict(zip(keys, [count > 0 for count in results]))
    
def convert_to_serializable(obj):
    """
    Convert any object to a JSON-serializable format.
    Recursively handles nested objects, lists, and dictionaries.
    """
    if isinstance(obj, DjangoModelStub):
        result = {}
        for key, value in obj.__dict__.items():
            result[key] = convert_to_serializable(value)
        return result
    elif isinstance(obj, dict):
        return {k: convert_to_serializable(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_to_serializable(item) for item in obj]
    elif isinstance(obj, (str, int, float, bool, type(None))):
        return obj
    else:
        try:
            if hasattr(obj, '__dict__'):
                return convert_to_serializable(obj.__dict__)
            elif hasattr(obj, '__iter__') and not isinstance(obj, (str, bytes)):
                return [convert_to_serializable(item) for item in obj]
            else:
                return str(obj)
        except:
            return str(obj)

def success_response(data):
    serializable_data = convert_to_serializable(data)
    result = {
        'statusCode': 200,
        'body': json.dumps(serializable_data) 
    }
    print(result)
    return result 

def error_response(status_code, message):
    result = {
        'statusCode': status_code,
        'body': json.dumps({'error': message})
    }
    print(result)
    return result