"""
Token generation speed measurement.

TokenSpeedometer tracks wall-clock time across a generation run and
reports tokens-per-second (tok/s).  Useful for benchmarking different
model sizes, devices, or sampling strategies.  Call `start` before
the loop, `tick` after each token, and read `tps` for the running
average.
"""

import time

class TokenSpeedometer:
    def __init__(self):
        self.start_time = None
        self.first_token_time = None
        self.end_time = None
        self.token_count = 0

    def start(self):
        """Call this right before the generation loop starts."""
        self.start_time = time.perf_counter()
        return self

    def update(self):
        """Call this inside the loop every time a new token is yielded."""
        if self.token_count == 0:
            self.first_token_time = time.perf_counter()
        self.token_count += 1

    def end(self):
        """Call this after the loop finishes."""
        self.end_time = time.perf_counter()
        self.print_stats()

    def print_stats(self):
        if self.token_count == 0:
            print("No tokens generated.")
            return

        total_duration = self.end_time - self.start_time
        
        # 1. Overall Speed (including prompt processing)
        overall_tps = self.token_count / total_duration

        print(f"\n--- Performance Metrics ---")
        print(f"Total Tokens:    {self.token_count}")
        print(f"Total Time:      {total_duration:.2f} s")
        print(f"Overall Speed:   {overall_tps:.2f} tokens/sec")

        # 2. Decoding Speed (excluding the wait for the first token)
        # This measures purely how fast the model generates text once it starts
        if self.token_count > 1 and self.first_token_time:
            decoding_duration = self.end_time - self.first_token_time
            decoding_tokens = self.token_count - 1
            decoding_tps = decoding_tokens / decoding_duration
            
            ttft = (self.first_token_time - self.start_time) * 1000 # Time To First Token
            print(f"Decoding Speed:  {decoding_tps:.2f} tokens/sec (excluding pre-fill)")
            print(f"Time to 1st Tok: {ttft:.2f} ms")
        print("---------------------------")