# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: E501

"""Unit tests for MediaFileManager"""

from __future__ import annotations

import io
import random
import unittest
from unittest import TestCase, mock
from unittest.mock import MagicMock, call, mock_open

import pytest

from streamlit.runtime.media_file_manager import MediaFileManager
from streamlit.runtime.media_file_storage import MediaFileKind, MediaFileStorageError
from streamlit.runtime.memory_media_file_storage import (
    MemoryFile,
    MemoryMediaFileStorage,
    _calculate_file_id,
)
from tests.exception_capturing_thread import call_on_threads


def random_coordinates():
    return f"{random.randint(1, 4)}.{(random.randint(1, 12), random.randint(1, 12))}.{random.randint(1, 99)}"


# Smallest possible "real" media files for a handful of different formats.
# Sourced from https://github.com/mathiasbynens/small
AUDIO_FIXTURES = {
    "wav": {
        "content": b"RIFF$\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00D\xac\x00\x00\x88X\x01\x00\x02\x00\x10\x00data\x00\x00\x00\x00",
        "mimetype": "audio/wav",
    },
    "mp3": {
        "content": b"\xff\xe3\x18\xc4\x00\x00\x00\x03H\x00\x00\x00\x00LAME3.98.2\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00",
        "mimetype": "audio/mp3",
    },
}


VIDEO_FIXTURES = {
    "mp4": {
        "content": b"\x00\x00\x00\x1cftypisom\x00\x00\x02\x00isomiso2mp41\x00\x00\x00\x08free\x00\x00\x02\xefmdat!\x10\x05",
        "mimetype": "video/mp4",
    },
    "webm": {
        "content": b'\x1aE\xdf\xa3@ B\x86\x81\x01B\xf7\x81\x01B\xf2\x81\x04B\xf3\x81\x08B\x82@\x04webmB\x87\x81\x02B\x85\x81\x02\x18S\x80g@\x8d\x15I\xa9f@(*\xd7\xb1@\x03\x0fB@M\x80@\x06whammyWA@\x06whammyD\x89@\x08@\x8f@\x00\x00\x00\x00\x00\x16T\xaek@1\xae@.\xd7\x81\x01c\xc5\x81\x01\x9c\x81\x00"\xb5\x9c@\x03und\x86@\x05V_VP8%\x86\x88@\x03VP8\x83\x81\x01\xe0@\x06\xb0\x81\x08\xba\x81\x08\x1fC\xb6u@"\xe7\x81\x00\xa3@\x1c\x81\x00\x00\x800\x01\x00\x9d\x01*\x08\x00\x08\x00\x01@&%\xa4\x00\x03p\x00\xfe\xfc\xf4\x00\x00',
        "mimetype": "video/webm",
    },
}


IMAGE_FIXTURES = {
    "png": {
        "content": b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82",
        "mimetype": "image/png",
    },
    "jpg": {
        "content": b"\xff\xd8\xff\xdb\x00C\x00\x03\x02\x02\x02\x02\x02\x03\x02\x02\x02\x03\x03\x03\x03\x04\x06\x04\x04\x04\x04\x04\x08\x06\x06\x05\x06\t\x08\n\n\t\x08\t\t\n\x0c\x0f\x0c\n\x0b\x0e\x0b\t\t\r\x11\r\x0e\x0f\x10\x10\x11\x10\n\x0c\x12\x13\x12\x10\x13\x0f\x10\x10\x10\xff\xc9\x00\x0b\x08\x00\x01\x00\x01\x01\x01\x11\x00\xff\xcc\x00\x06\x00\x10\x10\x05\xff\xda\x00\x08\x01\x01\x00\x00?\x00\xd2\xcf \xff\xd9",
        "mimetype": "image/jpg",
    },
}

TEXT_FIXTURES = {
    "txt": {"content": b"Hello world", "mimetype": "text/plain"},
    "csv": {
        "content": b"""
                    Foo, Bar
                    123, 456
                    789, 111""",
        "mimetype": "text/csv",
    },
}

ALL_FIXTURES = {}
ALL_FIXTURES.update(AUDIO_FIXTURES)
ALL_FIXTURES.update(VIDEO_FIXTURES)
ALL_FIXTURES.update(IMAGE_FIXTURES)
ALL_FIXTURES.update(TEXT_FIXTURES)


