"""
Redis backend implementation for InferenceKit.
Supports distributed caching using Redis with vector similarity search capabilities.
"""

import asyncio
import json
import pickle
from typing import Optional, Dict, Any, List, Tuple, Union
from datetime import datetime, timedelta
from uuid import UUID, uuid4
from functools import wraps

import redis
import numpy as np
from redis.exceptions import ConnectionError as RedisConnectionError
from redis.exceptions import TimeoutError as RedisTimeoutError
from redis.exceptions import RedisError

from ..types import (
    Request,
    Response,
    CacheEntry,
    ProviderConfig,
    CacheConfig,
    VectorData,
    SimilarityResult,
)
from ..exceptions import (
    BackendError,
    ConnectionError,
    CacheError,
    SerializationError,
)


class RedisBackend:
    """Redis backend for distributed caching with vector similarity search."""

    def __init__(self, cache_config: CacheConfig, provider_config: ProviderConfig):
        self.cache_config = cache_config
        self.provider_config = provider_config
        self.redis_client: Optional[redis.Redis] = None
        self.vector_index: Optional[redis.client.Pipeline] = None

        # Default Redis connection settings
        self.redis_host = provider_config.base_url or "localhost"
        self.redis_port = 6379
        self.redis_db = 0
        self.redis_password = provider_config.api_key if provider_config.api_key else None
        self.redis_timeout = provider_config.timeout or 10.0
        self.max_retries = provider_config.max_retries or 3

    async def initialize(self) -> None:
        """Initialize the Redis connection."""
        try:
            self.redis_client = redis.Redis(
                host=self.redis_host,
                port=self.redis_port,
                db=self.redis_db,
                password=self.redis_password,
                socket_timeout=self.redis_timeout,
                retry_on_timeout=True,
                max_connections=self.max_retries,
                decode_responses=True,
            )

            # Test connection
            self.redis_client.ping()

            # Initialize vector index if Redis modules are available
            self._initialize_vector_index()

        except RedisConnectionError as e:
            raise ConnectionError(f"Failed to connect to Redis server: {e}")
        except RedisTimeoutError as e:
            raise ConnectionError(f"Redis connection timed out: {e}")
        except RedisError as e:
            raise ConnectionError(f"Redis connection error: {e}")

    def _initialize_vector_index(self) -> None:
        """Initialize vector index using Redis modules if available."""
        try:
            # Check if RedisAI or other vector modules are available
            self.redis_client.execute_command("AI.TENSORSET", "test", "FLOAT", "1", "1", "VALUES", "1.0")
            self.vector_index = self.redis_client
        except RedisError:
            # Vector modules not available, fall back to basic implementation
            self.vector_index = None

    async def shutdown(self) -> None:
        """Shutdown the Redis connection."""
        if self.redis_client:
            self.redis_client.close()
            self.redis_client = None

    async def save_request(self, request: Request) -> str:
        """Save a request to Redis."""
        if not self.redis_client:
            raise ConnectionError("Redis client not initialized")

        try:
            request_id = str(request.id)
            request_data = {
                "id": request_id,
                "model": request.model,
                "provider": request.provider.value,
                "type": request.type.value,
                "prompt": request.prompt,
                "parameters": request.parameters,
                "metadata": request.metadata,
                "created_at": request.created_at.isoformat(),
            }

            # Store request with TTL if configured
            ttl = self.cache_config.ttl_seconds
            if ttl:
                self.redis_client.setex(
                    f"request:{request_id}",
                    ttl,
                    json.dumps(request_data)
                )
            else:
                self.redis_client.set(
                    f"request:{request_id}",
                    json.dumps(request_data)
                )

            # Store metadata for listing/searching
            self.redis_client.hset(
                "metadata:requests",
                request_id,
                json.dumps({
                    "created_at": request.created_at.isoformat(),
                    "type": request.type.value,
                    "provider": request.provider.value,
                })
            )

            return request_id

        except (RedisError, json.JSONDecodeError) as e:
            raise BackendError(f"Failed to save request: {e}")

    async def save_response(self, request_id: str, response: Response) -> str:
        """Save a response to Redis."""
        if not self.redis_client:
            raise ConnectionError("Redis client not initialized")

        try:
            response_data = {
                "request_id": request_id,
                "content": response.content,
                "metadata": response.metadata,
                "created_at": response.created_at.isoformat(),
            }

            # Store response with TTL if configured
            ttl = self.cache_config.ttl_seconds
            if ttl:
                self.redis_client.setex(
                    f"response:{request_id}",
                    ttl,
                    json.dumps(response_data)
                )
            else:
                self.redis_client.set(
                    f"response:{request_id}",
                    json.dumps(response_data)
                )

            # Update metadata to indicate response exists
            self.redis_client.hset(
                "metadata:requests",
                request_id,
                json.dumps({
                    "created_at": datetime.utcnow().isoformat(),
                    "type": "response",
                    "provider": "redis",
                    "response_at": response.created_at.isoformat(),
                })
            )

            return request_id

        except (RedisError, json.JSONDecodeError) as e:
            raise BackendError(f"Failed to save response: {e}")

    async def get_request(self, request_id: str) -> Optional[Request]:
        """Get a request from Redis."""
        if not self.redis_client:
            return None

        try:
            request_data = self.redis_client.get(f"request:{request_id}")
            if not request_data:
                return None

            data = json.loads(request_data)
            return Request(
                id=UUID(data["id"]),
                model=data["model"],
                provider=data["provider"],
                type=data["type"],
                prompt=data["prompt"],
                parameters=data["parameters"],
                metadata=data["metadata"],
                created_at=datetime.fromisoformat(data["created_at"]),
            )

        except (RedisError, json.JSONDecodeError, KeyError, ValueError) as e:
            raise BackendError(f"Failed to get request: {e}")

    async def get_response(self, request_id: str) -> Optional[Response]:
        """Get a response from Redis."""
        if not self.redis_client:
            return None

        try:
            response_data = self.redis_client.get(f"response:{request_id}")
            if not response_data:
                return None

            data = json.loads(response_data)
            return Response(
                request_id=UUID(data["request_id"]),
                content=data["content"],
                metadata=data["metadata"],
                created_at=datetime.fromisoformat(data["created_at"]),
            )

        except (RedisError, json.JSONDecodeError, KeyError, ValueError) as e:
            raise BackendError(f"Failed to get response: {e}")

    async def get_request_response(self, request_id: str) -> Optional[Tuple[Request, Response]]:
        """Get both request and response."""
        request = await self.get_request(request_id)
        response = await self.get_response(request_id)

        if request and response:
            return (request, response)
        return None

    async def list_requests(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """List requests with pagination."""
        if not self.redis_client:
            return []

        try:
            # Get all request IDs sorted by creation time
            request_ids = self.redis_client.hkeys("metadata:requests")
            if not request_ids:
                return []

            # Sort by creation time (Redis doesn't support sorting by hash field values)
            # This is a simplified approach - in production you might want to use sorted sets
            requests = []
            for request_id in request_ids:
                metadata = self.redis_client.hget("metadata:requests", request_id)
                if metadata:
                    metadata_dict = json.loads(metadata)
                    requests.append({
                        "id": request_id.decode(),
                        "created_at": metadata_dict.get("created_at"),
                        "type": metadata_dict.get("type"),
                        "provider": metadata_dict.get("provider"),
                        "has_response": "response_at" in metadata_dict,
                    })

            # Sort by created_at descending
            requests.sort(key=lambda x: x["created_at"], reverse=True)

            # Apply pagination
            paginated_requests = requests[offset:offset + limit]

            return paginated_requests

        except (RedisError, json.JSONDecodeError) as e:
            raise BackendError(f"Failed to list requests: {e}")

    async def search_requests(self, query: str, limit: int = 10) -> List[Dict[str, Any]]:
        """Search requests by prompt content."""
        if not self.redis_client:
            return []

        try:
            # This is a basic implementation - for production use, consider using Redis search module
            request_ids = self.redis_client.hkeys("metadata:requests")
            results = []
            query_lower = query.lower()

            for request_id in request_ids:
                request_data = self.redis_client.get(f"request:{request_id.decode()}")
                if request_data:
                    data = json.loads(request_data)
                    if query_lower in data.get("prompt", "").lower():
                        metadata = self.redis_client.hget("metadata:requests", request_id.decode())
                        if metadata:
                            metadata_dict = json.loads(metadata)
                            results.append({
                                "id": request_id.decode(),
                                "created_at": metadata_dict.get("created_at"),
                                "prompt": data.get("prompt", ""),
                                "has_response": "response_at" in metadata_dict,
                            })

                if len(results) >= limit:
                    break

            return results[:limit]

        except (RedisError, json.JSONDecodeError) as e:
            raise BackendError(f"Failed to search requests: {e}")

    async def count_requests(self) -> int:
        """Get total number of requests."""
        if not self.redis_client:
            return 0

        try:
            return self.redis_client.hlen("metadata:requests")
        except RedisError as e:
            raise BackendError(f"Failed to count requests: {e}")

    async def count_responses(self) -> int:
        """Get total number of responses."""
        if not self.redis_client:
            return 0

        try:
            # Count responses by checking which requests have response data
            request_ids = self.redis_client.hkeys("metadata:requests")
            response_count = 0

            for request_id in request_ids:
                if self.redis_client.exists(f"response:{request_id.decode()}"):
                    response_count += 1

            return response_count

        except RedisError as e:
            raise BackendError(f"Failed to count responses: {e}")

    async def get_stats(self) -> Dict[str, Any]:
        """Get backend statistics."""
        if not self.redis_client:
            return {}

        try:
            total_requests = await self.count_requests()
            total_responses = await self.count_responses()

            return {
                "total_requests": total_requests,
                "total_responses": total_responses,
                "completion_rate": total_responses / total_requests if total_requests > 0 else 0,
                "cache_config": self.cache_config.model_dump(),
                "redis_info": self.redis_client.info(),
                "memory_usage": self.redis_client.memory_usage("metadata:requests") if self.redis_client else 0,
            }

        except RedisError as e:
            raise BackendError(f"Failed to get stats: {e}")

    async def clear(self) -> None:
        """Clear all data from Redis."""
        if not self.redis_client:
            return

        try:
            # Get all keys matching our pattern
            keys = self.redis_client.keys("request:*") + self.redis_client.keys("response:*") + ["metadata:requests"]
            if keys:
                self.redis_client.delete(*keys)
        except RedisError as e:
            raise BackendError(f"Failed to clear Redis: {e}")

    async def delete_request(self, request_id: str) -> bool:
        """Delete a specific request and its response."""
        if not self.redis_client:
            return False

        try:
            # Delete request data
            self.redis_client.delete(f"request:{request_id}")

            # Delete response data
            self.redis_client.delete(f"response:{request_id}")

            # Delete metadata
            self.redis_client.hdel("metadata:requests", request_id)

            return True

        except RedisError as e:
            raise BackendError(f"Failed to delete request: {e}")

    async def delete_response(self, request_id: str) -> bool:
        """Delete a specific response."""
        if not self.redis_client:
            return False

        try:
            # Delete response data
            result = self.redis_client.delete(f"response:{request_id}")
            return result > 0

        except RedisError as e:
            raise BackendError(f"Failed to delete response: {e}")

    async def exists(self, request_id: str) -> bool:
        """Check if a request exists."""
        if not self.redis_client:
            return False

        try:
            return self.redis_client.exists(f"request:{request_id}") > 0
        except RedisError as e:
            raise BackendError(f"Failed to check existence: {e}")

    async def get_size(self) -> int:
        """Get the number of stored requests."""
        return await self.count_requests()

    async def save_vector_data(self, vector_data: VectorData, ttl: Optional[int] = None) -> str:
        """Save vector data for similarity search."""
        if not self.vector_index:
            raise CacheError("Vector index not available - Redis modules not loaded")

        try:
            vector_id = str(uuid4())

            # Store vectors in RedisAI tensor format
            vectors = np.array(vector_data.vectors, dtype=np.float32)

            # Create tensor
            self.vector_index.execute_command(
                "AI.TENSORSET",
                f"vectors:{vector_id}",
                "FLOAT",
                "1",
                str(len(vector_data.vectors)),
                "ROWS",
                str(len(vector_data.vectors[0])),
                "COLS",
                "VALUES",
                *vectors.flatten().tolist()
            )

            # Store metadata
            self.vector_index.hset(
                "metadata:vectors",
                vector_id,
                json.dumps({
                    "count": len(vector_data.vectors),
                    "dimension": len(vector_data.vectors[0]) if vector_data.vectors else 0,
                    "ids": vector_data.ids,
                    "metadata": vector_data.metadata,
                })
            )

            if ttl:
                self.vector_index.expire(f"vectors:{vector_id}", ttl)
                self.vector_index.expire("metadata:vectors", ttl)

            return vector_id

        except RedisError as e:
            raise BackendError(f"Failed to save vector data: {e}")

    async def search_similar_vectors(self, query_vector: List[float],
                                   limit: int = 10,
                                   threshold: float = 0.8) -> SimilarityResult:
        """Search for similar vectors using vector similarity."""
        if not self.vector_index:
            raise CacheError("Vector index not available - Redis modules not loaded")

        try:
            # Create query tensor
            query_tensor = np.array([query_vector], dtype=np.float32)

            # Get all vector IDs
            vector_ids = self.vector_index.hkeys("metadata:vectors")
            results = []

            for vector_id in vector_ids:
                vector_id_str = vector_id.decode()

                # Get vector tensor
                vector_tensor = self.vector_index.execute_command(
                    "AI.TENSORGET",
                    f"vectors:{vector_id_str}",
                    "VALUES"
                )

                # Calculate similarity (dot product for simplicity)
                vector_array = np.array(vector_tensor, dtype=np.float32)
                similarity = np.dot(query_vector, vector_array[0])

                if similarity >= threshold:
                    metadata = self.vector_index.hget("metadata:vectors", vector_id_str)
                    if metadata:
                        metadata_dict = json.loads(metadata)
                        results.append({
                            "id": vector_id_str,
                            "similarity": similarity,
                            "metadata": metadata_dict,
                        })

            # Sort by similarity
            results.sort(key=lambda x: x["similarity"], reverse=True)

            return SimilarityResult(
                request_id=uuid4(),
                similar_entries=[],  # This would be populated with actual CacheEntry objects
                scores=[r["similarity"] for r in results],
                best_match=None,
                best_score=results[0]["similarity"] if results else None,
            )

        except RedisError as e:
            raise BackendError(f"Failed to search similar vectors: {e}")

    async def batch_save_requests(self, requests: List[Request]) -> List[str]:
        """Batch save multiple requests."""
        if not self.redis_client:
            raise ConnectionError("Redis client not initialized")

        try:
            request_ids = []

            with self.redis_client.pipeline() as pipe:
                for request in requests:
                    request_id = str(request.id)
                    request_data = {
                        "id": request_id,
                        "model": request.model,
                        "provider": request.provider.value,
                        "type": request.type.value,
                        "prompt": request.prompt,
                        "parameters": request.parameters,
                        "metadata": request.metadata,
                        "created_at": request.created_at.isoformat(),
                    }

                    pipe.setex(
                        f"request:{request_id}",
                        self.cache_config.ttl_seconds or 3600,
                        json.dumps(request_data)
                    )

                    pipe.hset(
                        "metadata:requests",
                        request_id,
                        json.dumps({
                            "created_at": request.created_at.isoformat(),
                            "type": request.type.value,
                            "provider": request.provider.value,
                        })
                    )

                    request_ids.append(request_id)

                pipe.execute()

            return request_ids

        except RedisError as e:
            raise BackendError(f"Failed to batch save requests: {e}")

    async def batch_save_responses(self, responses: List[Tuple[str, Response]]) -> List[str]:
        """Batch save multiple responses."""
        if not self.redis_client:
            raise ConnectionError("Redis client not initialized")

        try:
            with self.redis_client.pipeline() as pipe:
                for request_id, response in responses:
                    response_data = {
                        "request_id": request_id,
                        "content": response.content,
                        "metadata": response.metadata,
                        "created_at": response.created_at.isoformat(),
                    }

                    pipe.setex(
                        f"response:{request_id}",
                        self.cache_config.ttl_seconds or 3600,
                        json.dumps(response_data)
                    )

                    # Update metadata
                    pipe.hset(
                        "metadata:requests",
                        request_id,
                        json.dumps({
                            "created_at": datetime.utcnow().isoformat(),
                            "type": "response",
                            "provider": "redis",
                            "response_at": response.created_at.isoformat(),
                        })
                    )

                pipe.execute()

            return [request_id for request_id, _ in responses]

        except RedisError as e:
            raise BackendError(f"Failed to batch save responses: {e}")

    def health_check(self) -> Dict[str, Any]:
        """Perform health check on Redis connection."""
        try:
            if not self.redis_client:
                return {"status": "disconnected", "error": "Redis client not initialized"}

            # Test connection
            self.redis_client.ping()

            # Get Redis info
            info = self.redis_client.info()

            return {
                "status": "healthy",
                "redis_version": info.get("redis_version"),
                "connected_clients": info.get("connected_clients"),
                "used_memory": info.get("used_memory_human"),
                "keyspace": info.get("keyspace", {}).get("db0", {}).get("keys", 0),
            }

        except RedisError as e:
            return {
                "status": "unhealthy",
                "error": str(e)
            }

    async def get_health(self) -> Dict[str, Any]:
        """Get health status asynchronously."""
        return self.health_check()