import time
import logging
from aiva.metrics.email_metrics import EmailMetrics  # Import EmailMetrics

logger = logging.getLogger(__name__)

class ThrottleManager:
    def __init__(self, rates=None, metrics: EmailMetrics = None):
        # Default rates (emails per second, per domain)
        self.rates = rates or {
            'default': 10,
            'gmail.com': 5,
            'yahoo.com': 3
        }
        self.last_sent = {}
        self.metrics = metrics or EmailMetrics()  # Instantiate EmailMetrics if not provided

    def should_throttle(self, domain):
        """Checks if sending to the domain should be throttled."""
        now = time.time()
        if domain not in self.last_sent:
            self.last_sent[domain] = 0

        rate = self.rates.get(domain, self.rates['default'])
        time_since_last = now - self.last_sent[domain]

        if time_since_last < (1 / rate):
            self.metrics.record_throttled(domain)
            return True
        return False

    def update_last_sent(self, domain):
        """Updates the last sent time for the given domain."""
        self.last_sent[domain] = time.time()
        self.metrics.record_sent(domain)

    def send_email(self, domain, email_content):
        """Simulates sending an email, applying throttling."""
        if self.should_throttle(domain):
            delay = (1 / self.rates.get(domain, self.rates['default'])) - (time.time() - self.last_sent.get(domain, 0))
            logger.warning(f"Throttling email to {domain}, waiting {delay:.2f} seconds.")
            time.sleep(delay)
        
        start_time = time.time()
        # Simulate sending the email
        print(f"Sending email to {domain}: {email_content}")
        end_time = time.time()
        
        self.update_last_sent(domain)
        self.metrics.record_delay(domain, end_time - start_time)

    def record_error(self, domain):
        self.metrics.record_error(domain)

# Example usage (can be removed in production, for demonstration only)
if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    metrics = EmailMetrics()
    throttle_manager = ThrottleManager(metrics=metrics)

    for i in range(5):
        throttle_manager.send_email('gmail.com', f'Test email {i}')
        throttle_manager.send_email('example.com', f'Test email {i}')
        time.sleep(0.1)

    print(metrics.get_metrics())