class MediaFileManagerTest(TestCase):
    def setUp(self):
        super().setUp()
        self.storage = MemoryMediaFileStorage("/mock/endpoint")
        self.media_file_manager = MediaFileManager(self.storage)
        random.seed(1337)

    def _add_file_and_get_object(
        self,
        content: bytes,
        mimetype: str,
        coordinates: str,
        filename: str | None = None,
    ) -> MemoryFile:
        """Add a new file to our test manager and return its MediaFile object."""
        file_id = _calculate_file_id(content, mimetype, filename)
        self.media_file_manager.add(content, mimetype, coordinates, filename)
        return self.storage.get_file(file_id)

    def test_calculate_file_id(self):
        """Test that file_id generation from data works as expected."""

        fake_bytes = "\x00\x00\xff\x00\x00\xff\x00\x00\xff\x00\x00\xff\x00".encode()
        test_hash = "2ba850426b188d25adc5a37ad313080c346f5e88e069e0807d0cdb2b"
        assert test_hash == _calculate_file_id(fake_bytes, "media/any")

        # Make sure we get different file ids for files with same bytes but diff't mimetypes.
        assert _calculate_file_id(fake_bytes, "audio/wav") != _calculate_file_id(
            fake_bytes, "video/mp4"
        )

        # Make sure we get different file ids for files with same bytes and mimetypes but diff't filenames.
        assert _calculate_file_id(
            fake_bytes, "audio/wav", filename="name1.wav"
        ) != _calculate_file_id(fake_bytes, "audio/wav", filename="name2.wav")

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session_id"),
    )
    def test_reject_null_files(self):
        """MediaFileManager.add raises a TypeError if it's passed None."""
        with pytest.raises(TypeError):
            self.media_file_manager.add(None, "media/any", random_coordinates())

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_add_binary_files(self):
        """Test that we can add binary files to the manager."""
        storage_load_spy = MagicMock(side_effect=self.storage.load_and_get_id)
        self.storage.load_and_get_id = storage_load_spy

        sample_coords = set()
        while len(sample_coords) < len(ALL_FIXTURES):
            sample_coords.add(random_coordinates())

        for sample in ALL_FIXTURES.values():
            content = sample["content"]
            assert isinstance(content, bytes)
            mimetype = sample["mimetype"]
            media_file = self._add_file_and_get_object(
                content, mimetype, sample_coords.pop()
            )
            assert media_file is not None

            # Ensure MediaFileStorage.load_and_get_id was called as expected.
            storage_load_spy.assert_called_once_with(
                content, mimetype, MediaFileKind.MEDIA, None
            )
            storage_load_spy.reset_mock()

        # There should be as many files in MFM as we added.
        assert len(self.media_file_manager._file_metadata) == len(ALL_FIXTURES)

        # There should only be 1 session with registered files.
        assert len(self.media_file_manager._files_by_session_and_coord) == 1

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    @mock.patch(
        "streamlit.runtime.memory_media_file_storage.open",
        mock_open(read_data=b"mock_test_file"),
        create=True,
    )
    def test_add_file_by_name(self):
        """Test that we can add files by filename."""
        storage_load_spy = MagicMock(side_effect=self.storage.load_and_get_id)
        self.storage.load_and_get_id = storage_load_spy

        self.media_file_manager.add(
            "mock/file/path.png", "image/png", random_coordinates()
        )

        # We should have a single file in the MFM.
        assert len(self.media_file_manager._file_metadata) == 1

        # And it should be registered to our session
        assert (
            len(self.media_file_manager._files_by_session_and_coord["mock_session"])
            == 1
        )

        # Ensure MediaFileStorage.load_and_get_id was called as expected.
        storage_load_spy.assert_called_once_with(
            "mock/file/path.png", "image/png", MediaFileKind.MEDIA, None
        )

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session_id"),
    )
    def test_add_files_same_coord(self):
        """We can add multiple files that share the same coordinate."""
        coord = random_coordinates()

        for sample in ALL_FIXTURES.values():
            self.media_file_manager.add(sample["content"], sample["mimetype"], coord)

        # There should be 6 files in MFM.
        assert len(self.media_file_manager._file_metadata) == len(ALL_FIXTURES)

        # There should only be 1 session with registered files.
        assert len(self.media_file_manager._files_by_session_and_coord) == 1

        # There should only be 1 coord in that session.
        assert (
            len(self.media_file_manager._files_by_session_and_coord["mock_session_id"])
            == 1
        )

        self.media_file_manager.clear_session_refs()
        self.media_file_manager.remove_orphaned_files()

        # There should be only 0 file in MFM.
        assert len(self.media_file_manager._file_metadata) == 0

        # There should only be 0 session with registered files.
        assert len(self.media_file_manager._files_by_session_and_coord) == 0

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session_id"),
    )
    def test_add_file_already_exists_same_coord(self):
        """Adding a file that already exists results in just a single file in
        the manager.
        """
        sample = IMAGE_FIXTURES["png"]
        coord = random_coordinates()

        self.media_file_manager.add(sample["content"], sample["mimetype"], coord)
        file_id = _calculate_file_id(sample["content"], sample["mimetype"])
        assert file_id in self.media_file_manager._file_metadata

        self.media_file_manager.add(sample["content"], sample["mimetype"], coord)
        assert file_id in self.media_file_manager._file_metadata

        # There should only be 1 file in MFM.
        assert len(self.media_file_manager._file_metadata) == 1

        # There should only be 1 session with registered files.
        assert len(self.media_file_manager._files_by_session_and_coord) == 1

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session_id"),
    )
    def test_add_file_already_exists_different_coord(self):
        """Adding a file that already exists, but with different coordinates,
        results in just a single file in the manager.
        """
        sample = IMAGE_FIXTURES["png"]

        coord = random_coordinates()
        self.media_file_manager.add(sample["content"], sample["mimetype"], coord)
        file_id = _calculate_file_id(sample["content"], sample["mimetype"])
        assert file_id in self.media_file_manager._file_metadata

        coord = random_coordinates()
        self.media_file_manager.add(sample["content"], sample["mimetype"], coord)
        assert file_id in self.media_file_manager._file_metadata

        # There should only be 1 file in MFM.
        assert len(self.media_file_manager._file_metadata) == 1

        # There should only be 1 session with registered files.
        assert len(self.media_file_manager._files_by_session_and_coord) == 1

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session_id"),
    )
    def test_remove_orphaned_files_in_empty_manager(self):
        """Calling clear_session_refs/remove_orphaned_files in an empty manager
        is a no-op.
        """
        storage_delete_spy = MagicMock(side_effect=self.storage.delete_file)
        self.storage.delete_file = storage_delete_spy

        assert len(self.media_file_manager._file_metadata) == 0
        assert len(self.media_file_manager._files_by_session_and_coord) == 0

        self.media_file_manager.clear_session_refs()
        self.media_file_manager.remove_orphaned_files()

        assert len(self.media_file_manager._file_metadata) == 0
        assert len(self.media_file_manager._files_by_session_and_coord) == 0

        # MediaFileStorage.delete_file should not have been called, because
        # no files were actually deleted.
        storage_delete_spy.assert_not_called()

    @mock.patch("streamlit.runtime.media_file_manager._get_session_id")
    def test_remove_orphaned_files_multiple_sessions(self, mock_get_session_id):
        """clear_session_refs/remove_orphaned_files behaves correctly when multiple
        sessions are referencing some of the same files.
        """
        storage_delete_spy = MagicMock(side_effect=self.storage.delete_file)
        self.storage.delete_file = storage_delete_spy

        # Have two sessions add the same set of files
        for session_id in ("mock_session_1", "mock_session_2"):
            mock_get_session_id.return_value = session_id
            for sample in VIDEO_FIXTURES.values():
                coord = random_coordinates()
                self.media_file_manager.add(
                    sample["content"], sample["mimetype"], coord
                )

        assert len(self.media_file_manager._file_metadata) == len(VIDEO_FIXTURES)

        file_ids = list(self.media_file_manager._file_metadata.keys())

        # Remove session1's references
        mock_get_session_id.return_value = "mock_session_1"
        self.media_file_manager.clear_session_refs()
        self.media_file_manager.remove_orphaned_files()

        # The files are all still referenced by session_2
        assert len(self.media_file_manager._file_metadata) == len(VIDEO_FIXTURES)

        # MediaFileStorage.delete_file should not have been called yet...
        storage_delete_spy.assert_not_called()

        # Remove session2's references, but don't call "remove_orphaned_files" yet...
        mock_get_session_id.return_value = "mock_session_2"
        self.media_file_manager.clear_session_refs()

        # The files still exist, because they've only been de-referenced and not
        # removed.
        assert len(self.media_file_manager._file_metadata) == len(VIDEO_FIXTURES)

        # MediaFileStorage.delete_file should not have been called yet...
        storage_delete_spy.assert_not_called()

        # After a final call to remove_orphaned_files, the files should be gone.
        self.media_file_manager.remove_orphaned_files()
        assert len(self.media_file_manager._file_metadata) == 0

        # MediaFileStorage.delete_file should have been called once for each
        # file.
        storage_delete_spy.assert_has_calls(
            [call(file_id) for file_id in file_ids], any_order=True
        )


