"""Expose verified speech generation through Transformers AutoClasses."""
from __future__ import annotations
import hashlib
import json
import tempfile
from dataclasses import dataclass
from pathlib import Path
import torch
from torch import Tensor
from transformers import AutoTokenizer, HiggsAudioV2TokenizerConfig, HiggsAudioV2TokenizerModel, PreTrainedConfig, PreTrainedModel
from transformers.utils.generic import ModelOutput
from .adaptation import inject_adapter
from .backbone import Generator
from .codec import Codec
from .configuration import Geometry
from .synthesis import InferenceSettings, SpeechRequest, read_waveform, render_speech

def compute_weight_identity(checkpoint: str, quantization: object, has_adapter: bool) -> str:
    payload = {'checkpoint': checkpoint, 'quantization': quantization, 'has_adapter': has_adapter}
    return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()

class SpeechConfig(PreTrainedConfig):
    model_type = 'agbalu_speech'

    def __init__(self, geometry: dict[str, object] | None=None, audio_configuration: dict[str, object] | None=None, inference: dict[str, object] | None=None, checkpoint_sha256: str='0' * 64, source_stamp: dict[str, object] | None=None, has_adapter: bool=True, **kwargs: object) -> None:
        self.geometry = Geometry.model_validate(geometry or {}).model_dump()
        self.audio_configuration = audio_configuration or {}
        self.inference = InferenceSettings.model_validate(inference or {}).model_dump()
        self.checkpoint_sha256 = checkpoint_sha256
        self.source_stamp = source_stamp or {}
        self.has_adapter = has_adapter
        self.sampling_rate = 24000
        super().__init__(**kwargs)

@dataclass
class SpeechOutput(ModelOutput):
    waveform: Tensor | None = None
    sampling_rate: int | None = None
    report: dict[str, object] | None = None

class SpeechModel(PreTrainedModel):
    config: SpeechConfig
    config_class = SpeechConfig
    main_input_name = 'text'
    _keep_in_fp32_modules = ['audio_tokenizer']
    _keep_in_fp32_modules_strict = ['audio_tokenizer', 'factor_a', 'factor_b']
    _no_split_modules = ['Qwen3DecoderLayer', 'AdaptedLinear']

    def __init__(self, config: SpeechConfig) -> None:
        super().__init__(config)
        self.generator = Generator(Geometry.model_validate(config.geometry))
        if config.has_adapter:
            inject_adapter(self.generator)
        self.audio_tokenizer = HiggsAudioV2TokenizerModel(HiggsAudioV2TokenizerConfig(**config.audio_configuration))
        self.post_init()

    def save_pretrained(self, save_directory: str | Path, **kwargs: object) -> None:
        commit: object = getattr(self.config, '_commit_hash', None)
        tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path, revision=commit if isinstance(commit, str) else 'main', trust_remote_code=True)
        super().save_pretrained(save_directory, **kwargs)
        tokenizer.save_pretrained(save_directory)

    @torch.inference_mode()
    def forward(self, text: str, reference_audio: str | Path, reference_text: str, output_path: str | Path | None=None) -> SpeechOutput:
        return self.generate(text, reference_audio, reference_text, output_path)

    @torch.inference_mode()
    def generate(self, text: str, reference_audio: str | Path, reference_text: str, output_path: str | Path | None=None) -> SpeechOutput:
        request = SpeechRequest(text=text, reference=Path(reference_audio), reference_text=reference_text, checkpoint_sha256=compute_weight_identity(self.config.checkpoint_sha256, self.config.to_dict().get('quantization_config'), self.config.has_adapter))
        if not request.reference.is_file():
            message = f'reference audio missing: {request.reference}'
            raise FileNotFoundError(message)
        settings = InferenceSettings.model_validate(self.config.inference)
        device = self.generator.audio_embedding.weight.device
        commit: object = getattr(self.config, '_commit_hash', None)
        tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path, revision=commit if isinstance(commit, str) else 'main', trust_remote_code=True)
        with tempfile.TemporaryDirectory(prefix='speech-') as temporary:
            destination = Path(output_path) if output_path is not None else Path(temporary) / 'speech.wav'
            report = render_speech(self.generator, Codec(self.audio_tokenizer, device), tokenizer, request, settings, destination, provenance_stamp=self.config.source_stamp)
            waveform = torch.from_numpy(read_waveform(destination)).unsqueeze(0)
        return SpeechOutput(waveform=waveform, sampling_rate=24000, report=report)
SpeechConfig.register_for_auto_class()
SpeechModel.register_for_auto_class('AutoModelForTextToWaveform')