"""EmbeddingGemma 2 adapter. Neural computation runs exclusively in Core AI."""

import json
from pathlib import Path

import numpy as np
from coreai.runtime import AIModel, NDArray, SpecializationOptions
from transformers import AutoProcessor


class EmbeddingGemma2CoreAI:
    @classmethod
    async def load(cls, bundle):
        self = cls()
        self.root = Path(bundle)
        self.manifest = json.loads((self.root / "manifest.json").read_text())
        self.processor = AutoProcessor.from_pretrained(self.root)
        self.prompts = json.loads(
            (self.root / "config_sentence_transformers.json").read_text()
        )["prompts"]
        self.models = {}
        self.functions = {}
        for name, asset in [
            ("text", "embedding"),
            ("vision", "vision"),
            ("audio", "audio"),
        ]:
            self.models[name] = await AIModel.load(
                self.root / f"{asset}.aimodel", SpecializationOptions.default()
            )
            self.functions[name] = self.models[name].load_function("main")
        return self

    async def call(self, name, inputs):
        result = await self.functions[name](
            {k: NDArray(np.ascontiguousarray(v)) for k, v in inputs.items()}
        )
        return [result[k].numpy() for k in self.manifest[f"{name}_outputs"]]

    async def encode(
        self,
        text=None,
        *,
        images=None,
        videos=None,
        audio=None,
        sampling_rate=16000,
        prompt_name=None,
        truncate_dim=768,
    ):
        if not 1 <= truncate_dim <= 768:
            raise ValueError("truncate_dim must be between 1 and 768")
        if prompt_name:
            if images is not None or videos is not None or audio is not None:
                raise ValueError("Task prompts apply to text-only inputs")
            text = self.prompts[prompt_name] + (text or "")
        kwargs = {"return_tensors": "np"}
        if text is not None:
            kwargs["text"] = text
        if images is not None:
            kwargs["images"] = images
        if videos is not None:
            if not isinstance(videos, np.ndarray) or videos.ndim != 4:
                raise ValueError(
                    "Video must be an already sampled NumPy [frames,height,width,channels] array"
                )
            kwargs["videos"] = videos
            kwargs["videos_kwargs"] = {"do_sample_frames": False}
        if audio is not None:
            samples = np.asarray(audio)
            if sampling_rate != 16000:
                raise ValueError("Audio must be sampled at 16000 Hz")
            if samples.ndim != 1 or len(samples) > 179200:
                raise ValueError("One mono clip of at most 11.2 seconds is supported")
            kwargs.update(audio=samples, sampling_rate=sampling_rate)
        encoded = self.processor(**kwargs)
        return await self.encode_processed(encoded, truncate_dim=truncate_dim)

    async def encode_processed(self, e, truncate_dim=768):
        ids = np.asarray(e["input_ids"], dtype=np.int32)
        if ids.shape[0] != 1:
            raise ValueError(
                "Batch size must be one; call encode separately per example"
            )
        n = ids.shape[1]
        cap = self.manifest["sequence_length"]
        if n > cap:
            raise ValueError(
                f"Input needs {n} tokens, capacity {cap}; input was not truncated"
            )
        media = np.zeros((1, cap, 512), np.float32)
        media_mask = np.zeros((1, cap, 1), bool)
        p = self.processor
        for pixels_key, pos_key, token in [
            ("pixel_values", "image_position_ids", p.image_token_id),
            ("pixel_values_videos", "video_position_ids", p.video_token_id),
        ]:
            if pixels_key not in e:
                continue
            features = []
            for pixels, positions in zip(e[pixels_key], e[pos_key]):
                count = len(pixels)
                capacity = self.manifest["vision_shape"][1]
                if count > capacity:
                    raise ValueError("Vision patch capacity exceeded")
                pix = np.pad(
                    np.asarray(pixels, dtype=np.float32),
                    ((0, capacity - count), (0, 0)),
                )
                pos = np.pad(
                    np.asarray(positions, dtype=np.int32),
                    ((0, capacity - count), (0, 0)),
                    constant_values=-1,
                )
                values, valid = await self.call(
                    "vision",
                    {"pixel_values": pix[None], "pixel_position_ids": pos[None]},
                )
                features.append(values[valid.astype(bool)])
            features = np.concatenate(features, axis=0)
            slots = np.flatnonzero(ids[0] == token)
            if len(slots) != len(features):
                raise ValueError("Visual placeholder count differs from encoder output")
            media[0, slots] = features
            media_mask[0, slots] = True
        if "input_features" in e:
            features = []
            for mel, mask in zip(e["input_features"], e["input_features_mask"]):
                count = len(mel)
                capacity = self.manifest["audio_shape"][1]
                if count > capacity:
                    raise ValueError("Audio feature capacity exceeded")
                mel = np.pad(
                    np.asarray(mel, dtype=np.float32), ((0, capacity - count), (0, 0))
                )
                mask = np.pad(np.asarray(mask, dtype=bool), (0, capacity - count))
                values, valid = await self.call(
                    "audio",
                    {"input_features": mel[None], "input_features_mask": mask[None]},
                )
                features.append(values[valid.astype(bool)])
            features = np.concatenate(features, axis=0)
            slots = np.flatnonzero(ids[0] == p.audio_token_id)
            if len(slots) != len(features):
                raise ValueError("Audio placeholder count differs from encoder output")
            media[0, slots] = features
            media_mask[0, slots] = True
        padded = np.pad(ids, ((0, 0), (0, cap - n)))
        mask = np.pad(
            np.asarray(e["attention_mask"], dtype=np.int32), ((0, 0), (0, cap - n))
        )
        value = (
            await self.call(
                "text",
                {
                    "input_ids": padded,
                    "attention_mask": mask,
                    "media_embeddings": media,
                    "media_mask": media_mask,
                },
            )
        )[0][0, :truncate_dim]
        return value / np.linalg.norm(value)