"""Small native API; model objects remain editable and hookable PyTorch modules."""

from contextlib import nullcontext
from dataclasses import dataclass
import json
from pathlib import Path
from typing import Sequence

import torch

from .audio import Request, load_audio, normalize_prompt


@dataclass
class AudioSlot:
    """One ordered source; token_positions [U_i] index the padded input row."""

    sample_index: int
    audio_index: int
    token_positions: torch.Tensor
    samples_16k: int
    windows: int | None


@dataclass
class PreparedBatch:
    """Ready-to-forward tensors, validated user prompt strings and source slots.

    inputs contains input_ids/attention_mask [B,L]. MF adds mel [W,128,3000]
    and frame mask [W,3000]. MOSS adds mel [A,128,F_max], lengths [A],
    audio_input_mask [B,L], audio_group_sizes [B], audio_token_lens [A].
    All tensors live on the model device; floating features use model dtype.
    slots identify valid audio-token positions, never padding or timestamp digits.
    """

    inputs: dict[str, torch.Tensor]
    slots: list[AudioSlot]
    prompts: list[str]
    windows_per_source: list[int]


@dataclass
class AudioEmbedding:
    """One source's unpooled token features, with padding already removed.

    encoder [U_i,1280], projected [U_i,D], token_positions [U_i]. D=3584
    for MF, 4096 for MOSS. MOSS has three deepstack_encoder [U_i,1280] and
    deepstack_projected [U_i,4096] tensors; MF has empty DeepStack tuples.
    Encoder != projected != language-contextual representations.
    """

    sample_index: int
    audio_index: int
    token_positions: torch.Tensor
    encoder: torch.Tensor
    projected: torch.Tensor
    deepstack_encoder: tuple[torch.Tensor, ...]
    deepstack_projected: tuple[torch.Tensor, ...]


