"""Encode and reconstruct native eight-stream speech codes."""
from __future__ import annotations
from pathlib import Path
import torch
from safetensors import safe_open
from torch import Tensor, nn
from torch.nn import functional
from transformers import HiggsAudioV2TokenizerConfig, HiggsAudioV2TokenizerModel
from .layout import CODEBOOKS, validate_codes
from .waveforms import Waveform, validate_waveform
PREFIX = 'tied.embedding.modality_embeddings.0.model.'
UNUSED_PRETRAINING_KEY = 'semantic_model.masked_spec_embed'
AUDIO_AXES = 3

class Codec:

    def __init__(self, model: HiggsAudioV2TokenizerModel, device: torch.device) -> None:
        nn.Module.to(model, device=device, dtype=torch.float32)
        self.model = model
        self.model.train(False)
        self.model.requires_grad_(False)
        self.device = device

    @torch.inference_mode()
    def encode(self, waveform: Waveform) -> Tensor:
        validate_waveform(waveform, 24000)
        audio = torch.from_numpy(waveform.copy()).to(self.device).reshape(1, 1, -1)
        audio = functional.pad(audio, (0, max(0, 24000 - audio.shape[-1])))
        output = self.model.encode(audio)
        codes = getattr(output, 'audio_codes', None)
        if not isinstance(codes, Tensor) or codes.ndim != AUDIO_AXES or codes.shape[:2] != (1, CODEBOOKS):
            message = f'unexpected codec encoder output: {type(codes).__name__}'
            raise ValueError(message)
        result = codes[0].transpose(0, 1).long().cpu()
        validate_codes(result)
        return result

    @torch.inference_mode()
    def decode(self, codes: Tensor) -> Tensor:
        validate_codes(codes)
        output = self.model.decode(codes.transpose(0, 1).unsqueeze(0).to(self.device))
        values = getattr(output, 'audio_values', None)
        if not isinstance(values, Tensor) or values.ndim != AUDIO_AXES or values.shape[:2] != (1, 1):
            message = f'unexpected codec decoder output: {type(values).__name__}'
            raise ValueError(message)
        waveform = values[0, 0].float().cpu()
        if not waveform.numel() or not bool(torch.isfinite(waveform).all().item()):
            message = 'codec decoder returned empty/non-finite waveform'
            raise ValueError(message)
        return waveform

def load_codec(checkpoint_path: Path, config_path: Path, device: torch.device) -> tuple[Codec, dict[str, object]]:
    config = HiggsAudioV2TokenizerConfig.from_json_file(config_path)
    with torch.device('meta'):
        model = HiggsAudioV2TokenizerModel(config)
    expected = model.state_dict()
    tensors: dict[str, Tensor] = {}
    dropped: list[str] = []
    with safe_open(checkpoint_path, framework='pt', device='cpu') as checkpoint:
        names = checkpoint.keys()
        for name in names:
            if not name.startswith(PREFIX):
                continue
            key = name.removeprefix(PREFIX)
            if key == UNUSED_PRETRAINING_KEY:
                dropped.append(name)
                continue
            tensor = checkpoint.get_tensor(name)
            if key not in expected or tensor.shape != expected[key].shape:
                message = f'unexpected codec tensor {checkpoint_path}:{name}, shape={tensor.shape}'
                raise ValueError(message)
            tensors[key] = tensor.float()
    model.load_state_dict(tensors, strict=True, assign=True)
    if any((value.is_meta for value in model.buffers())):
        materialized = HiggsAudioV2TokenizerModel(config)
        materialized.load_state_dict(tensors, strict=True)
        model = materialized
    return (Codec(model, device), {'loaded': sorted(tensors), 'dropped': dropped})