# tools/metrics.py
import cProfile
import pstats
import io 
from functools import wraps
from collections import defaultdict
import logging

class Metrics:
    def __init__(self, logger=None):
        self.call_count = defaultdict(int)
        self.total_time = defaultdict(float)
        self.logger = logger if logger else logging.getLogger(__name__)
        if not self.logger.handlers:
            self.logger.addHandler(logging.NullHandler())
        self.logger.debug("Metrics instance initialized.")

    def measure(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            self.call_count[func.__name__] += 1
            
            pr = cProfile.Profile()
            pr.enable()
            result = func(*args, **kwargs)
            pr.disable()
            
            ps = pstats.Stats(pr)
            
            func_code = func.__code__
            func_key_tuple = (func_code.co_filename, func_code.co_firstlineno, func_code.co_name)
            
            time_spent_for_func = 0.0
            if func_key_tuple in ps.stats:
                time_spent_for_func = ps.stats[func_key_tuple][3] # [3] is cumulative time (ct)
            else:
                # Fallback: try to find by function name if exact key fails (e.g. due to decorators changing code object details slightly)
                # This is less precise and might pick up other functions if names are not unique across files.
                found_by_name = False
                for key, stat in ps.stats.items():
                    if key[2] == func.__name__: # key[2] is function name
                        time_spent_for_func = stat[3] # cumulative time
                        self.logger.debug(f"Found stats for {func.__name__} by name {key} after primary key failed.")
                        found_by_name = True
                        break
                if not found_by_name:
                    self.logger.warning(
                        f"Could not find exact cProfile stats for {func.__name__} with key {func_key_tuple} or by name. "
                        f"Time for this call will be recorded as 0. This might occur for non-Python functions or due to complex decorators."
                    )
            
            self.total_time[func.__name__] += time_spent_for_func
            self.logger.debug(f"Measured cumulative time for {func.__name__}: {time_spent_for_func:.6f}s")
            
            return result
        return wrapper

    def get_metrics(self):
        metrics_data = {}
        for func_name in self.call_count:
            count = self.call_count[func_name]
            total_t = self.total_time[func_name]
            metrics_data[func_name] = {
                'call_count': count,
                'total_time': round(total_t, 6),
                'average_time': round(total_t / count, 6) if count > 0 else 0
            }
        return metrics_data

    def clear_metrics(self):
        self.call_count.clear()
        self.total_time.clear()
        self.logger.info("Metrics cleared.")

# Global instance for convenience
_metrics_instance_logger = logging.getLogger(__name__ + ".global_instance")
if not _metrics_instance_logger.handlers:
    _metrics_instance_logger.addHandler(logging.NullHandler())
metrics = Metrics(logger=_metrics_instance_logger)