from typing import Dict, Any, Optional
from pydantic import BaseModel
from inferencekit.core.types import Request, Response, CacheConfig


class MemoryBackend:
    """In-memory backend for storing requests and responses"""

    def __init__(self, cache_config: CacheConfig):
        self.cache_config = cache_config
        self.store: Dict[str, Dict[str, Any]] = {}
        self.metadata: Dict[str, Dict[str, Any]] = {}

    def save_request(self, request: Request) -> str:
        """Save a request to the backend"""
        request_id = str(request.id)
        self.store[request_id] = request.model_dump()
        self.metadata[request_id] = {
            "created_at": request.created_at.isoformat(),
            "type": request.type,
            "provider": request.provider.value
        }
        return request_id

    def save_response(self, request_id: str, response: Response) -> str:
        """Save a response to the backend"""
        if request_id not in self.store:
            raise ValueError(f"Request {request_id} not found")

        self.store[request_id]["response"] = response.model_dump()
        self.metadata[request_id]["response_at"] = response.created_at.isoformat()
        return request_id

    def get_request(self, request_id: str) -> Optional[Request]:
        """Get a request from the backend"""
        data = self.store.get(request_id)
        if not data or "response" not in data:
            return None

        try:
            return Request(**data)
        except Exception:
            return None

    def get_response(self, request_id: str) -> Optional[Response]:
        """Get a response from the backend"""
        data = self.store.get(request_id)
        if not data or "response" not in data:
            return None

        try:
            return Response(**data["response"])
        except Exception:
            return None

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

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

    def list_requests(self, limit: int = 100, offset: int = 0) -> list[Dict[str, Any]]:
        """List requests with pagination"""
        sorted_keys = sorted(
            self.metadata.keys(),
            key=lambda k: self.metadata[k]["created_at"],
            reverse=True
        )

        paginated_keys = sorted_keys[offset:offset + limit]
        results = []

        for key in paginated_keys:
            metadata = self.metadata[key]
            results.append({
                "id": key,
                "created_at": metadata["created_at"],
                "type": metadata["type"],
                "provider": metadata["provider"],
                "has_response": "response_at" in metadata
            })

        return results

    def search_requests(self, query: str, limit: int = 10) -> list[Dict[str, Any]]:
        """Search requests by prompt content"""
        results = []
        query_lower = query.lower()

        for request_id, data in self.store.items():
            request_data = data
            if query_lower in request_data.get("prompt", "").lower():
                metadata = self.metadata[request_id]
                results.append({
                    "id": request_id,
                    "created_at": metadata["created_at"],
                    "prompt": request_data.get("prompt", ""),
                    "has_response": "response_at" in metadata
                })

            if len(results) >= limit:
                break

        return results[:limit]

    def count_requests(self) -> int:
        """Get total number of requests"""
        return len(self.store)

    def count_responses(self) -> int:
        """Get total number of responses"""
        return sum(1 for data in self.store.values() if "response" in data)

    def get_stats(self) -> Dict[str, Any]:
        """Get backend statistics"""
        total_requests = self.count_requests()
        total_responses = self.count_responses()

        return {
            "total_requests": total_requests,
            "total_responses": total_responses,
            "completion_rate": total_requests / total_responses if total_responses > 0 else 0,
            "cache_config": self.cache_config.model_dump(),
            "memory_usage": self._estimate_memory_usage()
        }

    def _estimate_memory_usage(self) -> int:
        """Estimate memory usage in bytes"""
        size = 0
        for data in self.store.values():
            size += len(str(data).encode('utf-8'))
        return size

    def clear(self) -> None:
        """Clear all data from the backend"""
        self.store.clear()
        self.metadata.clear()

    def delete_request(self, request_id: str) -> bool:
        """Delete a specific request and its response"""
        if request_id in self.store:
            del self.store[request_id]
        if request_id in self.metadata:
            del self.metadata[request_id]
        return True