import time
import torch
import tiktoken
import torch.nn.functional as F
from .models.model import GPT, GPTConfig

def generate_text_stream(prompt, model_path, max_new_tokens=50, temperature=0.5, top_k=40, repetition_penalty=1.0, stop_on_eot=False, num_probs=0, device='cpu'):
    # load checkpoint safely onto the requested hardware device
    checkpoint = torch.load(model_path, map_location=device, weights_only=True)
    
    # parse the internal state dict regardless of how the training script packaged it
    if 'model_state_dict' in checkpoint:
        state_dict = checkpoint['model_state_dict']
    elif 'model' in checkpoint:
        state_dict = checkpoint['model']
    else:
        state_dict = checkpoint
        
    # strip out any compilation prefixes leftover from pytorch two point zero optimizations
    unwanted_prefix = '_orig_mod.'
    for k, v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
            
    # initialize config and restore any custom saved hyper parameters
    config = GPTConfig()
    if 'config' in checkpoint:
        for k, v in checkpoint['config'].items():
            setattr(config, k, v)
            
    # dynamically read the vocab size from the checkpoint embedding weights
    # this gracefully handles our padding trick without hardcoding numbers
    vocab_size = state_dict['transformer.wte.weight'].shape[0]
    config.vocab_size = vocab_size
    
    # map states to the architecture and freeze gradients
    model = GPT(config)
    model.load_state_dict(state_dict)
    model.eval()
    model.to(device)
    
    # setup tokenizer and define structural boundary tokens
    enc = tiktoken.get_encoding('gpt2')
    eot_token = 50256
    input_tokens = enc.encode(prompt)
    idx = torch.tensor(input_tokens, dtype=torch.long, device=device).unsqueeze(0)
    
    # cache initial input length to calculate exact generation volume later
    input_len = idx.size(1)
    
    if device == 'cuda':
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    
    generated_text_so_far = ""
    past_key_values = None
    
    # start the autoregressive decoding loop
    with torch.no_grad():
        with torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.float16):
            for i in range(max_new_tokens):
                # crop context to max block size if we exceed the physical attention window
                if past_key_values is not None:
                    idx_cond = idx[:, -1:]
                else:
                    idx_cond = idx[:, -config.block_size:]
                    
                logits, _, past_key_values = model(idx_cond, past_key_values=past_key_values, use_cache=True)
                
                # isolate the raw activations at the final time step
                next_token_logits = logits[:, -1, :]
                
                # apply the repetition penalty to heavily suppress tokens we already generated
                if repetition_penalty != 1.0:
                    for token_id in set(idx[0, input_len:].tolist()):
                        if next_token_logits[0, token_id] > 0:
                            next_token_logits[0, token_id] /= repetition_penalty
                        else:
                            next_token_logits[0, token_id] *= repetition_penalty
                                
                # scale by temperature to control the deterministic boundary
                next_token_logits = next_token_logits / temperature
                
                # truncate long tail probabilities via top k filtering
                if top_k is not None:
                    v, _ = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1)))
                    next_token_logits[next_token_logits < v[:, [-1]]] = -float('Inf')
                    
                probs = F.softmax(next_token_logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1)
                
                top_tokens = []
                if num_probs > 0:
                    top_k_probs, top_k_indices = torch.topk(probs, min(num_probs, probs.size(-1)))
                    top_k_probs_list = top_k_probs[0].tolist()
                    top_k_indices_list = top_k_indices[0].tolist()
                    for p, t_id in zip(top_k_probs_list, top_k_indices_list):
                        try:
                            token_bytes = enc.decode_single_token_bytes(t_id)
                            token_str = token_bytes.decode('utf-8', errors='replace')
                        except Exception:
                            token_str = str(t_id)
                        top_tokens.append({"token": token_str, "prob": p})
                
                idx = torch.cat((idx, idx_next), dim=1)
                
                if device == 'cuda':
                    torch.cuda.synchronize()
                t1 = time.perf_counter()
                
                actual_tokens_generated = i + 1
                inference_time = t1 - t0
                tokens_per_sec = actual_tokens_generated / inference_time if inference_time > 0 else 0
                
                # slice the tensor to ONLY decode the newly generated tokens, ignoring the prompt
                output_tokens = idx[0, input_len:].tolist()
                current_text = enc.decode(output_tokens).replace("<|endoftext|>", "")
                
                new_chunk = current_text[len(generated_text_so_far):]
                generated_text_so_far = current_text
                
                is_done = (i == max_new_tokens - 1) or (stop_on_eot and idx_next.item() == eot_token)
                
                yield {
                    "text": new_chunk,
                    "time": inference_time,
                    "tps": tokens_per_sec,
                    "length": actual_tokens_generated,
                    "top_tokens": top_tokens,
                    "is_done": is_done
                }
                
                # trigger early termination if the model emits an end of text token
                if stop_on_eot and idx_next.item() == eot_token:
                    break