from __future__ import annotations

import numpy as np
import torch
from transformers import Gemma4Processor
from transformers.feature_extraction_utils import BatchFeature

from .parakeet_projector import PARAKEET_NAME, valid_frames_for

TARGET_SR = 16000


def _load_audio(item, target_sr: int = TARGET_SR):
    if isinstance(item, dict):
        arr = np.asarray(item.get("array"), dtype="float32")
        sr = int(item.get("sampling_rate") or target_sr)
    elif isinstance(item, (str, bytes)):
        import librosa
        arr, sr = librosa.load(item, sr=target_sr, mono=True)
    else:
        arr, sr = np.asarray(item, dtype="float32"), target_sr
    if arr.ndim > 1:
        arr = arr.mean(axis=1)
    if sr != target_sr:
        import librosa
        arr = librosa.resample(arr.astype("float32"), orig_sr=sr,
                               target_sr=target_sr)
    return arr.astype("float32")


def _collect_audio(messages):
    out = []
    for m in messages:
        content = m.get("content")
        if not isinstance(content, (list, tuple)):
            continue
        for part in content:
            if isinstance(part, dict) and part.get("type") in ("audio",
                                                               "input_audio"):
                out.append(part.get("audio", part.get("url")))
    return out


class LFG3Processor(Gemma4Processor):
    parakeet_name = PARAKEET_NAME

    def _feature_extractor(self):
        if getattr(self, "_parakeet_fe", None) is None:
            from transformers import AutoFeatureExtractor
            self._parakeet_fe = AutoFeatureExtractor.from_pretrained(
                self.parakeet_name
            )
        return self._parakeet_fe

    def _scaffold_ids(self):
        tok = self.tokenizer
        ids = tuple(tok.convert_tokens_to_ids(t)
                    for t in ("<|audio|>", "<|audio>", "<audio|>"))
        if any(i is None or i == tok.unk_token_id for i in ids):
            raise ValueError(f"audio scaffold tokens missing from vocab: {ids}")
        return ids

    def _batch_features(self, clips):
        fe = self._feature_extractor()
        per_clip, lengths = [], []
        for clip in clips:
            out = fe(clip, sampling_rate=TARGET_SR, return_tensors="pt")
            feats = out["input_features"][0]
            mask = out.get("attention_mask")
            n = int(mask[0].sum()) if mask is not None else feats.shape[0]
            per_clip.append(feats[:n])
            lengths.append(n)

        width = max(lengths)
        padded = torch.zeros((len(per_clip), width, per_clip[0].shape[-1]),
                             dtype=per_clip[0].dtype)
        mask = torch.zeros((len(per_clip), width), dtype=torch.long)
        for i, (feats, n) in enumerate(zip(per_clip, lengths)):
            padded[i, :n] = feats
            mask[i, :n] = 1
        return {"input_features": padded, "encoder_attention_mask": mask}

    def _splice(self, text: str, clip) -> tuple[list[int], int]:
        ids = self.tokenizer(text, add_special_tokens=False)["input_ids"]
        if clip is None:
            return ids, 0
        marker_id, boa_id, eoa_id = self._scaffold_ids()
        n = valid_frames_for(len(clip), TARGET_SR)
        try:
            pos = ids.index(marker_id)
        except ValueError as exc:
            raise ValueError(
                "audio supplied but no <|audio|> marker in the rendered "
                "prompt; include a {'type': 'audio'} part in the message"
            ) from exc
        return ids[:pos] + [boa_id] + [marker_id] * n + [eoa_id] + ids[pos + 1:], n

    def apply_chat_template(
        self,
        conversations,
        tokenize: bool = True,
        return_dict: bool = True,
        return_tensors: str = "pt",
        padding: bool = False,
        add_generation_prompt: bool = True,
        enable_thinking: bool = False,
        audio=None,
        **kwargs,
    ):
        render = super().apply_chat_template
        batched = bool(conversations) and isinstance(conversations[0],
                                                     (list, tuple))
        convs = list(conversations) if batched else [conversations]

        texts = [render(c, tokenize=False,
                        add_generation_prompt=add_generation_prompt,
                        enable_thinking=enable_thinking, **kwargs)
                 for c in convs]
        if not tokenize:
            return texts if batched else texts[0]

        rows, clips, frames = [], [], []
        for conv, text in zip(convs, texts):
            found = _collect_audio(conv) if audio is None else list(audio)
            if len(found) > 1:
                raise ValueError("one audio clip per conversation is supported; "
                                 f"got {len(found)}")
            clip = _load_audio(found[0]) if found else None
            ids, n = self._splice(text, clip)
            rows.append(ids)
            if clip is not None:
                clips.append(clip)
                frames.append(n)

        pad_id = self.tokenizer.pad_token_id or 0
        width = max(len(r) for r in rows)
        input_ids = [[pad_id] * (width - len(r)) + r for r in rows]
        attn = [[0] * (width - len(r)) + [1] * len(r) for r in rows]

        data = {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attn, dtype=torch.long),
        }
        if clips:
            if len(clips) != len(rows):
                raise ValueError("every conversation in a batch must carry "
                                 "audio, or none of them may")
            data.update(self._batch_features(clips))
            data["valid_frames"] = torch.tensor(frames, dtype=torch.long)
        return BatchFeature(data=data)