# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""OmniInteract dataset discovery for the standard serving benchmark."""

from __future__ import annotations

import json
import logging
import os
import random
import shutil
import tarfile
import tempfile
from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from vllm.benchmarks.datasets import BenchmarkDataset, SampleRequest
from vllm.tokenizers import TokenizerLike
from vllm.transformers_utils.repo_utils import hf_fs

OMNIINTERACT_SUBSETS = ("1q1a", "1q1a_math", "1qna")
DEFAULT_OMNIINTERACT_REPO = "lucky-lance/OmniInteract"
DEFAULT_MAX_VIDEO_DURATION_S = 3600.0
logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class OmniInteractCase:
    subset: str
    video_rel: str
    video_path: Path
    annotation_path: Path
    scene_type: str


@dataclass(frozen=True)
class OmniInteractSessionOptions:
    output_root: Path
    timeout_s: float
    media_timeout_s: float
    ref_audio: str
    require_response: bool = False
    max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S


@dataclass(frozen=True)
class OmniInteractPreparedInput:
    duration_s: float
    pcm16: bytes = field(repr=False)
    video_frames: tuple[str | None, ...] = field(repr=False)
    ref_audio_data_url: str = field(repr=False)


@dataclass
class OmniInteractSampleRequest(SampleRequest):
    omniinteract_case: OmniInteractCase | None = None
    omniinteract_options: OmniInteractSessionOptions | None = None
    omniinteract_prepared_input: OmniInteractPreparedInput | None = field(default=None, repr=False)


class OmniInteractDataset(BenchmarkDataset):
    def __init__(
        self,
        *,
        data_root: str | None,
        dataset_repo: str = DEFAULT_OMNIINTERACT_REPO,
        subsets: Sequence[str] = OMNIINTERACT_SUBSETS,
        random_seed: int = 0,
        disable_shuffle: bool = False,
    ) -> None:
        super().__init__(
            dataset_path=data_root or dataset_repo,
            random_seed=random_seed,
            disable_shuffle=disable_shuffle,
        )
        self.root = resolve_omniinteract_root(data_root, dataset_repo)
        self.subsets = tuple(subsets)

    def sample(
        self,
        tokenizer: TokenizerLike | None,
        num_requests: int,
        *,
        request_id_prefix: str = "",
        options: OmniInteractSessionOptions,
        **_: Any,
    ) -> list[SampleRequest]:
        del tokenizer
        cases = discover_omniinteract_cases(
            self.root,
            self.subsets,
            num_prompts=num_requests,
            seed=self.random_seed,
            disable_shuffle=self.disable_shuffle,
        )
        return [
            OmniInteractSampleRequest(
                prompt="",
                prompt_len=0,
                expected_output_len=0,
                multi_modal_data=None,
                request_id=f"{request_id_prefix}{index}",
                omniinteract_case=case,
                omniinteract_options=options,
            )
            for index, case in enumerate(cases)
        ]


def _data_dir(root: Path) -> Path:
    for candidate in (root, root / "data"):
        if any((candidate / subset).is_dir() for subset in OMNIINTERACT_SUBSETS):
            return candidate
    raise FileNotFoundError(f"OmniInteract data not found under {root}")


def _confined_path(root: Path, relative: object, *, field: str) -> Path:
    text, resolved_root = str(relative or ""), root.resolve()
    path = Path(text)
    destination = (root / path).resolve()
    if not text or path.is_absolute() or ".." in path.parts or not destination.is_relative_to(resolved_root):
        raise ValueError(f"Unsafe OmniInteract {field} path: {text!r}")
    return destination


def _safe_extract(archive: tarfile.TarFile, target: Path) -> None:
    root = target.resolve()
    members = archive.getmembers()
    for member in members:
        path = Path(member.name)
        if (
            path.is_absolute()
            or ".." in path.parts
            or not (root / path).resolve().is_relative_to(root)
            or not (member.isdir() or member.isfile())
        ):
            raise ValueError(f"Unsafe path in OmniInteract archive: {member.name!r}")
    archive.extractall(target, members=members, filter="data")


def _extract_archive(archive: Path, target: Path) -> Path:
    if target.is_symlink():
        raise ValueError(f"Refusing to extract through symlink: {target}")
    stat = archive.stat()
    published = target / f"{stat.st_size}-{stat.st_mtime_ns}"
    if published.is_symlink():
        raise ValueError(f"Refusing to use symlinked extraction: {published}")
    if published.is_dir():
        try:
            return _data_dir(published)
        except FileNotFoundError:
            pass
    target.mkdir(parents=True, exist_ok=True)
    staging = Path(tempfile.mkdtemp(prefix=".tmp-", dir=target))
    try:
        with tarfile.open(archive, "r:*") as handle:
            _safe_extract(handle, staging)
        root = _data_dir(staging)
        try:
            staging.replace(published)
            return published / root.relative_to(staging)
        except OSError:
            if not published.is_dir():
                raise
            return _data_dir(published)
    finally:
        shutil.rmtree(staging, ignore_errors=True)


