import json
import numpy as np
import torch
from torch.utils.data import Dataset
from video_inputs import build_video_item, prepare
from modeling_skivlm import POSE_START, POSE_PATCH, POSE_END

IGNORE = -100
QUESTION = "Describe what the person is doing in this video."


def pose_placeholder(num_pose_tokens):
    return POSE_START + POSE_PATCH * num_pose_tokens + POSE_END


class MotionXSkiDataset(Dataset):
    def __init__(self, manifest_path, pose_path, processor, use_skeleton=True,
                 num_pose_tokens=8, video_fps=1.0, max_frames=8,
                 max_pixels=200704, limit=0, seed=0):
        items = json.load(open(manifest_path))
        if limit and limit < len(items):
            rng = np.random.default_rng(seed)
            idx = rng.choice(len(items), size=limit, replace=False)
            items = [items[k] for k in sorted(idx)]
        self.items = items
        self.pose = np.load(pose_path, mmap_mode="r")
        self.processor = processor
        self.use_skeleton = use_skeleton
        self.num_pose_tokens = num_pose_tokens
        self.video_fps = video_fps
        self.max_frames = max_frames
        self.max_pixels = max_pixels

    def __len__(self):
        return len(self.items)

    def _messages(self, m):
        text = QUESTION
        if self.use_skeleton:
            text = pose_placeholder(self.num_pose_tokens) + "\n" + QUESTION
        return [{
            "role": "user",
            "content": [build_video_item(m["video"], self.video_fps,
                                         self.max_frames, self.max_pixels),
                        {"type": "text", "text": text}],
        }]

    def __getitem__(self, i):
        for attempt in range(8):
            try:
                return self._build(i)
            except Exception as e:
                print("[dataset] skipping", self.items[i]["stem"], type(e).__name__, flush=True)
                i = int(np.random.default_rng(i + attempt).integers(len(self.items)))
        raise RuntimeError("too many unreadable videos")

    def _build(self, i):
        m = self.items[i]
        messages = self._messages(m)
        prompt = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True)
        answer = m["text"] + self.processor.tokenizer.eos_token
        video_inputs, vkw = prepare(messages, m["video"])

        enc = self.processor(text=[prompt + answer], videos=video_inputs,
                             return_tensors="pt", **vkw)
        prompt_only = self.processor(text=[prompt], videos=video_inputs,
                                     return_tensors="pt", **vkw)
        n_prompt = prompt_only["input_ids"].shape[1]
        input_ids = enc["input_ids"][0]
        labels = input_ids.clone()
        labels[:n_prompt] = IGNORE

        x = torch.from_numpy(np.asarray(self.pose[m["i"]], dtype=np.float32))
        x = x.permute(2, 0, 1).unsqueeze(-1)
        return {
            "input_ids": input_ids,
            "attention_mask": enc["attention_mask"][0],
            "labels": labels,
            "pixel_values_videos": enc.get("pixel_values_videos"),
            "video_grid_thw": enc.get("video_grid_thw"),
            "pose": x,
        }


def make_collator(pad_token_id):
    def collate(batch):
        maxlen = max(b["input_ids"].shape[0] for b in batch)

        def pad(x, val):
            return torch.cat([x, x.new_full((maxlen - x.shape[0],), val)])

        out = {
            "input_ids": torch.stack([pad(b["input_ids"], pad_token_id) for b in batch]),
            "attention_mask": torch.stack([pad(b["attention_mask"], 0) for b in batch]),
            "labels": torch.stack([pad(b["labels"], IGNORE) for b in batch]),
            "pose": torch.stack([b["pose"] for b in batch]),
        }
        if batch[0]["pixel_values_videos"] is not None:
            out["pixel_values_videos"] = torch.cat([b["pixel_values_videos"] for b in batch], dim=0)
            out["video_grid_thw"] = torch.cat([b["video_grid_thw"] for b in batch], dim=0)
        return out

    return collate