import os
from typing import Literal, overload

import av
import av.container

from .._features import require_video
from .._typing import PathLike
from ..resource_cache import ResourceCache
from .input_container_mixin import InputContainerMixin

# Ensure video dependencies are available
require_video()

DEFAULT_CACHE_SIZE = int(os.environ.get("AV_CACHE_SIZE", 10))

# Global container cache for efficient video file access
_container_cache: ResourceCache["MockedInputContainer"] = ResourceCache(max_size=DEFAULT_CACHE_SIZE)


@overload
def open(
    file: PathLike, mode: Literal["r"], *, keep_av_open: bool = False, **kwargs
) -> av.container.InputContainer: ...


@overload
def open(file: PathLike, mode: Literal["w"], **kwargs) -> av.container.OutputContainer: ...


def open(file: PathLike, mode: Literal["r", "w"], *, keep_av_open: bool = False, **kwargs):
    """
    Open video container with optional caching for read operations.

    Args:
        file: Video file path or URL
        mode: Access mode ('r' for read, 'w' for write)
        keep_av_open: Enable caching for read containers
        **kwargs: Additional arguments passed to av.open
    """
    if mode == "r":
        if not keep_av_open:
            # Direct access without caching
            return av.open(file, "r", **kwargs)

        # Use cached container when keep_av_open=True
        cache_key = str(file)
        if cache_key not in _container_cache:
            # Create new cached container
            return MockedInputContainer(file, **kwargs)
        else:
            # Reuse existing cached container and increment reference count
            return _container_cache.acquire_entry(cache_key)
    else:
        return av.open(file, mode, **kwargs)


def cleanup_cache():
    """Clear all cached video containers from memory."""
    _container_cache.clear()


class MockedInputContainer(InputContainerMixin):
    """Cached wrapper for PyAV InputContainer with reference counting."""

    def __init__(self, file: PathLike, **kwargs):
        self._cache_key = str(file)
        self._container: av.container.InputContainer = av.open(file, "r", **kwargs)
        _container_cache.add_entry(self._cache_key, self)

    def __enter__(self) -> "MockedInputContainer":
        return self

    def close(self):
        """Release container reference and cleanup when no longer needed. Safe to call multiple times."""
        if self._cache_key in _container_cache and _container_cache[self._cache_key].refs > 0:
            _container_cache.release_entry(self._cache_key)


__all__ = ["open", "cleanup_cache"]
