from unittest.mock import patch

import pytest
from data_import.functions import _async_import_background_streaming
from data_import.models import FileUpload
from data_import.uploader import load_tasks_for_async_import_streaming
from django.core.files.base import ContentFile
from organizations.tests.factories import OrganizationFactory
from projects.models import ProjectImport
from projects.tests.factories import ProjectFactory
from rest_framework.exceptions import ValidationError
from users.tests.factories import UserFactory

pytestmark = pytest.mark.django_db


@pytest.fixture
def user():
    return UserFactory()


@pytest.fixture
def organization():
    return OrganizationFactory()


@pytest.fixture
def project(user, organization):
    # Simple config with a single text field
    return ProjectFactory(
        created_by=user,
        organization=organization,
        label_config='<View><Text name="text" value="$text"/></View>',
    )


def create_file_upload(user, project, body: bytes, name: str):
    return FileUpload.objects.create(user=user, project=project, file=ContentFile(body, name=name))


class TestJSONStreamingReader:
    def test_array_of_objects_wraps_data(self, user, project):
        content = b'[{"text":"A"},{"text":"B"}]'
        fu = FileUpload.objects.create(user=user, project=project, file=ContentFile(content, name='tasks.json'))

        batches = list(fu.read_tasks_list_from_json_streaming(batch_size=1))
        # Flatten
        tasks = [t for batch in batches for t in batch]

        assert len(tasks) == 2
        assert tasks[0]['data'] == {'text': 'A'}
        assert tasks[1]['data'] == {'text': 'B'}

    def test_array_of_objects_with_data_preserved(self, user, project):
        content = b'[{"data":{"text":"A"}},{"data":{"text":"B"}}]'
        fu = FileUpload.objects.create(user=user, project=project, file=ContentFile(content, name='tasks.json'))

        batches = list(fu.read_tasks_list_from_json_streaming(batch_size=2))
        tasks = [t for batch in batches for t in batch]

        assert len(tasks) == 2
        assert tasks[0]['data'] == {'text': 'A'}
        assert tasks[1]['data'] == {'text': 'B'}

    def test_single_object(self, user, project):
        content = b'{"text":"A"}'
        fu = FileUpload.objects.create(user=user, project=project, file=ContentFile(content, name='task.json'))

        batches = list(fu.read_tasks_list_from_json_streaming(batch_size=10))
        assert len(batches) == 1
        assert len(batches[0]) == 1
        assert batches[0][0]['data'] == {'text': 'A'}

    def test_invalid_array_of_strings_raises(self, user, project):
        content = b'["A","B"]'
        fu = FileUpload.objects.create(user=user, project=project, file=ContentFile(content, name='tasks.json'))

        with pytest.raises(ValidationError):
            list(fu.read_tasks_list_from_json_streaming(batch_size=2))

    def test_invalid_top_level_type_raises(self, user, project):
        content = b'"A"'
        fu = FileUpload.objects.create(user=user, project=project, file=ContentFile(content, name='tasks.json'))

        with pytest.raises(ValidationError):
            list(fu.read_tasks_list_from_json_streaming(batch_size=2))

    def test_read_tasks_streaming_batches(self, user, project):
        items = ','.join([f'{{"text":"T{i}"}}' for i in range(7)]).encode('utf-8')
        content = b'[' + items + b']'
        fu = FileUpload.objects.create(user=user, project=project, file=ContentFile(content, name='tasks.json'))

        batches = list(fu.read_tasks_streaming(batch_size=3))
        # Expect 3,3,1
        sizes = [len(b) for b in batches]
        assert sizes == [3, 3, 1]
        # Spot check the first element
        assert batches[0][0]['data'] == {'text': 'T0'}


