import random
import logging
from typing import Any, Dict, Optional
from uuid import UUID, uuid4

logger = logging.getLogger(__name__)

# Using Protocol for interface definition for more flexibility (duck typing)
# Alternatively, an Abstract Base Class (ABC) could be used.
from typing import Protocol

from tensorus.metadata.schemas import DataType # For mock data generation


class TensorStorageConnector(Protocol):
    """
    Interface for a Tensor Storage Connector.

    Defines the contract for classes that handle the storage, retrieval,
    and deletion of tensor data.
    """

    def store_tensor(self, tensor_id: UUID, tensor_data: Any) -> bool:
        """
        Stores the given tensor data, associating it with the provided tensor_id.

        Args:
            tensor_id: The unique identifier for the tensor.
            tensor_data: The tensor data to store. The actual format depends on
                         the implementation (e.g., NumPy array, PyTorch tensor, file path).

        Returns:
            True if storage was successful, False otherwise.
        """
        ...

    def retrieve_tensor(self, tensor_id: UUID) -> Optional[Any]:
        """
        Retrieves the tensor data associated with the given tensor_id.

        Args:
            tensor_id: The unique identifier for the tensor.

        Returns:
            The tensor data if found, None otherwise.
        """
        ...

    def delete_tensor(self, tensor_id: UUID) -> bool:
        """
        Deletes the tensor data associated with the given tensor_id.

        Args:
            tensor_id: The unique identifier for the tensor.

        Returns:
            True if deletion was successful or if the tensor didn't exist,
            False if an error occurred during deletion.
        """
        ...

    def get_tensor_details(self, tensor_id: UUID) -> Optional[Dict[str, Any]]:
        """
        Retrieves basic details about the tensor data itself,
        which might be stored alongside or derived from the tensor.

        This is useful for auto-populating parts of the TensorDescriptor
        if the tensor already exists in storage.

        Args:
            tensor_id: The unique identifier for the tensor.

        Returns:
            A dictionary containing details like 'shape', 'data_type', 'byte_size',
            or None if the tensor is not found or details cannot be determined.
            Example: {"shape": [100, 100], "data_type": "float32", "byte_size": 40000}
        """
        ...


class MockTensorStorageConnector:
    """
    Mock implementation of the TensorStorageConnector.

    Simulates tensor storage using an in-memory dictionary.
    This is intended for testing and development purposes.
    """

    def __init__(self):
        self._mock_db: Dict[UUID, Any] = {}
        self._mock_details_db: Dict[UUID, Dict[str, Any]] = {}

    def clear_all_data(self):
        """Clears all data from the mock storage."""
        self._mock_db.clear()
        self._mock_details_db.clear()
        logger.info("MockStorage: All mock tensor data and details cleared.")

    def store_tensor(self, tensor_id: UUID, tensor_data: Any) -> bool:
        """Stores tensor data in the mock database."""
        logger.info("MockStorage: Storing tensor %s", tensor_id)
        self._mock_db[tensor_id] = tensor_data

        # Simulate deriving and storing details if not already present
        if tensor_id not in self._mock_details_db:
            # If tensor_data is a dict with details, use them
            if isinstance(tensor_data, dict) and all(k in tensor_data for k in ['shape', 'data_type', 'byte_size']):
                 self._mock_details_db[tensor_id] = {
                    "shape": tensor_data.get("shape", [random.randint(1,100), random.randint(1,100)]),
                    "data_type": tensor_data.get("data_type", random.choice(list(DataType))),
                    "byte_size": tensor_data.get("byte_size", random.randint(1000, 100000))
                }
            else: # Generate random details for mock purposes
                self._mock_details_db[tensor_id] = {
                    "shape": [random.randint(1, 100) for _ in range(random.randint(1, 4))],
                    "data_type": random.choice(list(DataType)).value,
                    "byte_size": random.randint(1000, 1000000),
                    "mock_info": "Data generated by mock store_tensor"
                }
        logger.debug("MockStorage: Stored details for %s: %s", tensor_id, self._mock_details_db[tensor_id])
        return True

    def retrieve_tensor(self, tensor_id: UUID) -> Optional[Any]:
        """Retrieves tensor data from the mock database."""
        logger.debug("MockStorage: Retrieving tensor %s", tensor_id)
        return self._mock_db.get(tensor_id)

    def delete_tensor(self, tensor_id: UUID) -> bool:
        """Deletes tensor data from the mock database."""
        logger.info("MockStorage: Deleting tensor %s", tensor_id)
        if tensor_id in self._mock_db:
            del self._mock_db[tensor_id]
            if tensor_id in self._mock_details_db:
                del self._mock_details_db[tensor_id]
            return True
        return False # Tensor not found, considered successful for deletion idempotency

    def get_tensor_details(self, tensor_id: UUID) -> Optional[Dict[str, Any]]:
        """
        Retrieves mock tensor details.

        If details were generated during store_tensor, those are returned.
        Otherwise, it can generate some on-the-fly or indicate not found.
        """
        logger.debug("MockStorage: Getting details for tensor %s", tensor_id)
        if tensor_id in self._mock_details_db:
            return self._mock_details_db[tensor_id]

        # Fallback: if tensor exists in _mock_db but not _mock_details_db (e.g. added externally)
        # This part is less likely to be hit with current store_tensor logic
        if tensor_id in self._mock_db:
            # For this mock, let's generate some details if they weren't explicitly stored
            # In a real scenario, this might involve inspecting the actual tensor data
            details = {
                "shape": [random.randint(1, 50), random.randint(1, 50)],
                "data_type": random.choice(list(DataType)).value,
                "byte_size": random.randint(500, 50000),
                "mock_info": "Data generated by mock get_tensor_details (on-the-fly)"
            }
            self._mock_details_db[tensor_id] = details # Cache it
            return details

        return None

# Global instance for easy access in this phase of development
# In a production system, dependency injection would be preferred.
mock_tensor_connector_instance = MockTensorStorageConnector()
