import mlx.core as mx
import mlx.nn as nn
from mlx_lm.models.base import BaseModelArgs
from mlx_lm.utils import load

from subsampling import ConvSubsampling
from modules import ConformerLayer
from attention import RelPositionalEncoding

class FastConformerEncoder(nn.Module):
    def __init__(
        self,
        feat_in: int = 80,
        n_layers: int = 18,
        d_model: int = 512,
        ff_expansion_factor: int = 4,
        n_heads: int = 8,
        conv_kernel_size: int = 9,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.d_model = d_model
        # FastConformer uses 8x downsampling via striding
        self.pre_encode = ConvSubsampling(
            subsampling='dw_striding',
            subsampling_factor=8,
            feat_in=feat_in,
            feat_out=d_model,
            conv_channels=256
        )
        
        self.pos_enc = RelPositionalEncoding(d_model=d_model, max_len=5000)
        
        d_ff = d_model * ff_expansion_factor
        self.layers = [
            ConformerLayer(
                d_model=d_model,
                d_ff=d_ff,
                n_heads=n_heads,
                conv_kernel_size=conv_kernel_size,
                dropout=dropout
            ) for _ in range(n_layers)
        ]

    def __call__(self, x, lengths=None):
        # x: (batch, time, feat_in)
        if lengths is None:
            lengths = mx.array([x.shape[1]] * x.shape[0])
            
        x, lengths = self.pre_encode(x, lengths)
        
        # Add positional encoding
        x, pos_emb = self.pos_enc(x)
        
        for layer in self.layers:
            x = layer(x, pos_emb=pos_emb)
            
        return x, lengths

class CanaryModel(nn.Module):
    """
    Hybrid ASR-LLM Model connecting FastConformer to Qwen via a linear projection.
    """
    def __init__(self, encoder: FastConformerEncoder, llm_dim: int):
        super().__init__()
        self.encoder = encoder
        
        # Audio to Text projection (LoRA or standard linear depending on NeMo config)
        self.audio_encoder_proj = nn.Linear(1024, llm_dim)
        
    def encode_audio(self, audio_features, lengths=None):
        """
        Passes audio features through the encoder and projects them to the LLM embedding space.
        """
        encoded_audio, lengths = self.encoder(audio_features, lengths)
        audio_embeds = self.audio_encoder_proj(encoded_audio)
        return audio_embeds, lengths

    # Note: Text generation logic will be handled externally by feeding 
    # the audio_embeds directly into the mlx_lm Qwen decoder prompt embedding!