import torch
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
from transformers.modeling_outputs import Seq2SeqLMOutput, BaseModelOutput
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
from transformers.models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

class LanceASRConfig(PretrainedConfig):
    model_type = "lance_asr"
    is_encoder_decoder = True
    def __init__(self, vocab_size=50257, hidden_size=256, num_layers=4, num_heads=4, num_mel_bins=128, architectures=["LanceASR"], **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.num_mel_bins = num_mel_bins
        self.architectures = architectures
        self.is_encoder_decoder = True
        self.decoder_start_token_id = kwargs.get("decoder_start_token_id", 0)

class LanceASR(PreTrainedModel, GenerationMixin):
    config_class = LanceASRConfig
    _supports_cache_class = False

    def __init__(self, config):
        config.is_encoder_decoder = True
        super().__init__(config)
        self.config = config
        
        # Audio feature extraction (Conv subsampling)
        self.conv1 = nn.Conv1d(config.num_mel_bins, config.hidden_size, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(config.hidden_size, config.hidden_size, kernel_size=3, stride=2, padding=1)
        
        # Text embedding
        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)

        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=config.num_heads, batch_first=True),
            num_layers=config.num_layers
        )
        self.decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model=config.hidden_size, nhead=config.num_heads, batch_first=True),
            num_layers=config.num_layers
        )

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
        self.loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
        
        # Generation config defaults
        self.generation_config.max_new_tokens = 250
        self.generation_config.temperature = 0.8
        self.generation_config.do_sample = True
        self.generation_config.decoder_start_token_id = self.config.decoder_start_token_id

        self.init_weights()

        self.to(torch.bfloat16)

    def get_encoder(self):
        class EncoderWrapper(nn.Module):
            def __init__(self, model):
                super().__init__()
                self.model = model
                self.main_input_name = "input_features"
            def forward(self, input_features, attention_mask=None, **kwargs):
                return self.model.forward_encoder(input_features)
            def __call__(self, *args, **kwargs):
                return self.forward(*args, **kwargs)
        return EncoderWrapper(self)

    def forward_encoder(self, input_features):
        hidden_states = nn.functional.gelu(self.conv1(input_features))
        hidden_states = nn.functional.gelu(self.conv2(hidden_states))
        
        inputs_embeds = hidden_states.permute(0, 2, 1)
        encoder_outputs = self.encoder(inputs_embeds)
        return BaseModelOutput(last_hidden_state=encoder_outputs)

    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
        shifted_labels = labels.new_zeros(labels.shape)
        shifted_labels[..., 1:] = labels[..., :-1].clone()
        shifted_labels[..., 0] = self.config.decoder_start_token_id
        shifted_labels.masked_fill_(shifted_labels == -100, 0)
        return shifted_labels

    def forward(self, input_features=None, decoder_input_ids=None, input_ids=None, encoder_outputs=None, labels=None, return_dict=True, use_cache=False, **kwargs):
        if decoder_input_ids is None and input_ids is not None:
            decoder_input_ids = input_ids
            
        if decoder_input_ids is None and labels is not None:
            decoder_input_ids = self.prepare_decoder_input_ids_from_labels(labels)
            
        if encoder_outputs is None and input_features is not None:
            encoder_outputs = self.forward_encoder(input_features)
            
        memory = encoder_outputs.last_hidden_state if hasattr(encoder_outputs, "last_hidden_state") else (encoder_outputs[0] if isinstance(encoder_outputs, tuple) else encoder_outputs)

        if decoder_input_ids is not None:
            decoder_embeds = self.embedding(decoder_input_ids)
        else:
            raise ValueError("decoder_input_ids must be provided")
            
        seq_len = decoder_embeds.size(1)
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(seq_len).to(device=decoder_embeds.device, dtype=decoder_embeds.dtype)

        decoder_output = self.decoder(
            tgt=decoder_embeds, 
            memory=memory, 
            tgt_mask=tgt_mask, 
            tgt_is_causal=True
        )

        logits = self.lm_head(decoder_output)
        loss = None

        if labels is not None:
            loss = self.loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))

        if return_dict:
            return Seq2SeqLMOutput(loss=loss, logits=logits, encoder_last_hidden_state=memory)

        return (loss, logits) if loss is not None else logits

    def prepare_inputs_for_generation(self, decoder_input_ids, past_key_values=None, attention_mask=None, encoder_outputs=None, **kwargs):
        return {
            "decoder_input_ids": decoder_input_ids,
            "encoder_outputs": encoder_outputs,
        }

    def _reorder_cache(self, past_key_values, beam_idx):
        pass

CONFIG_MAPPING.register("lance_asr", LanceASRConfig)
try:
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.register(LanceASRConfig, LanceASR)
except Exception:
    pass
LanceASRConfig.register_for_auto_class("AutoConfig")
LanceASR.register_for_auto_class("AutoModelForSeq2SeqLM")