class MediaFileManagerThreadingTest(unittest.TestCase):
    # The number of threads to run our tests on
    NUM_THREADS = 50

    def setUp(self):
        super().setUp()
        self.storage = MemoryMediaFileStorage("/mock/endpoint")
        self.media_file_manager = MediaFileManager(self.storage)
        random.seed(1337)

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session_id"),
    )
    def test_add_file_multiple_threads(self):
        """We can safely call `add` from multiple threads simultaneously."""

        def add_file(ii: int) -> None:
            coord = random_coordinates()
            data = bytes(f"{ii}", "utf-8")
            self.media_file_manager.add(data, "image/png", coord)

        call_on_threads(add_file, num_threads=self.NUM_THREADS)
        assert len(self.media_file_manager._file_metadata) == self.NUM_THREADS

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session_id"),
    )
    def test_clear_files_multiple_threads(self):
        """We can safely clear session refs and remove orphaned files
        from multiple threads simultaneously.
        """
        # Add a bunch of files
        for sample in ALL_FIXTURES.values():
            self.media_file_manager.add(
                sample["content"], sample["mimetype"], random_coordinates()
            )
        assert len(ALL_FIXTURES) == len(self.media_file_manager._file_metadata)

        # Remove those files from multiple threads
        def remove_files(_: int) -> None:
            self.media_file_manager.clear_session_refs("mock_session_id")
            self.media_file_manager.remove_orphaned_files()

        call_on_threads(remove_files, num_threads=self.NUM_THREADS)

        # Our files should be gone!
        assert len(self.media_file_manager._file_metadata) == 0