def resolve_omniinteract_root(data_root: str | None, dataset_repo: str = DEFAULT_OMNIINTERACT_REPO) -> Path:
    """Resolve a local tree/archive or download the official archive."""

    if data_root:
        local = Path(data_root).expanduser().resolve()
        if local.is_file():
            return _extract_archive(local, local.parent / f".{local.stem}.vllm_omni_extracted")
        if not local.is_dir():
            raise FileNotFoundError(f"--dataset-path does not exist: {local}")
        try:
            return _data_dir(local)
        except FileNotFoundError:
            for name in ("data.tar.gz", "data.tar"):
                if (archive := local / name).is_file():
                    return _extract_archive(archive, local / ".vllm_omni_extracted")
            raise

    cache = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
    target = cache / "vllm_omni" / "omniinteract" / dataset_repo.replace("/", "__")
    downloads = target / ".downloads"
    downloads.mkdir(parents=True, exist_ok=True)
    filesystem, errors = hf_fs(), []
    for name in ("data.tar.gz", "data.tar"):
        archive = downloads / name
        try:
            if not archive.is_file():
                descriptor, temporary_name = tempfile.mkstemp(prefix=f".{name}.", dir=downloads)
                os.close(descriptor)
                temporary = Path(temporary_name)
                try:
                    filesystem.get_file(f"datasets/{dataset_repo}/{name}", str(temporary))
                    if not temporary.stat().st_size:
                        raise ValueError("downloaded archive is empty")
                    temporary.replace(archive)
                finally:
                    temporary.unlink(missing_ok=True)
            return _extract_archive(archive, target)
        except Exception as exc:  # noqa: BLE001, PERF203 - Hub errors vary by version
            errors.append(f"{name}: {exc}")
    raise FileNotFoundError(f"Could not download OmniInteract from {dataset_repo!r}: {'; '.join(errors)}")


def _mapping_cases(root: Path, subset: str) -> list[OmniInteractCase]:
    mapping_path = root / "video_json_map.json"
    if not mapping_path.is_file():
        raise FileNotFoundError(f"Missing OmniInteract mapping: {mapping_path}")
    mapping = json.loads(mapping_path.read_text())
    entries = mapping.get("entries") if isinstance(mapping, dict) else None
    if not isinstance(entries, list):
        raise ValueError(f"Invalid OmniInteract mapping: {mapping_path}")
    cases = []
    for row in entries:
        if not isinstance(row, dict):
            raise ValueError(f"Invalid OmniInteract mapping row in {mapping_path}")
        video_rel, annotation_rel = str(row.get("video") or ""), str(row.get("annotation") or "")
        video = _confined_path(root, video_rel, field="video")
        annotation = _confined_path(root, annotation_rel, field="annotation")
        if not video.is_file() or not annotation.is_file():
            raise FileNotFoundError(f"OmniInteract mapping references missing files: {video_rel!r}, {annotation_rel!r}")
        cases.append(
            OmniInteractCase(
                subset,
                video_rel,
                video,
                annotation,
                str(row.get("scene_type") or "multi_turn").lower(),
            )
        )
    return cases


def _one_to_many_cases(root: Path) -> list[OmniInteractCase]:
    videos, annotations = root / "videos_bench", root / "annotations"
    if not videos.is_dir() or not annotations.is_dir():
        raise FileNotFoundError(f"Invalid OmniInteract 1qna layout under {root}")
    cases = []
    for video in sorted(videos.rglob("*.mp4")):
        if not video.resolve().is_relative_to(videos.resolve()):
            raise ValueError(f"Unsafe OmniInteract video path: {video}")
        relative = video.relative_to(videos)
        annotation = (annotations / relative).with_suffix(".json").resolve()
        if not annotation.is_relative_to(annotations.resolve()):
            raise ValueError(f"Unsafe OmniInteract annotation path: {relative}")
        if not annotation.is_file():
            raise FileNotFoundError(f"Missing OmniInteract annotation: {annotation}")
        cases.append(OmniInteractCase("1qna", f"videos_bench/{relative.as_posix()}", video, annotation, "1qna"))
    return cases


def discover_omniinteract_cases(
    root: Path,
    subsets: Sequence[str],
    *,
    num_prompts: int,
    seed: int = 0,
    disable_shuffle: bool = False,
) -> list[OmniInteractCase]:
    invalid = set(subsets) - set(OMNIINTERACT_SUBSETS)
    if invalid:
        raise ValueError(f"Unsupported OmniInteract subsets: {sorted(invalid)}")
    if not subsets:
        raise ValueError("At least one OmniInteract subset is required")
    if len(set(subsets)) != len(subsets):
        raise ValueError("OmniInteract subsets must not contain duplicates")
    if num_prompts < 0:
        raise ValueError("num_prompts must be non-negative")
    data_root, cases = _data_dir(root.resolve()), []
    for subset in subsets:
        subset_root = data_root / subset
        selected = _one_to_many_cases(subset_root) if subset == "1qna" else _mapping_cases(subset_root, subset)
        if not selected:
            raise ValueError(f"No OmniInteract sessions found for requested subset {subset!r}")
        cases.extend(selected)
    if len({case.video_path for case in cases}) != len(cases):
        raise ValueError("OmniInteract dataset contains duplicate video paths")
    if not disable_shuffle:
        random.Random(seed).shuffle(cases)
    if num_prompts:
        if num_prompts > len(cases):
            logger.warning(
                "Requested %d OmniInteract prompts but only %d are available; using all cases", num_prompts, len(cases)
            )
        cases = cases[:num_prompts]
    return cases


def case_manifest(case: OmniInteractCase, output_dir: Path) -> dict[str, Any]:
    return {
        "sample_id": f"{case.subset}__{output_dir.name}",
        "video": str(case.video_path),
        "gt_json": str(case.annotation_path.resolve()),
        "model_json": str((output_dir / "wav_transcript.json").resolve()),
        "scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type,
    }