class TestEndToEndStreamingFromUploads:
    def test_load_tasks_from_uploaded_files_streaming_real_files(self, user, project):
        content1 = b'[{"text":"A1"},{"text":"A2"},{"text":"A3"},{"text":"A4"}]'
        content2 = b'[{"text":"B1"},{"text":"B2"},{"text":"B3"},{"text":"B4"}]'
        FileUpload.objects.create(user=user, project=project, file=ContentFile(content1, name='a.json'))
        FileUpload.objects.create(user=user, project=project, file=ContentFile(content2, name='b.json'))

        gen = FileUpload.load_tasks_from_uploaded_files_streaming(project, batch_size=3)
        batches = list(gen)

        # Expect 3,3,2 tasks across two files aggregated by batch size
        sizes = [len(b[0]) for b in batches]
        assert sizes == [3, 3, 2]

        # All tasks must have file_upload_id set
        for batch_tasks, _, _ in batches:
            for t in batch_tasks:
                assert 'file_upload_id' in t
                assert t['data'] and isinstance(t['data'], dict)


class TestLoadTasksForAsyncImportStreaming:
    def test_from_file_upload_ids_batches_and_metadata(self, user, project, settings):
        settings.IMPORT_BATCH_SIZE = 3
        fu1 = create_file_upload(user, project, b'[{"text":"A1"},{"text":"A2"},{"text":"A3"}]', 'a.json')
        fu2 = create_file_upload(user, project, b'[{"text":"B1"},{"text":"B2"},{"text":"B3"},{"text":"B4"}]', 'b.json')

        pimport = ProjectImport.objects.create(project=project, file_upload_ids=[fu1.id, fu2.id])

        gen = load_tasks_for_async_import_streaming(pimport, user, batch_size=3)
        batches = list(gen)

        # Expect 3,3,1 aggregated across files with batch size 3
        sizes = [len(b[0]) for b in batches]
        assert sizes == [3, 3, 1]

        # Each batch returns file_upload_ids and metadata
        for batch_tasks, file_upload_ids, found_formats, data_columns in batches:
            assert file_upload_ids == [fu1.id, fu2.id]
            assert isinstance(found_formats, dict)
            assert isinstance(data_columns, list)

    def test_from_url_with_json_string(self, user, project, settings):
        settings.IMPORT_BATCH_SIZE = 10
        # Use JSON-in-URL mode (no network). Function will create an inplace file upload and parse it.
        url_json = '[{"text":"U1"},{"text":"U2"}]'
        pimport = ProjectImport.objects.create(project=project, url=url_json)

        gen = load_tasks_for_async_import_streaming(pimport, user, batch_size=2)
        batches = list(gen)

        assert len(batches) == 1
        tasks, file_upload_ids, found_formats, data_columns = batches[0]
        assert [t['data'] for t in tasks] == [{'text': 'U1'}, {'text': 'U2'}]
        assert isinstance(file_upload_ids, list)
        assert isinstance(found_formats, dict)
        assert isinstance(data_columns, list)

    def test_from_tasks_inline(self, user, project, settings):
        settings.IMPORT_BATCH_SIZE = 2
        tasks = [{'data': {'text': 'T1'}}, {'data': {'text': 'T2'}}, {'data': {'text': 'T3'}}]
        pimport = ProjectImport.objects.create(project=project, tasks=tasks)

        gen = load_tasks_for_async_import_streaming(pimport, user, batch_size=2)
        batches = list(gen)

        sizes = [len(b[0]) for b in batches]
        assert sizes == [2, 1]


class TestAsyncImportBackgroundStreaming:
    @patch('data_import.functions.flag_set', return_value=False)
    def test_counts_and_status_without_commit(self, mock_flag, user, project, settings):
        settings.IMPORT_BATCH_SIZE = 4
        fu1 = create_file_upload(user, project, b'[{"text":"A1"},{"text":"A2"}]', 'a.json')
        fu2 = create_file_upload(user, project, b'[{"text":"B1"},{"text":"B2"},{"text":"B3"}]', 'b.json')

        pimport = ProjectImport.objects.create(
            project=project,
            file_upload_ids=[fu1.id, fu2.id],
            commit_to_project=False,
            return_task_ids=False,
        )

        _async_import_background_streaming(pimport, user)

        pimport.refresh_from_db()
        assert pimport.status == ProjectImport.Status.COMPLETED
        assert pimport.task_count == 5
        # found_formats and data_columns are populated
        assert isinstance(pimport.found_formats, dict)
        assert pimport.found_formats.get('.json')
        assert isinstance(pimport.data_columns, (list, set))