import base64
import io
import unittest
from datetime import datetime
from unittest.mock import MagicMock, patch

import pytest
from io_storages.proxy_api import (
    ProjectResolveStorageUri,
    ResolveStorageUriAPIMixin,
    TaskResolveStorageUri,
)
from projects.models import Project
from rest_framework import status
from rest_framework.response import Response
from rest_framework.test import APIRequestFactory, force_authenticate
from tasks.models import Task


class TestResolveStorageUriAPIMixin(unittest.TestCase):
    def setUp(self):
        self.mixin = ResolveStorageUriAPIMixin()
        self.user = MagicMock()
        self.project = MagicMock()
        self.task = MagicMock()
        self.task.project = self.project
        # Set the __class__.__name__ to "Task" for type checks
        type(self.task).__name__ = 'Task'
        self.task.has_permission.return_value = True
        self.request = MagicMock()
        self.request.user = self.user
        self.storage = MagicMock()
        self.storage.presign = True

    def test_resolve_with_permission_denied(self):
        self.task.has_permission.return_value = False
        result = self.mixin.resolve(self.request, 'test_fileuri', self.task)
        assert result.status_code == status.HTTP_403_FORBIDDEN

    @patch('io_storages.proxy_api.flag_set')
    @patch('io_storages.proxy_api.get_storage_by_url')
    def test_resolve_with_base64_decoding(self, mock_get_storage, mock_flag_set):
        mock_flag_set.return_value = True
        mock_get_storage.return_value = self.storage
        fileuri = base64.urlsafe_b64encode(b'test_uri').decode()

        with patch.object(self.mixin, 'redirect_to_presign_url') as mock_redirect:
            mock_redirect.return_value = Response()
            self.mixin.resolve(self.request, fileuri, self.task)
            mock_redirect.assert_called_once_with('test_uri', self.task, 'Task')

    @patch('io_storages.proxy_api.flag_set')
    @patch('io_storages.proxy_api.get_storage_by_url')
    def test_resolve_with_url_unquote_fallback(self, mock_get_storage, mock_flag_set):
        mock_flag_set.return_value = True
        mock_get_storage.return_value = self.storage

        with patch.object(self.mixin, 'redirect_to_presign_url') as mock_redirect:
            mock_redirect.return_value = Response()
            # Non-base64 uri to trigger fallback
            self.mixin.resolve(self.request, 's3://bucket/file.jpg', self.task)
            mock_redirect.assert_called_once_with('s3://bucket/file.jpg', self.task, 'Task')

    @patch('io_storages.proxy_api.flag_set')
    @patch('io_storages.proxy_api.get_storage_by_url')
    def test_resolve_storage_not_found(self, mock_get_storage, mock_flag_set):
        mock_flag_set.return_value = True
        mock_get_storage.return_value = None
        result = self.mixin.resolve(self.request, 'fileuri', self.task)
        assert result.status_code == status.HTTP_404_NOT_FOUND

    @patch('io_storages.proxy_api.flag_set')
    @patch('io_storages.proxy_api.get_storage_by_url')
    def test_resolve_storage_no_presign_support(self, mock_get_storage, mock_flag_set):
        mock_flag_set.return_value = True
        mock_storage = MagicMock()
        delattr(mock_storage, 'presign')
        mock_get_storage.return_value = mock_storage
        result = self.mixin.resolve(self.request, 'fileuri', self.task)
        assert result.status_code == status.HTTP_404_NOT_FOUND

    @patch('io_storages.proxy_api.flag_set')
    @patch('io_storages.proxy_api.get_storage_by_url')
    def test_resolve_with_presign_true(self, mock_get_storage, mock_flag_set):
        mock_flag_set.return_value = True
        mock_storage = MagicMock()
        mock_storage.presign = True
        mock_get_storage.return_value = mock_storage

        with patch.object(self.mixin, 'redirect_to_presign_url') as mock_redirect:
            mock_redirect.return_value = Response()
            self.mixin.resolve(self.request, 'fileuri', self.task)
            mock_redirect.assert_called_once()

    @patch('io_storages.proxy_api.flag_set')
    @patch('io_storages.proxy_api.get_storage_by_url')
    def test_resolve_with_presign_false(self, mock_get_storage, mock_flag_set):
        mock_flag_set.return_value = True
        mock_storage = MagicMock()
        mock_storage.presign = False
        mock_get_storage.return_value = mock_storage
        project = self.task.project

        with patch.object(self.mixin, 'proxy_data_from_storage') as mock_proxy:
            mock_proxy.return_value = Response()
            self.mixin.resolve(self.request, 'fileuri', self.task)
            mock_proxy.assert_called_once_with(self.request, 'fileuri', project, mock_storage)

    def test_redirect_to_presign_url_success(self):
        self.task.resolve_storage_uri.return_value = {'url': 'https://example.com/file.jpg', 'presign_ttl': 60}
        result = self.mixin.redirect_to_presign_url('fileuri', self.task, 'Task')

        assert result.status_code == status.HTTP_303_SEE_OTHER
        assert result.url == 'https://example.com/file.jpg'
        assert result.headers['Cache-Control'] == 'no-store, max-age=3600'

    def test_redirect_to_presign_url_no_url(self):
        self.task.resolve_storage_uri.return_value = {'url': None}
        result = self.mixin.redirect_to_presign_url('fileuri', self.task, 'Task')
        assert result.status_code == status.HTTP_404_NOT_FOUND

    def test_redirect_to_presign_url_exception(self):
        self.task.resolve_storage_uri.side_effect = Exception('Error resolving URL')
        result = self.mixin.redirect_to_presign_url('fileuri', self.task, 'Task')
        assert result.status_code == status.HTTP_404_NOT_FOUND

    def test_proxy_data_from_storage_success(self):
        mock_storage = MagicMock()
        # Ensure get_bytes_stream returns a three-tuple, metadata can be empty initially
        mock_storage.get_bytes_stream.return_value = (io.BytesIO(b'test data'), 'image/jpeg', {})
        mock_project = MagicMock()

        with patch('io_storages.proxy_api.StreamingHttpResponse') as mock_response_class, patch(
            'io_storages.proxy_api.settings'
        ) as mock_settings:
            # Configure mock settings
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 1024 * 1024  # 1MB
            mock_settings.RESOLVER_PROXY_BUFFER_SIZE = 8192
            mock_settings.RESOLVER_PROXY_CACHE_TIMEOUT = 3600

            # Set up mock stream and response
            mock_stream = MagicMock()
            mock_metadata = {
                'StatusCode': 200,
                'ContentLength': 1000,
                'LastModified': datetime.now(),
                'ETag': '"abcdef123456"',
            }
            mock_storage.get_bytes_stream.return_value = (mock_stream, 'application/test', mock_metadata)

            # Set up mock response
            mock_response = MagicMock()
            mock_response.headers = {}
            mock_response_class.return_value = mock_response

            # Set up request with range header
            self.request.headers = {'Range': 'bytes=100-200'}

            # Call the method
            result = self.mixin.proxy_data_from_storage(self.request, 'uri', mock_project, mock_storage)

            # Verify the correct range header was passed
            mock_storage.get_bytes_stream.assert_called_once()
            args, kwargs = mock_storage.get_bytes_stream.call_args
            # First positional argument should be the URI
            self.assertEqual(args[0], 'uri')
            # Range header should be passed as a keyword argument and start with 'bytes='
            self.assertIn('range_header', kwargs)
            self.assertTrue(kwargs['range_header'].startswith('bytes='))

            # Verify the response was created with the stream
            mock_response_class.assert_called_once()
            # The first positional argument should be a generator returned by time_limited_chunker
            called_args, _ = mock_response_class.call_args
            streaming_generator = called_args[0]
            self.assertTrue(
                hasattr(streaming_generator, '__next__') and hasattr(streaming_generator, '__iter__'),
                'Expected a generator for streaming chunks',
            )

            # Verify correct headers are set
            self.assertEqual(result, mock_response)
            self.assertTrue('ETag' in mock_response.headers)

    def test_proxy_data_from_storage_no_data(self):
        mock_storage = MagicMock()
        # Return three-tuple with empty metadata when no data is available
        mock_storage.get_bytes_stream.return_value = (None, None, {})
        mock_project = MagicMock()

        result = self.mixin.proxy_data_from_storage(self.request, 'uri', mock_project, mock_storage)
        assert result.status_code == status.HTTP_424_FAILED_DEPENDENCY

    def test_proxy_data_from_storage_exception(self):
        mock_storage = MagicMock()
        mock_storage.get_bytes_stream.side_effect = Exception('Storage error')
        mock_project = MagicMock()

        result = self.mixin.proxy_data_from_storage(self.request, 'uri', mock_project, mock_storage)
        assert result.status_code == status.HTTP_424_FAILED_DEPENDENCY

    def test_time_limited_chunker_normal_case(self):
        """Test time_limited_chunker when all chunks are processed within timeout"""
        # Create a mock stream with iter_chunks method
        mock_stream = MagicMock()
        mock_stream.iter_chunks.return_value = [b'chunk1', b'chunk2', b'chunk3']

        # Set up settings
        with patch('io_storages.proxy_api.settings') as mock_settings, patch('time.monotonic') as mock_time:
            # Mock settings
            mock_settings.RESOLVER_PROXY_BUFFER_SIZE = 8192
            mock_settings.RESOLVER_PROXY_TIMEOUT = 20

            # Mock time to simulate being within timeout
            # Add an extra value for the 'finally' block
            mock_time.side_effect = [0, 1, 2, 3, 4]  # Start, three chunk iterations, finally

            # Run the chunker and collect all chunks
            chunks = list(self.mixin.time_limited_chunker(mock_stream))

            # Verify all chunks were yielded
            assert chunks == [b'chunk1', b'chunk2', b'chunk3']
            assert mock_stream.close.called

    def test_time_limited_chunker_timeout(self):
        """Test time_limited_chunker when timeout is reached during processing"""
        # Create a mock stream with iter_chunks method
        mock_stream = MagicMock()
        mock_stream.iter_chunks.return_value = [b'chunk1', b'chunk2', b'chunk3', b'chunk4', b'chunk5']

        # Set up settings
        with patch('io_storages.proxy_api.settings') as mock_settings, patch('time.monotonic') as mock_time:
            # Mock settings
            mock_settings.RESOLVER_PROXY_BUFFER_SIZE = 8192
            mock_settings.RESOLVER_PROXY_TIMEOUT = 10

            # Mock time to simulate exceeding timeout after second chunk
            # Add an extra value for the 'finally' block
            mock_time.side_effect = [0, 5, 9, 15, 20, 25]  # Start time, chunk checks, finally

            # Run the chunker and collect all chunks
            chunks = list(self.mixin.time_limited_chunker(mock_stream))

            # Verify only the chunks before timeout were yielded
            assert chunks == [b'chunk1', b'chunk2']
            assert mock_stream.close.called

    def test_time_limited_chunker_exception(self):
        """Test time_limited_chunker when an exception occurs during streaming"""
        # Create a mock stream with iter_chunks method that raises an exception
        mock_stream = MagicMock()
        mock_stream.iter_chunks.side_effect = Exception('Streaming error')

        # Set up settings
        with patch('io_storages.proxy_api.settings') as mock_settings, patch('time.monotonic') as mock_time:
            # Mock settings
            mock_settings.RESOLVER_PROXY_BUFFER_SIZE = 8192
            mock_settings.RESOLVER_PROXY_TIMEOUT = 20

            # Mock time - need two values: one for start and one for the finally block
            mock_time.side_effect = [0, 1]

            # Run the chunker and collect all chunks (should be empty due to exception)
            chunks = list(self.mixin.time_limited_chunker(mock_stream))

            # Verify no chunks were yielded and the stream was closed
            assert chunks == []
            assert mock_stream.close.called

    def test_override_range_header_no_header(self):
        """Test override_range_header when no Range header is present"""
        self.request.headers = {}
        result = self.mixin.override_range_header(self.request)
        assert result is None

    def test_override_range_header_header_probes(self):
        """Test override_range_header with header probe formats"""
        # Test bytes=0-
        self.request.headers = {'Range': 'bytes=0-'}
        with patch('io_storages.proxy_api.settings') as mock_settings, patch(
            'io_storages.proxy_api.parse_range'
        ) as mock_parse_range:
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 1024 * 1024
            # Mock the parse_range function to return a known value
            mock_parse_range.return_value = (0, '')

            result = self.mixin.override_range_header(self.request)
            assert result == 'bytes=0-'

        # Test bytes=0-0
        self.request.headers = {'Range': 'bytes=0-0'}
        with patch('io_storages.proxy_api.settings') as mock_settings, patch(
            'io_storages.proxy_api.parse_range'
        ) as mock_parse_range:
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 1024 * 1024
            # Mock the parse_range function to return a known value
            mock_parse_range.return_value = (0, 0)

            result = self.mixin.override_range_header(self.request)
            assert result == 'bytes=0-0'

    def test_override_range_header_start_no_end(self):
        """Test override_range_header with a start position but no end"""
        # Case: bytes=100-
        self.request.headers = {'Range': 'bytes=100-'}
        with patch('io_storages.proxy_api.settings') as mock_settings, patch(
            'io_storages.proxy_api.parse_range'
        ) as mock_parse_range:
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 1024 * 1024  # 1MB
            # Mock the parse_range function to return a known value
            mock_parse_range.return_value = (100, '')

            result = self.mixin.override_range_header(self.request)
            # Should add MAX_RANGE_SIZE to start
            assert result == f'bytes=100-{100 + 1024*1024}'

        # Case: bytes=100-0 (treated like bytes=100-)
        self.request.headers = {'Range': 'bytes=100-0'}
        with patch('io_storages.proxy_api.settings') as mock_settings, patch(
            'io_storages.proxy_api.parse_range'
        ) as mock_parse_range:
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 1024 * 1024  # 1MB
            # Mock the parse_range function to return a known value
            mock_parse_range.return_value = (100, 0)

            result = self.mixin.override_range_header(self.request)
            # Should add MAX_RANGE_SIZE to start
            assert result == f'bytes=100-{100 + 1024*1024}'

    def test_override_range_header_start_and_end(self):
        """Test override_range_header with start and end positions"""
        # Case: Range within limit
        self.request.headers = {'Range': 'bytes=100-5000'}
        with patch('io_storages.proxy_api.settings') as mock_settings, patch(
            'io_storages.proxy_api.parse_range'
        ) as mock_parse_range:
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 10000  # 10KB
            # Mock the parse_range function to return a known value
            mock_parse_range.return_value = (100, 5000)

            result = self.mixin.override_range_header(self.request)
            # Should remain unchanged as it's within limit
            assert result == 'bytes=100-5000'

        # Case: Range exceeding limit
        self.request.headers = {'Range': 'bytes=100-20000'}
        with patch('io_storages.proxy_api.settings') as mock_settings, patch(
            'io_storages.proxy_api.parse_range'
        ) as mock_parse_range:
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 10000  # 10KB
            # Mock the parse_range function to return a known value
            mock_parse_range.return_value = (100, 20000)

            result = self.mixin.override_range_header(self.request)
            # Should limit the range to MAX_RANGE_SIZE from start
            assert result == f'bytes=100-{100 + 10000}'

    def test_override_range_header_negative_start(self):
        """Test override_range_header with negative start position"""
        self.request.headers = {'Range': 'bytes=-1024'}
        with patch('io_storages.proxy_api.settings') as mock_settings, patch(
            'io_storages.proxy_api.parse_range'
        ) as mock_parse_range:
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 10000  # 10KB
            # Mock the parse_range function to return a negative start
            mock_parse_range.return_value = (-1024, None)

            result = self.mixin.override_range_header(self.request)
            # Should reset to 0 and add MAX_RANGE_SIZE
            assert result == f'bytes=0-{10000}'

    def test_override_range_header_unsupported_format(self):
        """Test override_range_header with unsupported range format"""
        self.request.headers = {'Range': 'invalid-range-format'}
        with patch('io_storages.proxy_api.settings') as mock_settings, patch(
            'io_storages.proxy_api.parse_range'
        ) as mock_parse_range:
            mock_settings.RESOLVER_PROXY_MAX_RANGE_SIZE = 10000  # 10KB
            # Mock parse_range to simulate failure with invalid format
            mock_parse_range.return_value = (0, None)

            result = self.mixin.override_range_header(self.request)
            # Should reset to default
            assert result == 'bytes=0-'