class Session:
    """Single-device inference session; not a training or concurrent-serving API.

    model is the actual torch module (not a remote text client). Users may edit
    its forward or attach hooks. Default methods disable gradients; grad=True
    enables ordinary autograd for feature/forward research without checkpointed
    training. Do not invoke the same MOSS instance concurrently: DeepStack hooks
    are scoped to one forward and removed on success or failure.
    """

    def __init__(self, model, processor, family: str):
        if family not in {"mf2601", "moss_music"}:
            raise ValueError("Require mf2601 or moss_music")
        self.model = model
        self.processor = processor
        self.family = family
        if family == "mf2601":
            from . import mf2601 as backend
            self.context_length = int(model.config.text_config.max_position_embeddings)
        else:
            from . import moss_music as backend
            self.context_length = int(model.config.language_config.max_position_embeddings)
        self.backend = backend

    def prepare(self, requests: Sequence[Request]) -> PreparedBatch:
        """Decode/expand B single-turn requests, return tensors and ordered slots.

        Reject marker mismatches, empty/nonfinite waveforms, unsupported native
        tokens and prompts leaving no output context. No truncation of tokens,
        source concatenation, deduplication or sorting is performed.
        """
        if not requests:
            raise ValueError("At least one Request is required")
        prompts = [normalize_prompt(r.prompt, len(r.audios), label_single_audio=self.family != "mf2601") for r in requests]
        groups = [[load_audio(a, crop_seconds=1200 if self.family == "mf2601" else None) for a in r.audios] for r in requests]
        inputs, lens, windows = self.backend.prepare(prompts, groups, self.processor)
        if inputs["input_ids"].shape[1] >= self.context_length:
            raise ValueError("Expanded multimodal prompt leaves no output context")
        device, dtype = self.model.device, self.model.dtype
        inputs = {k: v.to(device=device, dtype=dtype if v.is_floating_point() else v.dtype) for k, v in inputs.items()}
        token_id = self.model.config.audio_token_id if self.family == "mf2601" else self.processor.audio_token_id
        slots = []
        source_index = 0
        for sample_index, group in enumerate(groups):
            positions = torch.where((inputs["input_ids"][sample_index] == token_id) & inputs["attention_mask"][sample_index].bool())[0]
            offset = 0
            for audio_index, audio in enumerate(group):
                count = lens[source_index]
                source_positions = positions[offset:offset + count]
                if len(source_positions) != count or count <= 0:
                    raise ValueError("Audio marker expansion differs from source token lengths")
                slots.append(AudioSlot(sample_index, audio_index, source_positions, len(audio), windows[source_index] if windows else None))
                offset += count
                source_index += 1
            if offset != len(positions):
                raise ValueError("Unexpected extra audio tokens in expanded input")
        if source_index != len(lens):
            raise ValueError("Some source lengths were not consumed")
        return PreparedBatch(inputs, slots, prompts, windows)

    def _autocast(self):
        if self.model.device.type == "cuda" and self.model.dtype in {torch.bfloat16, torch.float16}:
            return torch.autocast("cuda", dtype=self.model.dtype)
        return nullcontext()

    def forward(self, batch: PreparedBatch, *, output_hidden_states: bool = False, grad: bool = False):
        """Return native ModelOutput: logits [B,1,V], optionally hidden [B,L,D].

        Runs prefill only with cache disabled. Hidden-state index 0 is embedded
        input; 1..K correspond to decoder-layer boundaries (last includes final
        normalization). Audio positions are language-contextual features, not raw
        audio. Full hidden states can be expensive for long prompts.
        """
        with torch.set_grad_enabled(grad), self._autocast():
            return self.model(**batch.inputs, use_cache=False, return_dict=True,
                              output_hidden_states=output_hidden_states, logits_to_keep=1)

    def extract_audio_embeddings(self, batch: PreparedBatch, *, grad: bool = False) -> list[AudioEmbedding]:
        """Encode sources without running the LLM; return one AudioEmbedding/slot.

        See AudioEmbedding for exact shapes. No mean pooling/detach is applied;
        grad=True preserves the autograd graph through encoder and projectors.
        """
        if not batch.slots:
            return []
        with torch.set_grad_enabled(grad), self._autocast():
            raw, projected, deep_raw, deep_projected = self.backend.audio_features(self.model, batch.inputs, batch.windows_per_source)
        total = sum(len(s.token_positions) for s in batch.slots)
        if raw.shape[0] != total or projected.shape[0] != total:
            raise ValueError("Feature lengths differ from expanded source token counts")
        if any(s.shape[0] != total for s in (*deep_raw, *deep_projected)):
            raise ValueError("DeepStack token count differs from final encoder stream")
        results = []
        offset = 0
        for slot in batch.slots:
            stop = offset + len(slot.token_positions)
            results.append(AudioEmbedding(slot.sample_index, slot.audio_index, slot.token_positions,
                           raw[offset:stop], projected[offset:stop], tuple(s[offset:stop] for s in deep_raw),
                           tuple(s[offset:stop] for s in deep_projected)))
            offset = stop
        return results

    def extract_contextual_embeddings(self, batch: PreparedBatch, layers: Sequence[int] = (-1,), *, grad: bool = False):
        """Return list[dict] per source: requested layer -> tensor [U_i,D].

        Uses forward output_hidden_states, including MOSS DeepStack injection.
        token_positions refer to the padded input row. Negative layer indexes
        follow Python indexing; invalid indexes fail, never silently clamp.
        """
        if not layers:
            raise ValueError("At least one hidden-state layer is required")
        outputs = self.forward(batch, output_hidden_states=True, grad=grad)
        if outputs.hidden_states is None:
            raise RuntimeError("Native model did not return language hidden states")
        return [{"sample_index": s.sample_index, "audio_index": s.audio_index,
                 "token_positions": s.token_positions,
                 "layers": {i: outputs.hidden_states[i][s.sample_index, s.token_positions] for i in layers}}
                for s in batch.slots]

    def generate(self, requests: Sequence[Request], *, profile: str = "greedy", max_new_tokens: int = 16384) -> list[str]:
        """Return one generated string/request using KV cache; no retry/rewrite.

        greedy decodes deterministically; sampled uses the base model's sampling
        defaults. In a batch, the token budget is capped by the longest prompt;
        use batch size 1 near the context limit.
        """
        if type(max_new_tokens) is not int or max_new_tokens < 1:
            raise ValueError("max_new_tokens must be a positive integer")
        if profile not in {"greedy", "sampled"}:
            raise ValueError("Unknown generation profile")
        batch = self.prepare(requests)
        width = batch.inputs["input_ids"].shape[1]
        budget = min(max_new_tokens, self.context_length - width)
        from transformers import GenerationConfig
        original = self.model.generation_config
        sampled = ({"temperature": 0.7 if self.family == "mf2601" else 1.0,
                    "top_p": 0.9 if self.family == "mf2601" else 0.8, "top_k": 50}
                   if profile == "sampled" else {})
        config = GenerationConfig(
            do_sample=profile == "sampled", max_new_tokens=budget,
            repetition_penalty=1.0,
            num_beams=1, num_return_sequences=1, use_cache=True,
            bos_token_id=original.bos_token_id, eos_token_id=original.eos_token_id,
            pad_token_id=self.processor.tokenizer.pad_token_id,
            **sampled,
        )
        # fork_rng restores the caller's RNG state; a sampled batch uses one RNG stream.
        devices = [self.model.device] if self.model.device.type == "cuda" else []
        with torch.random.fork_rng(devices=devices), torch.no_grad(), self._autocast():
            torch.manual_seed(3407)
            output = self.model.generate(**batch.inputs, generation_config=config, logits_to_keep=1)
        return [s.strip() for s in self.processor.batch_decode(output[:, width:], skip_special_tokens=True)]