class MediaFileManagerDeferredTest(TestCase):
    """Tests for deferred callable functionality in MediaFileManager."""

    def setUp(self):
        super().setUp()
        self.storage = MemoryMediaFileStorage("/mock/endpoint")
        self.media_file_manager = MediaFileManager(self.storage)

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_add_deferred_generates_file_id(self):
        """Test that add_deferred generates a unique file_id."""

        def callable1():
            return b"content1"

        def callable2():
            return b"content2"

        file_id1 = self.media_file_manager.add_deferred(
            callable1, "text/plain", random_coordinates()
        )
        file_id2 = self.media_file_manager.add_deferred(
            callable2, "text/plain", random_coordinates()
        )

        # File IDs should be different
        assert file_id1 != file_id2
        # File IDs should be non-empty
        assert file_id1 != ""
        assert file_id2 != ""

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_add_deferred_stores_callable(self):
        """Test that add_deferred stores the callable with metadata."""

        def generate_data():
            return b"test data"

        file_id = self.media_file_manager.add_deferred(
            generate_data, "application/pdf", random_coordinates(), file_name="test.pdf"
        )

        # Callable should be stored
        assert file_id in self.media_file_manager._deferred_callables
        deferred = self.media_file_manager._deferred_callables[file_id]
        assert deferred["callable"] == generate_data
        assert deferred["mimetype"] == "application/pdf"
        assert deferred["filename"] == "test.pdf"

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_add_deferred_tracks_session_coordinate(self):
        """Test that add_deferred tracks session and coordinate mapping."""
        coord = random_coordinates()

        def generate_data():
            return b"data"

        file_id = self.media_file_manager.add_deferred(
            generate_data, "text/plain", coord
        )

        # Should be tracked by session and coordinate
        assert (
            self.media_file_manager._files_by_session_and_coord["mock_session"][coord]
            == file_id
        )

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_calls_callable(self):
        """Test that execute_deferred invokes the callable."""
        call_count = 0

        def generate_data():
            nonlocal call_count
            call_count += 1
            return b"data"

        file_id = self.media_file_manager.add_deferred(
            generate_data, "text/plain", random_coordinates()
        )

        assert call_count == 0
        self.media_file_manager.execute_deferred(file_id)
        assert call_count == 1

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_returns_url(self):
        """Test that execute_deferred returns a valid URL."""

        def generate_data():
            return b"test data"

        file_id = self.media_file_manager.add_deferred(
            generate_data, "text/plain", random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_handles_str_return(self):
        """Test that execute_deferred handles string return values."""

        def generate_string():
            return "string data"

        file_id = self.media_file_manager.add_deferred(
            generate_string, "text/plain", random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_handles_bytes_return(self):
        """Test that execute_deferred handles bytes return values."""

        def generate_bytes():
            return b"bytes data"

        file_id = self.media_file_manager.add_deferred(
            generate_bytes, "application/octet-stream", random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_handles_bytesio_return(self):
        """Test that execute_deferred handles BytesIO return values."""

        def generate_bytesio():
            return io.BytesIO(b"bytesio data")

        file_id = self.media_file_manager.add_deferred(
            generate_bytesio, "application/octet-stream", random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_handles_buffered_reader_return(self):
        """Test that execute_deferred handles BufferedReader return values."""

        def generate_buffered_reader():
            return io.BufferedReader(io.BytesIO(b"buffered data"))

        file_id = self.media_file_manager.add_deferred(
            generate_buffered_reader, "application/octet-stream", random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_keeps_callable(self):
        """Test that execute_deferred keeps callable for multiple downloads."""

        def generate_data():
            return b"data"

        file_id = self.media_file_manager.add_deferred(
            generate_data, "text/plain", random_coordinates()
        )

        # Callable should exist before execution
        assert file_id in self.media_file_manager._deferred_callables

        self.media_file_manager.execute_deferred(file_id)

        # Callable should still exist after execution for multiple downloads
        assert file_id in self.media_file_manager._deferred_callables

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_file_not_found(self):
        """Test that execute_deferred raises error for non-existent file_id."""
        with pytest.raises(MediaFileStorageError, match=r"Deferred file .* not found"):
            self.media_file_manager.execute_deferred("nonexistent_id")

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_callable_raises_exception(self):
        """Test that execute_deferred propagates callable execution errors."""

        def failing_callable():
            raise ValueError("Test error")

        file_id = self.media_file_manager.add_deferred(
            failing_callable, "text/plain", random_coordinates()
        )

        with pytest.raises(MediaFileStorageError, match="Callable execution failed"):
            self.media_file_manager.execute_deferred(file_id)

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_callable_returns_invalid_type(self):
        """Test that execute_deferred handles invalid return types."""

        def invalid_callable():
            return 123  # Invalid type

        file_id = self.media_file_manager.add_deferred(
            invalid_callable, "text/plain", random_coordinates()
        )

        with pytest.raises(
            MediaFileStorageError, match="Callable returned unsupported type"
        ):
            self.media_file_manager.execute_deferred(file_id)

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_multiple_times_same_callable(self):
        """Test that the same deferred callable can be executed multiple times."""
        call_count = 0

        def generate_data():
            nonlocal call_count
            call_count += 1
            return f"data_{call_count}".encode()

        file_id = self.media_file_manager.add_deferred(
            generate_data, "text/plain", random_coordinates()
        )

        # First execution should work
        url1 = self.media_file_manager.execute_deferred(file_id)
        assert url1.startswith("/mock/endpoint/")
        assert call_count == 1

        # Second execution should also work (callable is kept)
        url2 = self.media_file_manager.execute_deferred(file_id)
        assert url2.startswith("/mock/endpoint/")
        assert call_count == 2

        # URLs can be different if the callable returns different data
        # but both should be valid

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_clear_session_refs_and_remove_orphaned_cleans_deferred(self):
        """Test that deferred callables are cleaned up after clear_session_refs + remove_orphaned_files."""

        def generate_data():
            return b"data"

        # Add a deferred callable
        file_id = self.media_file_manager.add_deferred(
            generate_data, "text/plain", random_coordinates()
        )

        # Callable should exist
        assert file_id in self.media_file_manager._deferred_callables
        assert len(self.media_file_manager._deferred_callables) == 1

        # Clear session refs (doesn't immediately delete callables)
        self.media_file_manager.clear_session_refs("mock_session")

        # Callable should still exist (not immediately deleted to avoid race conditions)
        assert file_id in self.media_file_manager._deferred_callables
        assert len(self.media_file_manager._deferred_callables) == 1

        # Remove orphaned files (this cleans up orphaned deferred callables)
        self.media_file_manager.remove_orphaned_files()

        # Now the deferred callable should be removed
        assert file_id not in self.media_file_manager._deferred_callables
        assert len(self.media_file_manager._deferred_callables) == 0

    @mock.patch("streamlit.runtime.media_file_manager._get_session_id")
    def test_remove_orphaned_only_cleans_unreferenced_deferred(
        self, mock_get_session_id
    ):
        """Test that only truly orphaned deferred callables are removed."""

        def generate_data():
            return b"data"

        # Add deferred callable for session 1
        mock_get_session_id.return_value = "session_1"
        file_id_1 = self.media_file_manager.add_deferred(
            generate_data, "text/plain", random_coordinates()
        )

        # Add deferred callable for session 2
        mock_get_session_id.return_value = "session_2"
        file_id_2 = self.media_file_manager.add_deferred(
            generate_data, "text/plain", random_coordinates()
        )

        # Both callables should exist
        assert file_id_1 in self.media_file_manager._deferred_callables
        assert file_id_2 in self.media_file_manager._deferred_callables
        assert len(self.media_file_manager._deferred_callables) == 2

        # Clear session 1 refs (doesn't immediately delete)
        self.media_file_manager.clear_session_refs("session_1")

        # Both callables still exist
        assert file_id_1 in self.media_file_manager._deferred_callables
        assert file_id_2 in self.media_file_manager._deferred_callables

        # Remove orphaned files - only session 1's callable should be cleaned
        self.media_file_manager.remove_orphaned_files()

        assert file_id_1 not in self.media_file_manager._deferred_callables
        assert file_id_2 in self.media_file_manager._deferred_callables
        assert len(self.media_file_manager._deferred_callables) == 1

        # Clear session 2 refs and remove orphans
        self.media_file_manager.clear_session_refs("session_2")
        self.media_file_manager.remove_orphaned_files()

        # Now both should be removed
        assert file_id_2 not in self.media_file_manager._deferred_callables
        assert len(self.media_file_manager._deferred_callables) == 0

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_handles_text_io_wrapper_return(self):
        """Test that execute_deferred handles TextIOWrapper (text stream) return values."""

        def generate_text_wrapper():
            # Create a TextIOWrapper over BytesIO containing UTF-8 text
            byte_stream = io.BytesIO(b"wrapped text")
            return io.TextIOWrapper(byte_stream, encoding="utf-8")

        file_id = self.media_file_manager.add_deferred(
            generate_text_wrapper, "text/plain", random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_infers_text_plain_for_string_when_mimetype_none(self):
        """If mimetype is None, infer text/plain for str returns."""

        def generate_text():
            return "hello world"

        file_id = self.media_file_manager.add_deferred(
            generate_text, None, random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")
        assert url.endswith(".txt")

        # Verify stored mimetype is text/plain
        filename = url.split("/")[-1]
        stored = self.storage.get_file(filename)
        assert stored.mimetype == "text/plain"

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_respects_provided_mimetype_over_inferred(self):
        """Test that provided mimetype is used even when data type suggests different."""

        def generate_text():
            return "hello world"

        # Even though data is string (would infer text/plain), use provided mimetype
        file_id = self.media_file_manager.add_deferred(
            generate_text, "text/csv", random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")

        # Verify stored mimetype is text/csv (the provided one)
        filename = url.split("/")[-1]
        stored = self.storage.get_file(filename)
        assert stored.mimetype == "text/csv"

    @mock.patch(
        "streamlit.runtime.media_file_manager._get_session_id",
        MagicMock(return_value="mock_session"),
    )
    def test_execute_deferred_infers_octet_stream_for_bytes_when_mimetype_none(self):
        """If mimetype is None, infer application/octet-stream for bytes returns."""

        def generate_bytes():
            return b"binary data"

        file_id = self.media_file_manager.add_deferred(
            generate_bytes, None, random_coordinates()
        )

        url = self.media_file_manager.execute_deferred(file_id)
        assert url.startswith("/mock/endpoint/")

        # Verify stored mimetype is application/octet-stream
        filename = url.split("/")[-1]
        stored = self.storage.get_file(filename)
        assert stored.mimetype == "application/octet-stream"
