from __future__ import annotations

import os

import torch
from transformers import Gemma4ForConditionalGeneration

from .configuration_lfg3 import LFG3Config
from .parakeet_projector import ParakeetAudioFrontEnd, merge_audio_into_embeds


class LFG3ForConditionalGeneration(Gemma4ForConditionalGeneration):
    config_class = LFG3Config

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        revision = kwargs.get("revision")
        model = super().from_pretrained(
            pretrained_model_name_or_path, *model_args, **kwargs
        )
        model._install_lfg3_audio(pretrained_model_name_or_path, revision=revision)
        return model

    def _install_lfg3_audio(self, name_or_path, revision=None):
        cfg = self.config
        embed = self.get_input_embeddings()
        device, dtype = embed.weight.device, embed.weight.dtype

        frontend = ParakeetAudioFrontEnd(
            parakeet_name=self._resolve_parakeet(name_or_path, cfg),
            hidden=getattr(cfg, "projector_hidden", 4096),
            out_dim=cfg.text_config.hidden_size,
            encoder_dtype=dtype,
        )
        path = self._resolve_repo_file(
            name_or_path, getattr(cfg, "projector_file", "projector_final.pt"),
            revision,
        )
        ckpt = torch.load(path, map_location="cpu")
        state = ckpt.get("state_dict", ckpt)
        frontend.projector.load_state_dict(state, strict=True)

        frontend.to(device).eval()
        for p in frontend.parameters():
            p.requires_grad_(False)
        self.audio_frontend = frontend

    @staticmethod
    def _resolve_parakeet(name_or_path, cfg):
        local = os.path.join(str(name_or_path), "parakeet")
        if os.path.isdir(local):
            return local
        return getattr(cfg, "parakeet_name", "nvidia/parakeet-tdt-0.6b-v3")

    @staticmethod
    def _resolve_repo_file(name_or_path, filename, revision=None):
        local = os.path.join(str(name_or_path), filename)
        if os.path.isfile(local):
            return local
        from huggingface_hub import hf_hub_download
        return hf_hub_download(str(name_or_path), filename, revision=revision)

    @staticmethod
    def _is_prefill(past_key_values) -> bool:
        if past_key_values is None:
            return True
        get_len = getattr(past_key_values, "get_seq_length", None)
        return get_len() == 0 if callable(get_len) else not past_key_values

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        inputs_embeds=None,
        input_features=None,
        valid_frames=None,
        encoder_attention_mask=None,
        past_key_values=None,
        labels=None,
        use_cache=None,
        logits_to_keep=0,
        **kwargs,
    ):
        if (input_features is not None and inputs_embeds is None
                and input_ids is not None and self._is_prefill(past_key_values)):
            inputs_embeds = merge_audio_into_embeds(
                self,
                self.audio_frontend,
                input_ids,
                input_features.to(self.audio_frontend.encoder_dtype),
                valid_frames,
                self.config.audio_token_id,
                encoder_attention_mask=encoder_attention_mask,
            )
            input_ids = None
        return super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            past_key_values=past_key_values,
            labels=labels,
            use_cache=use_cache,
            logits_to_keep=logits_to_keep,
            **kwargs,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
                                      **kwargs):
        model_inputs = super().prepare_inputs_for_generation(
            input_ids, past_key_values=past_key_values, **kwargs
        )
        audio_keys = ("input_features", "valid_frames", "encoder_attention_mask")
        if self._is_prefill(past_key_values):
            for k in audio_keys:
                if kwargs.get(k) is not None:
                    model_inputs[k] = kwargs[k]
        else:
            for k in audio_keys:
                model_inputs.pop(k, None)
        return model_inputs