def load_model(model_path: str | Path, *, device: str = "cuda:0", dtype: torch.dtype = torch.bfloat16,
               attention: str = "sdpa") -> Session:
    """Load a downloaded model folder (e.g. from snapshot_download).

    The returned Session exposes .model and .processor. No remote code or global
    patches are used. Missing, unexpected or mismatched tensors fail explicitly.
    """
    path = Path(model_path).expanduser()
    if not path.is_dir():
        raise FileNotFoundError("Pass a local model directory")
    if attention not in {"sdpa", "eager"}:
        raise ValueError("Supported attention: sdpa or eager")
    cfg = json.loads((path / "config.json").read_text())
    kwargs = dict(local_files_only=True, dtype=dtype, attn_implementation=attention, output_loading_info=True)
    if cfg.get("architectures") == ["MusicFlamingoForConditionalGeneration"]:
        from transformers import AutoProcessor, MusicFlamingoForConditionalGeneration
        processor = AutoProcessor.from_pretrained(path, local_files_only=True)
        model, loading = MusicFlamingoForConditionalGeneration.from_pretrained(path, **kwargs)
        from .mf2601 import patch_zero_token_window_timestamps
        patch_zero_token_window_timestamps(model)
        family = "mf2601"
    elif cfg.get("architectures") == ["MossMusicModel"]:
        from transformers import Qwen2Tokenizer
        from .vendor.configuration_moss_music import MossMusicConfig
        from .vendor.modeling_moss_music import MossMusicModel
        from .vendor.processing_moss_music import MelConfig, MossMusicProcessor
        from .moss_patches import (patch_moss_music_batch_forward, patch_moss_music_audio_encoder_official_chunks,
                                  patch_whisper_encoder_layer_outputs, reset_moss_music_sinusoid_buffer)

        class AttrDict(dict):
            def __getattr__(self, key):
                try:
                    return self[key]
                except KeyError as exc:
                    raise AttributeError(key) from exc

        markers = {cfg.get("audio_encoder_chunking"), cfg["audio_config"].get("audio_encoder_chunking")}
        if markers != {"official_chunked_400"}:
            raise ValueError("Expected a checkpoint trained with official 400-frame audio-encoder chunking")
        config = MossMusicConfig.from_dict(cfg)
        config.audio_config = AttrDict(config.audio_config)
        # The final encoder used eager attention; LLM SDPA is an inference backend
        # choice. Do not apply the persisted FA4 training backend to the encoder.
        config.audio_config["_attn_implementation"] = "eager"
        config.language_config._attn_implementation = attention
        config._attn_implementation = attention
        tokenizer = Qwen2Tokenizer.from_pretrained(path, local_files_only=True)
        processor = MossMusicProcessor(tokenizer, mel_config=MelConfig(), enable_time_marker=True)
        model, loading = MossMusicModel.from_pretrained(path, config=config, **kwargs)
        reset_moss_music_sinusoid_buffer(model)
        patch_whisper_encoder_layer_outputs(model)
        patch_moss_music_audio_encoder_official_chunks(model)
        patch_moss_music_batch_forward(model)
        family = "moss_music"
    else:
        raise ValueError("Unknown architecture; refuse to guess the model family")
    failures = {k: loading.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") if loading.get(k)}
    if failures:
        raise RuntimeError(f"Incomplete checkpoint loading: {failures}")
    model.to(device).eval()
    # Moving BF16 modules would downcast the nonpersistent sinusoid buffer.
    # Restore after placement too, without changing protected config files.
    if family == "moss_music":
        reset_moss_music_sinusoid_buffer(model)
    model.config.use_cache = True
    return Session(model, processor, family)