class TestTaskResolveStorageUri:
    @pytest.fixture
    def setup(self):
        # Create the necessary objects for testing without database
        self.factory = APIRequestFactory()
        self.user = MagicMock()
        self.task = MagicMock(spec=Task)
        self.view = TaskResolveStorageUri.as_view()

    @patch('io_storages.proxy_api.Task.objects.get')
    def test_get_with_missing_params(self, mock_task_get, setup):
        # Mock the database query
        mock_task_get.return_value = self.task

        # Test missing fileuri parameter
        request = self.factory.get('/task/1/resolve/')
        force_authenticate(request, user=self.user)
        response = self.view(request, task_id=1)
        assert response.status_code == status.HTTP_400_BAD_REQUEST

        # Test missing task_id parameter
        request = self.factory.get('/task/resolve/?fileuri=test')
        force_authenticate(request, user=self.user)
        response = self.view(request)
        assert response.status_code == status.HTTP_400_BAD_REQUEST

    @patch('io_storages.proxy_api.Task.objects.get')
    def test_get_task_not_found(self, mock_task_get, setup):
        # Mock the database query to raise DoesNotExist
        mock_task_get.side_effect = Task.DoesNotExist

        request = self.factory.get('/task/999/resolve/?fileuri=test')
        force_authenticate(request, user=self.user)
        response = self.view(request, task_id=999)

        assert response.status_code == status.HTTP_404_NOT_FOUND

    @patch('io_storages.proxy_api.Task.objects.get')
    @patch.object(ResolveStorageUriAPIMixin, 'resolve')
    def test_get_success(self, mock_resolve, mock_task_get, setup):
        # Mock the database query and resolve method
        mock_task_get.return_value = self.task
        mock_resolve.return_value = Response(status=status.HTTP_200_OK)

        request = self.factory.get('/task/1/resolve/?fileuri=test')
        force_authenticate(request, user=self.user)
        response = self.view(request, task_id=1)

        mock_task_get.assert_called_once_with(pk=1)
        # Use any_call instead of assert_called_once_with to handle DRF request vs WSGIRequest
        assert mock_resolve.call_args is not None
        assert mock_resolve.call_args[0][1] == 'test'
        assert mock_resolve.call_args[0][2] == self.task
        assert response.status_code == status.HTTP_200_OK


