import time
from langchain_core.callbacks import BaseCallbackHandler
from src.utils.metrics import metrics_manager

class PerformanceCallback(BaseCallbackHandler):
    def __init__(self):
        self.start_time = None
        self.first_token_time = None
        self.end_time = None
        self.token_count = 0
        
    def on_llm_start(self, serialized, messages, **kwargs):
        """Called when LLM starts processing."""
        self.start_time = time.perf_counter()
        self.first_token_time = None 
        self.token_count = 0
        
    def on_llm_new_token(self, token: str, **kwargs):
        """Called for every new token generated."""
        if self.first_token_time is None:
            self.first_token_time = time.perf_counter()
        self.token_count += 1
        
    def on_llm_end(self, response, **kwargs):
        """Called when LLM finishes generation."""
        self.end_time = time.perf_counter()
        
        # Calculate final stats and push to global manager
        if self.start_time:
            ttft = (self.first_token_time - self.start_time) if self.first_token_time else 0
            
            # Generation time is strictly (First Token -> End), or (Start -> End) if no tokens (non-streaming failover)
            if self.first_token_time and self.end_time:
                gen_time = self.end_time - self.first_token_time
            elif self.end_time:
                gen_time = self.end_time - self.start_time
            else:
                gen_time = 0
                
            metrics_manager.log_llm_metrics(
                ttft_sec=ttft,
                gen_sec=gen_time,
                tokens=self.token_count
            )