import importlib
import json
import os
import pathlib
import sys

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Union

import numpy as np
import pytest

import torch

from torchcodec._core import get_ffmpeg_library_versions


# Decorator for skipping CUDA tests when CUDA isn't available. The tests are
# effectively marked to be skipped in pytest_collection_modifyitems() of
# conftest.py
def needs_cuda(test_item):
    return pytest.mark.needs_cuda(test_item)


def cpu_and_cuda():
    return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))


def get_ffmpeg_major_version():
    return int(get_ffmpeg_library_versions()["ffmpeg_version"].split(".")[0])


# For use with decoded data frames. On CPU Linux, we expect exact, bit-for-bit
# equality. On CUDA Linux, we expect a small tolerance.
# On other platforms (e.g. MacOS), we also allow a small tolerance. FFmpeg does
# not guarantee bit-for-bit equality across systems and architectures, so we
# also cannot. We currently use Linux on x86_64 as our reference system.
def assert_frames_equal(*args, **kwargs):
    if sys.platform == "linux":
        if args[0].device.type == "cuda":
            atol = 2
            if get_ffmpeg_major_version() == 4:
                assert_tensor_close_on_at_least(
                    args[0], args[1], percentage=95, atol=atol
                )
            else:
                torch.testing.assert_close(*args, **kwargs, atol=atol, rtol=0)
        else:
            torch.testing.assert_close(*args, **kwargs, atol=0, rtol=0)
    else:
        torch.testing.assert_close(*args, **kwargs, atol=3, rtol=0)
