# 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.

"""Unit tests for UploadedFileManager"""

from __future__ import annotations

import unittest

from streamlit.runtime.memory_uploaded_file_manager import MemoryUploadedFileManager
from streamlit.runtime.stats import CacheStat
from streamlit.runtime.uploaded_file_manager import UploadedFileRec
from tests.exception_capturing_thread import call_on_threads

FILE_1 = UploadedFileRec(file_id="url1", name="file1", type="type", data=b"file1")
FILE_2 = UploadedFileRec(file_id="url2", name="file2", type="type", data=b"file222")


class UploadedFileManagerTest(unittest.TestCase):
    def setUp(self):
        self.mgr = MemoryUploadedFileManager("/mock/upload")

    def test_added_file_id(self):
        """Presigned file URL should have a unique ID."""
        info1, info2 = self.mgr.get_upload_urls("session", ["name1", "name1"])
        assert info1.file_id != info2.file_id

    def test_retrieve_added_file(self):
        """An added file should maintain all its source properties
        except its ID."""
        self.mgr.add_file("session", FILE_1)
        self.mgr.add_file("session", FILE_2)

        file1_from_storage, *rest_files = self.mgr.get_files("session", ["url1"])
        assert len(rest_files) == 0
        assert file1_from_storage.file_id == FILE_1.file_id
        assert file1_from_storage.name == FILE_1.name
        assert file1_from_storage.type == FILE_1.type
        assert file1_from_storage.data == FILE_1.data

        file2_from_storage, *other_files = self.mgr.get_files("session", ["url2"])
        assert len(other_files) == 0
        assert file2_from_storage.file_id == FILE_2.file_id
        assert file2_from_storage.name == FILE_2.name
        assert file2_from_storage.type == FILE_2.type
        assert file2_from_storage.data == FILE_2.data

    def test_remove_file(self):
        # This should not error.
        self.mgr.remove_file("non-session", "non-file-id")

        self.mgr.add_file("session", FILE_1)
        self.mgr.remove_file("session", FILE_1.file_id)
        assert self.mgr.get_files("session", [FILE_1.file_id]) == []

        # Remove the file again. It doesn't exist, but this isn't an error.
        self.mgr.remove_file("session", FILE_1.file_id)
        assert self.mgr.get_files("session", [FILE_1.file_id]) == []

        self.mgr.add_file("session", FILE_1)
        self.mgr.add_file("session", FILE_2)
        self.mgr.remove_file("session", FILE_1.file_id)
        assert self.mgr.get_files("session", [FILE_1.file_id, FILE_2.file_id]) == [
            FILE_2
        ]

    def test_remove_session_files(self):
        # This should not error.
        self.mgr.remove_session_files("non-report")

        # Add two files with different session IDs, but the same widget ID.
        self.mgr.add_file("session1", FILE_1)
        self.mgr.add_file("session1", FILE_2)

        self.mgr.add_file("session2", FILE_1)

        self.mgr.remove_session_files("session1")
        assert self.mgr.get_files("session1", [FILE_1.file_id, FILE_2.file_id]) == []
        assert self.mgr.get_files("session2", [FILE_1.file_id]) == [FILE_1]

    def test_cache_stats_provider(self):
        """Test StatsProvider implementation."""

        # Test empty manager
        assert self.mgr.get_stats() == {}

        # Test manager with files
        self.mgr.add_file("session1", FILE_1)
        self.mgr.add_file("session1", FILE_2)

        expected_stat = CacheStat(
            category_name="UploadedFileManager",
            cache_name="",
            byte_length=len(FILE_1.data) + len(FILE_2.data),
        )
        expected = {expected_stat.family_name: [expected_stat]}
        assert expected == self.mgr.get_stats()


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

    def setUp(self) -> None:
        self.mgr = MemoryUploadedFileManager("/mock/upload")

    def test_add_file(self):
        """`add_file` is thread-safe."""
        # Call `add_file` from a bunch of threads
        added_files = []

        def add_file(index: int) -> None:
            file = UploadedFileRec(
                file_id=f"id_{index}",
                name=f"file_{index}",
                type="type",
                data=bytes(f"{index}", "utf-8"),
            )

            self.mgr.add_file("session", file)
            files_from_storage = self.mgr.get_files("session", [file.file_id])
            added_files.extend(files_from_storage)

        call_on_threads(add_file, num_threads=self.NUM_THREADS)

        # Ensure all our files are present
        for ii in range(self.NUM_THREADS):
            files = self.mgr.get_files("session", [f"id_{ii}"])
            assert len(files) == 1
            assert bytes(f"{ii}", "utf-8") == files[0].data

        # Ensure all files have unique IDs
        file_ids = set()
        file_ids.update(
            file_rec.file_id for file_rec in self.mgr.file_storage["session"].values()
        )
        assert len(file_ids) == self.NUM_THREADS

    def test_remove_file(self):
        """`remove_file` is thread-safe."""
        # Add a bunch of files to a single widget
        file_ids = []
        for ii in range(self.NUM_THREADS):
            file = UploadedFileRec(
                file_id=f"id_{ii}",
                name=f"file_{ii}",
                type="type",
                data=b"123",
            )
            self.mgr.add_file("session", file)
            file_ids.append(file.file_id)

        # Have each thread remove a single file
        def remove_file(index: int) -> None:
            file_id = file_ids[index]

            # Ensure our file exists
            get_files_result = self.mgr.get_files("session", [file_id])
            assert len(get_files_result) == 1

            # Remove our file and ensure our file no longer exists
            self.mgr.remove_file("session", file_id)
            get_files_result = self.mgr.get_files("session", [file_id])
            assert len(get_files_result) == 0

        call_on_threads(remove_file, self.NUM_THREADS)

        assert len(self.mgr.file_storage["session"]) == 0

    def test_remove_session_files(self):
        """`remove_session_files` is thread-safe."""
        # Add a bunch of files, each to a different session
        file_ids = []
        for ii in range(self.NUM_THREADS):
            file = UploadedFileRec(
                file_id=f"id_{ii}",
                name=f"file_{ii}",
                type="type",
                data=b"123",
            )
            self.mgr.add_file(f"session_{ii}", file)
            file_ids.append(file.file_id)

        # Have each thread remove its session's file
        def remove_session_files(index: int) -> None:
            session_id = f"session_{index}"
            # Our file should exist
            session_files = self.mgr.get_files(session_id, [f"id_{index}"])
            assert len(session_files) == 1
            assert file_ids[index] == session_files[0].file_id

            # Remove session files
            self.mgr.remove_session_files(session_id)

            # Our file should no longer exist
            session_files = self.mgr.get_files(session_id, [f"id_{index}"])
            assert len(session_files) == 0

        call_on_threads(remove_session_files, num_threads=self.NUM_THREADS)
