import sys
import torch
import types
from transformers.cache_utils import DynamicCache


class StreamingDisplay:
    """Real-time streaming display for token generation.
    
    Shows block decoding effect with visual markers and timing info.
    """
    
    def __init__(self, show_block_marker=True, show_speed=True):
        self.prev_text = ""
        self.show_block_marker = show_block_marker
        self.show_speed = show_speed
        self.start_time = None
        self.total_tokens = 0
        self.block_count = 0
        import time
        self.time = time
        self.start_time = time.time()
    
    def update(self, text, tokens_in_block=None):
        """Print only the new characters since last update.
        
        Args:
            text: Full decoded text so far
            tokens_in_block: Number of tokens decoded in this block (for display)
        """
        # Find the new part
        if text.startswith(self.prev_text):
            delta = text[len(self.prev_text):]
        else:
            delta = text
        
        if delta:
            self.block_count += 1
            
            if self.show_block_marker and tokens_in_block and tokens_in_block > 1:
                # Show block marker with color: green for multi-token blocks
                sys.stdout.write(f"\033[92m[+{tokens_in_block}]\033[0m")
            
            sys.stdout.write(delta)
            sys.stdout.flush()
        
        self.prev_text = text
    
    def finish(self, total_tokens=None):
        """Finish streaming with stats."""
        sys.stdout.write("\n")
        
        if self.show_speed and self.start_time:
            elapsed = self.time.time() - self.start_time
            if total_tokens and elapsed > 0:
                tps = total_tokens / elapsed
                sys.stdout.write(f"\033[90m[{total_tokens} tokens in {elapsed:.2f}s = {tps:.1f} tok/s]\033[0m\n")
        
        sys.stdout.flush()

def _build_causal_attention_mask(module, input_ids):
    """Create 4D causal attention mask compatible with eager attention."""
    assert input_ids.dim() == 2, "input_ids should be [batch, seq_len]"
    batch_size, seq_len = input_ids.shape
    device = input_ids.device
    dtype = next(module.parameters()).dtype if hasattr(module, 'parameters') else torch.float32

    causal = torch.tril(torch.ones((seq_len, seq_len), device=device, dtype=torch.bool))
    causal = causal.unsqueeze(0).unsqueeze(0).expand(batch_size, 1, seq_len, seq_len)

    neg_inf = torch.finfo(dtype).min
    mask = torch.zeros((batch_size, 1, seq_len, seq_len), dtype=dtype, device=device)
    mask = mask.masked_fill(~causal, neg_inf)
    return mask