@pytest.mark.django_db
class TestProjectResolveStorageUri:
    @pytest.fixture
    def setup(self):
        # Create the necessary objects for testing without database
        self.factory = APIRequestFactory()
        self.user = MagicMock()
        self.project = MagicMock()  # Avoid using spec=Project - it triggers database access
        self.view = ProjectResolveStorageUri.as_view()

    @patch('io_storages.proxy_api.Project.objects.get')
    def test_get_with_missing_params(self, mock_project_get, setup):
        # Mock the database query
        mock_project_get.return_value = self.project

        # Test missing fileuri parameter
        request = self.factory.get('/project/1/resolve/')
        force_authenticate(request, user=self.user)
        response = self.view(request, project_id=1)
        assert response.status_code == status.HTTP_400_BAD_REQUEST

        # Test missing project_id parameter
        request = self.factory.get('/project/resolve/?fileuri=test')
        force_authenticate(request, user=self.user)
        response = self.view(request)
        assert response.status_code == status.HTTP_400_BAD_REQUEST

    @patch('io_storages.proxy_api.Project.objects.get')
    def test_get_project_not_found(self, mock_project_get, setup):
        # Mock the database query to raise DoesNotExist
        mock_project_get.side_effect = Project.DoesNotExist

        request = self.factory.get('/project/999/resolve/?fileuri=test')
        force_authenticate(request, user=self.user)
        response = self.view(request, project_id=999)

        assert response.status_code == status.HTTP_404_NOT_FOUND

    @patch('io_storages.proxy_api.Project.objects.get')
    @patch.object(ResolveStorageUriAPIMixin, 'resolve')
    def test_get_success(self, mock_resolve, mock_project_get, setup):
        # Mock the database query and resolve method
        mock_project_get.return_value = self.project
        mock_resolve.return_value = Response(status=status.HTTP_200_OK)

        request = self.factory.get('/project/1/resolve/?fileuri=test')
        force_authenticate(request, user=self.user)
        response = self.view(request, project_id=1)

        mock_project_get.assert_called_once_with(pk=1)
        # Use any_call instead of assert_called_once_with to handle DRF request vs WSGIRequest
        assert mock_resolve.call_args is not None
        assert mock_resolve.call_args[0][1] == 'test'
        assert mock_resolve.call_args[0][2] == self.project
        assert response.status_code == status.HTTP_200_OK