"""IndicF5-Hinglish -- self-contained loader."""

import json
import os

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import PretrainedConfig, PreTrainedModel

from .cfm import CFM
from .dit import DiT
from .utils import get_tokenizer

SAMPLE_RATE = 24_000
HOP_LENGTH = 256
TARGET_RMS = 0.1

PREFIXES = ("ema_model._orig_mod.", "ema_model.", "_orig_mod.", "model.", "module.")
IGNORE = {"initted", "step",
          "mel_spec.mel_stft.mel_scale.fb",
          "mel_spec.mel_stft.spectrogram.window"}


class IndicF5HinglishConfig(PretrainedConfig):
    model_type = "indicf5_hinglish"

    def __init__(self, dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512,
                 conv_layers=4, mel_dim=100, **kw):
        self.dim, self.depth, self.heads = dim, depth, heads
        self.ff_mult, self.text_dim = ff_mult, text_dim
        self.conv_layers, self.mel_dim = conv_layers, mel_dim
        super().__init__(**kw)


def _strip_prefix(sd, target_keys):
    variants = {"": sd}
    for p in PREFIXES:
        s = {k[len(p):]: v for k, v in sd.items() if k.startswith(p)}
        if s:
            variants[p] = s
    prefix, best = max(variants.items(), key=lambda kv: len(set(kv[1]) & target_keys))
    return prefix, {k: v for k, v in best.items() if k not in IGNORE}


class IndicF5Hinglish(PreTrainedModel):
    config_class = IndicF5HinglishConfig

    def __init__(self, config, vocab_path=None, asset_dir=None):
        super().__init__(config)
        self.asset_dir = asset_dir
        if vocab_path is None:
            vocab_path = os.path.join(asset_dir or ".", "vocab.txt")

        vocab_char_map, vocab_size = get_tokenizer(vocab_path, tokenizer="custom")
        self.model = CFM(
            transformer=DiT(dim=config.dim, depth=config.depth, heads=config.heads,
                            ff_mult=config.ff_mult, text_dim=config.text_dim,
                            conv_layers=config.conv_layers,
                            text_num_embeds=vocab_size, mel_dim=config.mel_dim),
            mel_spec_kwargs=dict(n_fft=1024, hop_length=HOP_LENGTH, win_length=1024,
                                 n_mel_channels=config.mel_dim,
                                 target_sample_rate=SAMPLE_RATE,
                                 mel_spec_type="vocos"),
            odeint_kwargs=dict(method="euler"),
            vocab_char_map=vocab_char_map,
        )
        self._vocoder = None

    @classmethod
    def from_pretrained(cls, path, device=None, **kw):
        device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        if os.path.isdir(path):
            get = lambda f: os.path.join(path, f)
            asset_dir = path
        else:
            get = lambda f: hf_hub_download(path, filename=f)
            asset_dir = os.path.dirname(get("vocab.txt"))
            vj = get("voices/voices.json")
            for _v in json.load(open(vj)).values():
                get("voices/" + _v["file"])

        cfg_dict = json.load(open(get("config.json")))
        config = IndicF5HinglishConfig(**cfg_dict.get("model_config", {}))
        model = cls(config, vocab_path=get("vocab.txt"), asset_dir=asset_dir)

        sd = load_file(get("model.safetensors"))
        target = set(model.model.state_dict())
        prefix, sd = _strip_prefix(sd, target)

        missing, unexpected = model.model.load_state_dict(sd, strict=False)
        real = [k for k in missing if k not in IGNORE]
        loaded = len(target) - len(real)
        if loaded < 0.9 * len(target):
            raise RuntimeError(f"only {loaded}/{len(target)} tensors loaded "
                               f"(prefix={prefix!r}); missing e.g. {real[:5]}")
        print(f"loaded {loaded}/{len(target)} tensors (prefix={prefix!r})")
        return model.to(device).eval()

    @property
    def vocoder(self):
        if self._vocoder is None:
            from vocos import Vocos
            self._vocoder = Vocos.from_pretrained(
                "charactr/vocos-mel-24khz").to(self.device).eval()
        return self._vocoder

    def voices(self):
        """Available bundled voices, keyed by name."""
        return json.load(open(os.path.join(self.asset_dir, "voices", "voices.json")))

    def voice(self, key):
        """Returns (audio_path, transcript) for a bundled voice."""
        meta = self.voices()
        if key not in meta:
            raise KeyError(f"{key!r} not found; available: {list(meta)}")
        v = meta[key]
        return os.path.join(self.asset_dir, "voices", v["file"]), v["transcript"]

    def default_voice(self):
        meta = json.load(open(os.path.join(self.asset_dir, "voices", "voices.json")))
        v = meta[next(iter(meta))]
        return os.path.join(self.asset_dir, "voices", v["file"]), v["transcript"]

    def _load_ref(self, path):
        import soundfile as sf
        import librosa
        audio, sr = sf.read(path, dtype="float32")
        if audio.ndim > 1:
            audio = audio.mean(axis=1)
        if sr != SAMPLE_RATE:
            audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE)
        t = torch.from_numpy(audio).unsqueeze(0)
        rms = torch.sqrt(torch.mean(t ** 2))
        if rms < TARGET_RMS:
            t = t * TARGET_RMS / rms
        return t.to(self.device), rms

    @torch.inference_mode()
    def generate(self, text, ref_audio=None, ref_text=None, nfe_step=32,
                 cfg_strength=2.0, sway_sampling_coef=-1.0, speed=1.0):
        if ref_audio is None or ref_text is None:
            ref_audio, ref_text = self.default_voice()

        cond, rms = self._load_ref(ref_audio)
        ref_text = ref_text.strip()
        if not ref_text.endswith((" ", ".", "।", "!", "?")):
            ref_text += ". "

        ref_len = cond.shape[-1] // HOP_LENGTH
        duration = ref_len + int(
            ref_len / len(ref_text.encode()) * len(text.encode()) / speed)

        generated, _ = self.model.sample(
            cond=cond, text=[ref_text + text], duration=duration,
            steps=nfe_step, cfg_strength=cfg_strength,
            sway_sampling_coef=sway_sampling_coef)

        mel = generated.to(torch.float32)[:, ref_len:, :].permute(0, 2, 1)
        wave = self.vocoder.decode(mel).squeeze().cpu()
        if rms < TARGET_RMS:
            wave = wave * rms / TARGET_RMS
        return wave.numpy().astype(np.float32), SAMPLE_RATE