import time
import logging
from typing import Callable, Any, Dict

logger = logging.getLogger(__name__)

class MetricsCollector:
    """Collects metrics for RWL execution.
    """

    def __init__(self):
        self.metrics = {}

    def collect(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
        """Executes the given function and collects metrics.

        Args:
            func: The function to execute.
            *args: Positional arguments for the function.
            **kwargs: Keyword arguments for the function.

        Returns:
            The return value of the function.
        """
        start_time = time.time()
        success = False
        retry_count = 0
        result = None
        exception = None

        try:
            result = func(*args, **kwargs)
            success = True
        except Exception as e:
            exception = e
            logger.exception(f"Function {func.__name__} failed: {e}")
            raise  # Re-raise the exception to allow retry logic
        finally:
            end_time = time.time()
            execution_time = end_time - start_time

            self.metrics.setdefault(func.__name__, []).append({
                'execution_time': execution_time,
                'success': success,
                'retry_count': retry_count,  # To be updated by retry mechanism
                'timestamp': time.time(),
                'exception': str(exception) if exception else None, # Capture exception info
                'result': str(result) if success else None # Capture result if success
            })

        return result

    def update_retry_count(self, func_name: str, retry_count: int):
        """Updates the retry count for a specific function.

        Args:
            func_name: The name of the function.
            retry_count: The retry count.
        """
        if func_name in self.metrics and self.metrics[func_name]:
            self.metrics[func_name][-1]['retry_count'] = retry_count
        else:
            logger.warning(f"No metrics found for function {func_name} to update retry count.")

    def add_quality_score(self, func_name: str, quality_score: float):
        """Adds a quality score to the metrics for a specific function.

        Args:
            func_name: The name of the function.
            quality_score: The quality score.
        """
        if func_name in self.metrics and self.metrics[func_name]:
            self.metrics[func_name][-1]['quality_score'] = quality_score
        else:
            logger.warning(f"No metrics found for function {func_name} to add quality score.")

    def get_metrics(self) -> Dict[str, list]:
        """Returns all collected metrics.

        Returns:
            A dictionary containing the collected metrics.
        """
        return self.metrics