class Fast_dLLM_QwenForCausalLM:

    @torch.no_grad()
    def mdm_sample(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        pixel_values=None,
        image_grid_thw=None,
        mask_id=151665,
        threshold=0.9,
        output_file='decoded_output.txt',
        stop_token=151645,
        return_stats=False,
    ):
        self.model.bd_size = block_size
        tokens_per_step = []  # Track tokens decoded at each step
        total_steps = 0
        
        with open(output_file, 'w') as f:
            f.write('begin mdm sampling with sampling controls\n')


        num_blocks = max_tokens // block_size
        original_input_length = input_ids.shape[1]
        for block_idx in range(num_blocks):
            if stop_token in input_ids[:, original_input_length:]:
                break
            prompt_length = input_ids.shape[1]
            # Initialize x_init with mask_id tokens
            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
            x_init = torch.cat([input_ids, x_init], dim=1)
            x_t = x_init.clone()
            step = 0
            while True:
                mask_idx = (x_t[:, prompt_length:] == mask_id)
                if mask_idx.sum() == 0:
                    break

                if stop_token in x_t[:, prompt_length:]:
                    stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                    if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                        break

                #attention_mask = _build_causal_attention_mask(self, x_t)
                logits = self.forward(
                    input_ids=x_t,
                    #attention_mask=attention_mask,
                    pixel_values=pixel_values,
                    image_grid_thw=image_grid_thw,
                ).logits

                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                x_1 = logits.argmax(dim=-1)[:, prompt_length:]
                probs = torch.softmax(logits, dim=-1)[:, prompt_length:]
                x1_p = torch.gather(probs, dim=-1, index=x_1.unsqueeze(-1)).squeeze(-1)
                
                # Select tokens with confidence greater than threshold
                x1_p = torch.where(mask_idx, x1_p, -torch.inf)
                unmask_idx = (x1_p > threshold)
                if unmask_idx.sum() > 0:
                    num_unmasked = int(unmask_idx.sum())
                    x_t[:, prompt_length:][unmask_idx] = x_1[unmask_idx]
                    tokens_per_step.append(num_unmasked)
                    eos_prob_this_position = None
                    max_eos_prob = None
                else:
                    # Select the token with highest confidence
                    token_position = x1_p.argmax()
                    x_t[:, prompt_length:][0, token_position] = x_1[0, token_position]
                    tokens_per_step.append(1)
                    eos_prob_this_position = probs[0, token_position, stop_token]
                    max_eos_prob = probs[0, :, stop_token].max()
                    
                step += 1
                total_steps += 1
                with open(output_file, 'a') as f:
                    f.write('-' * 50 + '\n')
                    f.write(f'step {step}\n')
                    f.write(f'eos_prob_this_position: {eos_prob_this_position}\n')
                    f.write(f'max_eos_prob this block: {max_eos_prob}\n')
                    f.write('decoded x_t:\n')
                    f.write(tokenizer.decode(x_t[0]))
                    f.write('\n' + '-' * 50 + '\n')

            input_ids = x_t
        # Truncate at stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        
        if return_stats:
            gen_length = input_ids.shape[1] - original_input_length
            stats = {
                "tokens_per_step": tokens_per_step,
                "total_steps": total_steps,
                "gen_length": gen_length,
            }
            return input_ids, stats
        return input_ids


    @torch.no_grad()
    def mdm_sample_cache(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        pixel_values=None,
        image_grid_thw=None,
        mask_id=151665,
        threshold=0.9,
        stop_token=151645,
        return_stats=False,
    ):
        
        self.model.bd_size = block_size
        tokens_per_step = []  # Track tokens decoded at each step
        
        if input_ids.shape[1] > block_size:
            cached_ids = input_ids[:, :(input_ids.shape[1] // block_size * block_size)]
            #attention_mask = _build_causal_attention_mask(self, cached_ids)
            output = self.forward(
                input_ids=cached_ids,
                #attention_mask=attention_mask,
                pixel_values=pixel_values,
                image_grid_thw=image_grid_thw,
                use_cache=True,
                update_kv_cache=True,
            )
            logits, past_key_values = output.logits, output.past_key_values
            if input_ids.shape[1] % block_size == 0:
                next_token = logits[:, -1:, :].argmax(dim=-1)
                input_ids = torch.cat([input_ids, next_token], dim=1)
                tokens_per_step.append(1)  # AR step generates 1 token
        else:
            past_key_values = None

        num_blocks = max_tokens // block_size
        original_input_length = input_ids.shape[1]
        step = 0

        for _ in range(num_blocks):
            if stop_token in input_ids[:, original_input_length:]:
                break
            prompt_length = input_ids.shape[1]
            # Initialize x_init with mask_id tokens
            if prompt_length % block_size != 0:
                x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
                x_init = torch.cat([input_ids, x_init], dim=1)
            else:
                x_init = input_ids
            x_t = x_init.clone()
            while True:
                mask_idx = (x_t[:, -block_size:] == mask_id)
                if mask_idx.sum() == 0:
                    current_ids = x_t[:, -block_size:]
                    # attention_mask = _build_causal_attention_mask(self, current_ids)
                    output = self.forward(
                        input_ids=current_ids,
                        # attention_mask=attention_mask,
                        use_cache=True,
                        past_key_values=past_key_values,
                        update_kv_cache=True,
                        eval_bd_size=block_size,
                    )
                    logits, past_key_values = output.logits, output.past_key_values
                    next_token = logits[:, -1:, :].argmax(dim=-1)
                    x_t = torch.cat([x_t, next_token], dim=1)
                    step += 1
                    tokens_per_step.append(1)  # AR step generates 1 token
                    break

                if stop_token in x_t[:, prompt_length:]:
                    stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                    if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                        break

                current_ids = x_t[:, -block_size:]
                # attention_mask = _build_causal_attention_mask(self, current_ids)
                logits = self.forward(
                    input_ids=current_ids,
                    # attention_mask=attention_mask,
                    use_cache=True,
                    past_key_values=past_key_values,
                    update_kv_cache=False,
                    eval_bd_size=block_size,
                ).logits

                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                x_1 = logits.argmax(dim=-1)
                probs = torch.softmax(logits, dim=-1)
                x1_p = torch.gather(probs, dim=-1, index=x_1.unsqueeze(-1)).squeeze(-1)
                
                # Select tokens with confidence greater than threshold
                x1_p = torch.where(mask_idx, x1_p, -torch.inf)
                unmask_idx = (x1_p > threshold)
                if unmask_idx.sum() > 0:
                    num_unmasked = int(unmask_idx.sum())
                    x_t[:, -block_size:][unmask_idx] = x_1[unmask_idx]
                    tokens_per_step.append(num_unmasked)
                else:
                    # Select the token with highest confidence
                    token_position = x1_p.argmax()
                    x_t[:, -block_size:][0, token_position] = x_1[0, token_position]
                    tokens_per_step.append(1)
                step += 1

            input_ids = x_t
        # Truncate at stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        
        gen_length = input_ids.shape[1] - original_input_length
        print(f'step: {step} gen_length: {gen_length}')
        
        if return_stats:
            stats = {
                "tokens_per_step": tokens_per_step,
                "total_steps": step,
                "gen_length": gen_length,
            }
            return input_ids, stats
        return input_ids



    @torch.no_grad()
    def block_causal_sample_cache(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        pixel_values=None,
        image_grid_thw=None,
        mask_id=151665,
        threshold=1,
        stop_token=151645,
        return_stats=False,
        debug=False,
        force_left_to_right=False,
        output_file=None,
        sub_block_size=None
    ):
        if output_file is None:
            output_file = f'decoded_output_block_size_{block_size}_force_left_to_right_{force_left_to_right}.txt'
        self.model.bd_size = block_size
        original_input_length = input_ids.shape[1]
        tokens_per_step = []  # Track tokens decoded at each step
        
        if debug:
            with open(output_file, 'w') as f:
                f.write('begin block causal sampling with sampling controls\n')
                f.write(f"Generation configs: block_size={block_size}, force_left_to_right={force_left_to_right}, threshold={threshold}, mask_id={mask_id}, stop_token={stop_token}\n")
                f.write('-' * 50 + '\n')
                f.write(f'input_ids: {tokenizer.decode(input_ids[0])}\n')
                f.write('-' * 50 + '\n')

        output = self.forward(
            input_ids=input_ids,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            use_cache=True,
            update_kv_cache=True,
        )
        logits, past_key_values = output.logits, output.past_key_values
        next_token = logits[:, -1:, :].argmax(dim=-1)
        input_ids = torch.cat([input_ids, next_token], dim=1)
        tokens_per_step.append(1)  # AR step generates 1 token
        if debug:
            with open(output_file, 'a') as f:
                f.write('-' * 50 + '\n')
                f.write(f'step 1\n')
                f.write(f'decoded x_t: {tokenizer.decode(input_ids[0, original_input_length:])}\n')
                f.write('-' * 50 + '\n')

        num_blocks = max_tokens // block_size
        step = 1

        for _ in range(num_blocks):
            if stop_token in input_ids[:, original_input_length:]:
                break
            prompt_length = input_ids.shape[1]
            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-1), device=self.device, dtype=torch.long)
            x_init = torch.cat([input_ids, x_init], dim=1)
            x_t = x_init.clone()
            if sub_block_size is not None:
                sub_block_start_idx = 0
                sub_block_end_idx = sub_block_size
            while True:
                mask_idx = (x_t[:, -block_size:] == mask_id)
                if sub_block_size is not None:
                    if mask_idx[:, sub_block_start_idx:sub_block_end_idx].sum() == 0:
                        sub_block_start_idx += sub_block_size
                        sub_block_end_idx += sub_block_size
                    mask_idx[:, sub_block_end_idx:] = False
                if mask_idx.sum() == 0:
                    current_ids = x_t[:, -block_size:]
                    # attention_mask = _build_causal_attention_mask(self, current_ids)
                    output = self.forward(
                        input_ids=current_ids,
                        # attention_mask=attention_mask,
                        use_cache=True,
                        past_key_values=past_key_values,
                        update_kv_cache=True,
                        eval_bd_size=block_size,
                    )
                    logits, past_key_values = output.logits, output.past_key_values
                    next_token = logits[:, -1:, :].argmax(dim=-1)
                    x_t = torch.cat([x_t, next_token], dim=1)
                    step += 1
                    tokens_per_step.append(1)  # AR step generates 1 token
                    if debug:
                        with open(output_file, 'a') as f:
                            f.write('-' * 50 + '\n')
                            f.write(f'step {step}\n')
                            f.write(f'decoded x_t: {tokenizer.decode(x_t[0, original_input_length:])}\n')
                            f.write('-' * 50 + '\n')
                    break

                if stop_token in x_t[:, prompt_length:]:
                    stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                    if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                        break

                current_ids = x_t[:, -block_size:]
                # attention_mask = _build_causal_attention_mask(self, current_ids)
                logits = self.forward(
                    input_ids=current_ids,
                    # attention_mask=attention_mask,
                    use_cache=True,
                    past_key_values=past_key_values,
                    update_kv_cache=False,
                    eval_bd_size=block_size,
                ).logits
                step += 1

                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                x_1 = logits.argmax(dim=-1)
                probs = torch.softmax(logits, dim=-1)
                x1_p = torch.gather(probs, dim=-1, index=x_1.unsqueeze(-1)).squeeze(-1)
                
                # Select tokens with confidence greater than threshold
                x1_p = torch.where(mask_idx, x1_p, -torch.inf)
                if force_left_to_right:
                    # select the mask token to be predicted from the left to the right
                    unmask_idx = torch.zeros_like(mask_idx, dtype=torch.bool)
                    unmask_id = mask_idx.nonzero(as_tuple=False)[0]
                    unmask_idx[unmask_id[0], unmask_id[1]] = True
                else:
                    unmask_idx = (x1_p > threshold)
                if unmask_idx.sum() > 0:
                    num_unmasked = int(unmask_idx.sum())
                    x_t[:, -block_size:][unmask_idx] = x_1[unmask_idx]
                    tokens_per_step.append(num_unmasked)
                else:
                    # Select the token with highest confidence
                    token_position = x1_p.argmax()
                    x_t[:, -block_size:][0, token_position] = x_1[0, token_position]
                    tokens_per_step.append(1)

                if step == 6:
                    pass
                if debug:
                    with open(output_file, 'a') as f:
                        f.write('-' * 50 + '\n')
                        f.write(f'step {step}\n')
                        f.write(f'decoded x_t: {tokenizer.decode(x_t[0, original_input_length:])}\n')
                        f.write('-' * 50 + '\n')

            input_ids = x_t
        # Truncate at stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        
        gen_length = input_ids.shape[1] - original_input_length
        print(f'step: {step} gen_length: {gen_length}')
        
        if return_stats:
            stats = {
                "tokens_per_step": tokens_per_step,
                "total_steps": step,
                "gen_length": gen_length,
            }
            return input_ids, stats
        return input_ids



    @torch.no_grad()
    def speculative_block_causal_sample_cache(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        pixel_values=None,
        image_grid_thw=None,
        mask_id=151665,
        threshold=0,
        stop_token=151645,
        debug=False,
        output_file=None,
        sub_block_size=None,
        return_stats=False,
        stream=False,
        stream_callback=None  # callback(decoded_text, tokens_in_block) for custom streaming
    ):
        if output_file is None:
            output_file = f'decoded_output_speculative_block_size_{block_size}.txt'
        self.model.bd_size = block_size
        original_input_length = input_ids.shape[1]
        tokens_per_step = []
        
        # Initialize streaming display (use callback if provided, else built-in)
        streamer = StreamingDisplay() if (stream and stream_callback is None) else None
        
        if debug:
            with open(output_file, 'w') as f:
                f.write('begin speculative block causal sampling with sampling controls\n')
                f.write('-' * 50 + '\n')
                f.write(f'input_ids: {tokenizer.decode(input_ids[0])}\n')
                f.write('-' * 50 + '\n')

        output = self.forward(
            input_ids=input_ids,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            use_cache=True,
            update_kv_cache=True,
        )
        logits, past_key_values = output.logits, output.past_key_values
        next_token = logits[:, -1:, :].argmax(dim=-1)
        input_ids = torch.cat([input_ids, next_token], dim=1)
        tokens_per_step.append(1)
        if debug:
            with open(output_file, 'a') as f:
                f.write('-' * 50 + '\n')
                f.write(f'step 1\n')
                f.write(f'decoded x_t: {tokenizer.decode(input_ids[0, original_input_length:])}\n')
                f.write('-' * 50 + '\n')
        
        # Stream output (first token from AR)
        if streamer:
            decoded = tokenizer.decode(input_ids[0, original_input_length:], skip_special_tokens=True)
            streamer.update(decoded, tokens_in_block=1)
        elif stream_callback:
            decoded = tokenizer.decode(input_ids[0, original_input_length:], skip_special_tokens=True)
            stream_callback(decoded, 1)

        step = 1

        while True:
            prompt_length = input_ids.shape[1]
            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-1), device=self.device, dtype=torch.long)
            x_init = torch.cat([input_ids, x_init], dim=1)
            x_t = x_init.clone()
            mask_idx = (x_t[:, -block_size:] == mask_id)

            current_ids = x_t[:, -block_size:]
            # attention_mask = _build_causal_attention_mask(self, current_ids)
            logits = self.forward(
                input_ids=current_ids,
                # attention_mask=attention_mask,
                use_cache=True,
                past_key_values=past_key_values,
                update_kv_cache=False,
                eval_bd_size=block_size,
            ).logits
            tokens_per_step.append(0)
            step += 1

            logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
            x_1 = logits.argmax(dim=-1)
            probs = torch.softmax(logits, dim=-1)
            x1_p = torch.gather(probs, dim=-1, index=x_1.unsqueeze(-1)).squeeze(-1)
            
            # Select tokens with confidence greater than threshold
            x1_p = torch.where(mask_idx, x1_p, -torch.inf)
            unmask_idx = (x1_p > threshold)

            if unmask_idx.sum() > 0:
                x_t[:, -block_size:][unmask_idx] = x_1[unmask_idx]
            else:
                assert False, "No mask token to be predicted with threshold 0"

            if debug:
                with open(output_file, 'a') as f:
                    f.write('-' * 50 + '\n')
                    f.write(f'step {step}\n')
                    f.write(f'decoded x_t: {tokenizer.decode(x_t[0, original_input_length:])}\n')
                    f.write('-' * 50 + '\n')


            current_ids = x_t[:, -block_size:]
            output = self.forward(
                input_ids=current_ids,
                use_cache=True,
                past_key_values=past_key_values,
                update_kv_cache=True,
                eval_bd_size=block_size,
            )
            logits, past_key_values = output.logits, output.past_key_values
            ar_block_token = logits.argmax(dim=-1)
            accepted_token_num = 0
            for i in range(ar_block_token.shape[1]-1):
                if (ar_block_token[:, i] == current_ids[:, i+1]).all():
                    accepted_token_num += 1
                else:
                    break
            accepted_token_num += 1
            tokens_per_step.append(accepted_token_num)
            input_ids = torch.cat([input_ids, ar_block_token[:, :accepted_token_num]], dim=1)

            new_past_key_values = []
            for layer_num in range(len(past_key_values)):
                layer_past_key_values = ()
                for kv_idx in range(len(past_key_values[layer_num])):
                    layer_past_key_values += (past_key_values[layer_num][kv_idx][:, :, :input_ids.shape[1]-1, :],)
                new_past_key_values.append(layer_past_key_values)
            past_key_values = DynamicCache(new_past_key_values)
            
            step += 1
            if debug:
                with open(output_file, 'a') as f:
                    f.write('-' * 50 + '\n')
                    f.write(f'step {step}\n')
                    f.write(f'accepted_token_num: {accepted_token_num}\n')
                    f.write(f'decoded x_t: {tokenizer.decode(input_ids[0, original_input_length:])}\n')
                    f.write('-' * 50 + '\n')
            
            # Stream output - show accepted tokens in real-time with block size
            if streamer:
                decoded = tokenizer.decode(input_ids[0, original_input_length:], skip_special_tokens=True)
                streamer.update(decoded, tokens_in_block=accepted_token_num)
            elif stream_callback:
                decoded = tokenizer.decode(input_ids[0, original_input_length:], skip_special_tokens=True)
                stream_callback(decoded, accepted_token_num)
            
            if input_ids.shape[1] - original_input_length > max_tokens:
                break
            if stop_token in input_ids[:, prompt_length:]:
                stop_token_idx = (input_ids[:, prompt_length:] == stop_token).nonzero()[0][1]
                if (input_ids[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                    break

        # Truncate at stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        
        gen_length = input_ids.shape[1] - original_input_length
        
        # Final newline for stream mode with stats
        if streamer:
            streamer.finish(total_tokens=gen_length)
        
        print(f'step: {step} gen_length: {gen_length}')
        
        if return_stats:
            stats = {
                "tokens_per_step": tokens_per_step,
                "total_steps": step,
                "gen_length": gen_length,
            }
            return input_ids, stats
        return input_ids


def setup_model_with_custom_generation(model):
    """
    Set up custom generation functions for the model
    """
    # Add mdm_sample method with visualization
    model.mdm_sample = types.MethodType(Fast_dLLM_QwenForCausalLM.mdm_sample, model)
    return model
