import base64
import io
import uuid
import zipfile
from unittest.mock import MagicMock, patch
from uuid import UUID

import django
from celery.result import AsyncResult
from django.contrib.auth import get_user_model
from django.core.cache import cache
from django.test import TestCase, override_settings
from django.test.client import Client
from graphene.test import Client as GrapheneClient
from graphql_relay import to_global_id

from config.graphql.schema import schema
from opencontractserver.constants.zip_import import (
    BULK_UPLOAD_OWNER_CACHE_PREFIX,
    BULK_UPLOAD_OWNER_CACHE_TTL_SECONDS,
)
from opencontractserver.corpuses.models import Corpus
from opencontractserver.documents.models import Document
from opencontractserver.tasks.import_tasks import process_documents_zip
from opencontractserver.types.enums import PermissionTypes
from opencontractserver.utils.permissioning import set_permissions_for_obj_to_user

User = get_user_model()


# Create a TestContext class similar to other test files
class TestContext:
    def __init__(self, user):
        self.user = user


class BulkDocumentUploadTests(TestCase):
    """Test the bulk document upload feature."""

    def setUp(self):
        self.user = User.objects.create_superuser(
            username="testuser",
            password="testpass",
            is_usage_capped=False,  # Ensure user can perform bulk uploads
        )
        self.client = Client()
        self.client.force_login(self.user)

        # Create a test corpus
        self.corpus = Corpus.objects.create(
            title="Test Corpus",
            description="Test Corpus for bulk upload",
            creator=self.user,
        )

        # Explicitly grant permissions for the test user on the test corpus
        # Even superusers might need explicit object permissions for guardian checks
        set_permissions_for_obj_to_user(self.user, self.corpus, [PermissionTypes.CRUD])

        # Create test file data
        self.pdf_content = b"%PDF-1.7\n1 0 obj\n<</Type/Catalog/Pages 2 0 R>>\nendobj\n2 0 obj\n<</Type/Pages/Kids[3 0 R]/Count 1>>\nendobj\n3 0 obj\n<</Type/Page/MediaBox[0 0 612 792]/Parent 2 0 R/Resources<<>>>>\nendobj\nxref\n0 4\n0000000000 65535 f\n0000000010 00000 n\n0000000053 00000 n\n0000000102 00000 n\ntrailer\n<</Size 4/Root 1 0 R>>\nstartxref\n178\n%%EOF"  # noqa: E501
        self.txt_content = b"This is a simple text file for testing."
        self.docx_content = b"PK\x03\x04\x14\x00\x06\x00\x08\x00\x00\x00!\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0c\x00\x00\x00word/document.xml"  # Mock DOCX content  # noqa: E501
        self.unsupported_content = b"This is an unsupported file type"

    def create_test_zip(self):
        """Create a zip file with test documents."""
        zip_buffer = io.BytesIO()

        with zipfile.ZipFile(zip_buffer, "w") as zip_file:
            # Add various file types
            zip_file.writestr("test_document.pdf", self.pdf_content)
            zip_file.writestr("test_document.txt", self.txt_content)
            zip_file.writestr("test_document.docx", self.docx_content)
            zip_file.writestr("unsupported.xyz", self.unsupported_content)
            # Add a directory (should be skipped)
            zip_file.writestr("test_dir/", b"")
            # Add a file in subdirectory
            zip_file.writestr("subdir/nested.pdf", self.pdf_content)

        zip_buffer.seek(0)
        return base64.b64encode(zip_buffer.read()).decode("utf-8")

    def execute_mutation(
        self, base64_file_string: str, add_to_corpus_id: str = None
    ) -> dict:
        """Execute the UploadDocumentsZip mutation and return the resulting dict."""
        # Create a client with correct context
        client = GrapheneClient(schema, context_value=TestContext(self.user))

        mutation = """
            mutation UploadDocumentsZip($base64FileString: String!, $addToCorpusId: ID) {
                uploadDocumentsZip(
                    base64FileString: $base64FileString,
                    makePublic: true,
                    addToCorpusId: $addToCorpusId
                ) {
                    ok
                    message
                    jobId
                }
            }
        """
        variables = {"base64FileString": base64_file_string, "makePublic": True}
        if add_to_corpus_id:
            variables["addToCorpusId"] = add_to_corpus_id

        # Execute the mutation
        response = client.execute(mutation, variable_values=variables)
        # Safely access nested dictionary
        mutation_result = response.get("data", {}).get("uploadDocumentsZip")
        if mutation_result:
            return response["data"]["uploadDocumentsZip"]
        return None

    def execute_status_query(self, job_id: str) -> dict:
        """Execute the bulkDocumentUploadStatus query with a given job_id and return the result dict."""
        # Create a client with correct context
        client = GrapheneClient(schema, context_value=TestContext(self.user))

        query = """
            query BulkDocumentUploadStatus($jobId: String!) {
                bulkDocumentUploadStatus(jobId: $jobId) {
                    completed
                    success
                    totalFiles
                    processedFiles
                    skippedFiles
                    errorFiles
                    documentIds
                    errors
                }
            }
        """
        variables = {"jobId": job_id}

        # Execute the query
        response = client.execute(query, variable_values=variables)
        if (
            response
            and "data" in response
            and "bulkDocumentUploadStatus" in response["data"]
        ):
            return response["data"]["bulkDocumentUploadStatus"]
        return None

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    @patch("opencontractserver.tasks.import_tasks.process_documents_zip")
    def test_upload_documents_zip_mutation(self, mock_process_task):
        """Test the UploadDocumentsZip mutation."""
        # Setup mock for the async task
        mock_process_task.apply_async.return_value = MagicMock()
        mock_process_task.s.return_value = mock_process_task

        # Mock task ID
        job_id = "12345678-1234-5678-1234-567812345678"
        mock_uuid = patch("uuid.uuid4", return_value=UUID(job_id))
        mock_uuid.start()

        # Create zip file
        base64_zip = self.create_test_zip()

        # Execute mutation without corpus
        response = self.execute_mutation(base64_zip)
        self.assertIsNotNone(
            response, "Response from execute_mutation should not be None"
        )
        self.assertIn("ok", response)
        self.assertTrue(response["ok"])
        self.assertEqual(response["jobId"], job_id)

        # Verify temporary file was created
        # temp_files = TemporaryFileHandle.objects.all()
        # self.assertEqual(len(temp_files), 1)
        # NOTE: This assertion is removed as it's fragile with eager task execution/mocking.
        # The handle is created, but the eager task (even if mocked) might interact
        # with transactions in a way that makes this check unreliable here.
        # The end-to-end test implicitly covers file handling.

        # Execute mutation with corpus
        corpus_id = to_global_id("CorpusType", self.corpus.id)
        response = self.execute_mutation(base64_zip, corpus_id)
        self.assertIsNotNone(
            response, "Response from execute_mutation with corpus_id should not be None"
        )
        self.assertIn("ok", response)
        self.assertTrue(response["ok"])

        # Clean up mock
        mock_uuid.stop()

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_bulk_document_upload_status_query(self):
        """Test bulk document upload status query after task completion by patching the task's run method."""

        # Define the dummy result that our task should return
        dummy_result = {
            "job_id": "dummy_job_123",
            "success": True,
            "completed": True,
            "total_files": 5,
            "processed_files": 3,
            "skipped_files": 2,
            "error_files": 0,
            "document_ids": ["1", "2", "3"],
            "errors": [],
        }

        # Generate a unique task ID for this test
        test_task_id = f"test-task-{uuid.uuid4()}"

        # Directly set the AsyncResult for this task ID to ensure it's available
        AsyncResult(test_task_id).backend.store_result(
            test_task_id, dummy_result, "SUCCESS"
        )

        # Bind ownership for the IDOR-protected status resolver.
        cache.set(
            f"{BULK_UPLOAD_OWNER_CACHE_PREFIX}{test_task_id}",
            self.user.id,
            BULK_UPLOAD_OWNER_CACHE_TTL_SECONDS,
        )

        # Use our known task ID to perform the status query
        response = self.execute_status_query(test_task_id)
        self.assertIsNotNone(
            response, "Response from execute_status_query should not be None"
        )

        # Check response structure
        self.assertIn(
            "completed", response, f"Response should have 'completed' key: {response}"
        )
        self.assertIn(
            "success", response, f"Response should have 'success' key: {response}"
        )
        self.assertIn(
            "totalFiles", response, f"Response should have 'totalFiles' key: {response}"
        )
        self.assertIn(
            "processedFiles",
            response,
            f"Response should have 'processedFiles' key: {response}",
        )
        self.assertIn(
            "documentIds",
            response,
            f"Response should have 'documentIds' key: {response}",
        )

        # Verify that the status query returns the expected values
        self.assertTrue(response["completed"])
        self.assertTrue(response["success"])
        self.assertEqual(response["totalFiles"], 5)
        self.assertEqual(response["processedFiles"], 3)
        self.assertEqual(response["documentIds"], ["1", "2", "3"])

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_end_to_end_document_upload(self):
        """Test the full document upload process end-to-end with real task execution."""
        # We'll use CELERY_TASK_ALWAYS_EAGER to run tasks synchronously

        initial_doc_count = Document.objects.count()

        # Create the zip file
        base64_zip = self.create_test_zip()

        # Execute the mutation with corpus ID
        corpus_id = to_global_id("CorpusType", self.corpus.id)
        response = self.execute_mutation(base64_zip, corpus_id)
        self.assertIsNotNone(
            response, "Response from execute_mutation should not be None"
        )

        # Check response structure
        self.assertIn("ok", response, f"Response should have 'ok' key: {response}")
        self.assertTrue(response["ok"])

        # NOTE: In CELERY_TASK_ALWAYS_EAGER mode, the task runs synchronously
        # before control returns here. Querying status immediately can be unreliable.
        # We will rely on checking the final DB state instead.

        # Check if new documents were created
        new_doc_count = Document.objects.count()
        self.assertGreater(new_doc_count, initial_doc_count)

        # Verify documents are associated with the corpus
        corpus_docs = self.corpus._get_active_documents().count()
        self.assertGreater(corpus_docs, 0)

        # Verify document titles and content
        # With corpus isolation, we need to check corpus copies (have DocumentPath), not originals
        from opencontractserver.documents.models import DocumentPath

        corpus_doc_ids = DocumentPath.objects.filter(
            corpus=self.corpus, is_current=True, is_deleted=False
        ).values_list("document_id", flat=True)
        documents = Document.objects.filter(
            id__in=corpus_doc_ids, creator=self.user
        ).order_by("-created")

        # Filter to only documents that should match the pattern and check the first 3 of those
        test_docs = [doc for doc in documents if doc.title.startswith("test_document")][
            :3
        ]
        self.assertEqual(
            len(test_docs),
            3,
            "Should have found 3 documents starting with 'test_document'",
        )
        for doc in test_docs:
            self.assertTrue(doc.title.startswith("test_document"))
            self.assertIn(
                doc.file_type,
                [
                    "application/pdf",
                    "text/plain",
                    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                ],
            )

        # Verify the documents are linked to the corpus via DocumentPath
        for doc in test_docs:
            # Check that a DocumentPath exists linking the document to the corpus
            doc_path_exists = DocumentPath.objects.filter(
                document=doc,
                corpus=self.corpus,
                is_current=True,
                is_deleted=False,
            ).exists()
            self.assertTrue(
                doc_path_exists,
                f"Document {doc.id} should have a DocumentPath to corpus {self.corpus.id}",
            )

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_user_capped_limit_reached(self):
        """Test the case where a usage-capped user has already reached their document limit."""
        # Create a usage-capped user who has reached their limit
        capped_user = User.objects.create_user(
            username="capped_user",
            password="testpass",
            is_usage_capped=True,
        )

        # Set a low document cap for testing
        with override_settings(USAGE_CAPPED_USER_DOC_CAP_COUNT=1):
            # Create a document to reach the limit
            Document.objects.create(
                creator=capped_user,
                title="Existing document",
                description="Document that fills the user's quota",
            )

            # Create a test corpus for this user
            corpus = Corpus.objects.create(
                title="Capped User Corpus",
                description="Test Corpus for capped user",
                creator=capped_user,
            )
            set_permissions_for_obj_to_user(capped_user, corpus, [PermissionTypes.CRUD])

            # Create zip file
            base64_zip = self.create_test_zip()

            # Execute the mutation directly with process_documents_zip to check results
            from opencontractserver.corpuses.models import TemporaryFileHandle

            temp_file = TemporaryFileHandle.objects.create()
            temp_file.file.save("test.zip", io.BytesIO(base64.b64decode(base64_zip)))

            job_id = str(uuid.uuid4())
            results = process_documents_zip(
                temporary_file_handle_id=temp_file.id,
                user_id=capped_user.id,
                job_id=job_id,
                corpus_id=corpus.id,
            )

            # Verify the task completed but failed due to user cap
            self.assertTrue(results["completed"])
            self.assertFalse(results["success"])
            self.assertEqual(results["processed_files"], 0)
            self.assertTrue(
                any("maximum document limit" in error for error in results["errors"])
            )

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_user_capped_limit_reached_during_upload(self):
        """Test the case where a usage-capped user reaches their limit during upload."""
        # Create a usage-capped user who is just under their limit
        capped_user = User.objects.create_user(
            username="capped_user_mid",
            password="testpass",
            is_usage_capped=True,
        )

        # Set document cap to 1 for this test
        with override_settings(USAGE_CAPPED_USER_DOC_CAP_COUNT=1):
            # Create a test corpus for this user
            corpus = Corpus.objects.create(
                title="Capped User Corpus",
                description="Test Corpus for capped user",
                creator=capped_user,
            )
            set_permissions_for_obj_to_user(capped_user, corpus, [PermissionTypes.CRUD])

            # Create a zip with multiple files to trigger mid-upload cap
            zip_buffer = io.BytesIO()
            with zipfile.ZipFile(zip_buffer, "w") as zip_file:
                zip_file.writestr("doc1.pdf", self.pdf_content)
                zip_file.writestr(
                    "doc2.pdf", self.pdf_content
                )  # This should be skipped

            zip_buffer.seek(0)
            base64_zip = base64.b64encode(zip_buffer.read()).decode("utf-8")

            # Execute the task directly
            from opencontractserver.corpuses.models import TemporaryFileHandle

            temp_file = TemporaryFileHandle.objects.create()
            temp_file.file.save(
                "test_mid.zip", io.BytesIO(base64.b64decode(base64_zip))
            )

            job_id = str(uuid.uuid4())
            results = process_documents_zip(
                temporary_file_handle_id=temp_file.id,
                user_id=capped_user.id,
                job_id=job_id,
                corpus_id=corpus.id,
            )

            # Verify only one document was processed before hitting the limit
            self.assertTrue(results["completed"])
            self.assertEqual(results["processed_files"], 1)
            self.assertEqual(len(results["document_ids"]), 1)
            self.assertTrue(
                any(
                    "User document limit reached during processing" in error
                    for error in results["errors"]
                )
            )

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_unknown_binary_file_skipped(self):
        """Test that unknown binary files are skipped during processing."""
        # Create a zip with an unknown binary file
        unknown_binary = b"\x00\x01\x02\x03\xde\xad\xbe\xef"  # Non-text binary data

        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w") as zip_file:
            zip_file.writestr("unknown.bin", unknown_binary)

        zip_buffer.seek(0)
        base64_zip = base64.b64encode(zip_buffer.read()).decode("utf-8")

        # Execute the task directly
        from opencontractserver.corpuses.models import TemporaryFileHandle

        temp_file = TemporaryFileHandle.objects.create()
        temp_file.file.save(
            "test_unknown.zip", io.BytesIO(base64.b64decode(base64_zip))
        )

        job_id = str(uuid.uuid4())
        results = process_documents_zip(
            temporary_file_handle_id=temp_file.id,
            user_id=self.user.id,
            job_id=job_id,
        )

        # Verify the unknown file was skipped
        self.assertTrue(results["completed"])
        self.assertTrue(results["success"])
        self.assertEqual(results["total_files"], 1)
        self.assertEqual(results["skipped_files"], 1)
        self.assertEqual(results["processed_files"], 0)

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_file_processing_error(self):
        """Test handling of errors during file processing."""
        # Create a normal zip file
        base64_zip = self.create_test_zip()

        from opencontractserver.corpuses.models import TemporaryFileHandle

        temp_file = TemporaryFileHandle.objects.create()
        temp_file.file.save("test_error.zip", io.BytesIO(base64.b64decode(base64_zip)))

        # Mock the ContentFile constructor to raise an exception for one specific file
        # ContentFile is now used in versioning.py (via import_content -> import_document)
        original_content_file = django.core.files.base.ContentFile

        def mock_content_file(content, name=None):
            if name and name.endswith(".pdf"):
                raise OSError("Simulated error processing PDF file")
            return original_content_file(content, name)

        with patch(
            "opencontractserver.documents.versioning.ContentFile",
            side_effect=mock_content_file,
        ):
            job_id = str(uuid.uuid4())
            results = process_documents_zip(
                temporary_file_handle_id=temp_file.id,
                user_id=self.user.id,
                job_id=job_id,
            )

        # Verify error handling
        self.assertTrue(results["completed"])
        self.assertGreater(results["error_files"], 0)
        self.assertTrue(any("Error processing" in error for error in results["errors"]))

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_job_level_exception(self):
        """Test handling of a job-level exception."""
        # Create a zip file
        base64_zip = self.create_test_zip()

        from opencontractserver.corpuses.models import TemporaryFileHandle

        temp_file = TemporaryFileHandle.objects.create()
        temp_file.file.save(
            "test_job_error.zip", io.BytesIO(base64.b64decode(base64_zip))
        )

        # Force a job-level exception by patching the TemporaryFileHandle.objects.get method
        with patch(
            "opencontractserver.corpuses.models.TemporaryFileHandle.objects.get",
            side_effect=Exception("Simulated job-level error"),
        ):
            job_id = str(uuid.uuid4())
            results = process_documents_zip(
                temporary_file_handle_id=999999,  # This ID doesn't matter due to our mock
                user_id=self.user.id,
                job_id=job_id,
            )

        # Verify job-level error handling
        self.assertTrue(results["completed"])
        self.assertFalse(results["success"])
        self.assertEqual(results["processed_files"], 0)
        self.assertTrue(any("Job failed" in error for error in results["errors"]))


class BulkDocumentUploadStatusIDORTests(TestCase):
    """
    Regression tests for the bulk_document_upload_status IDOR.

    Before the fix, the resolver called celery_app.AsyncResult(job_id) with
    no ownership check — any authenticated user with a leaked job_id could
    read the job's results (including document_ids of the import).
    """

    def setUp(self):
        self.owner = User.objects.create_user(username="owner", password="x")
        self.attacker = User.objects.create_user(username="attacker", password="x")

    def _execute_status_query(self, user, job_id: str) -> dict:
        client = GrapheneClient(schema, context_value=TestContext(user))
        query = """
            query BulkDocumentUploadStatus($jobId: String!) {
                bulkDocumentUploadStatus(jobId: $jobId) {
                    completed
                    success
                    totalFiles
                    documentIds
                    errors
                }
            }
        """
        return client.execute(query, variable_values={"jobId": job_id})

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_attacker_cannot_read_other_users_bulk_upload_status(self):
        """A user without the ownership cache entry must get the uniform not-found response."""
        # Owner enqueues a job (we just simulate the cache binding + a stored result).
        job_id = f"test-task-{uuid.uuid4()}"
        cache.set(
            f"{BULK_UPLOAD_OWNER_CACHE_PREFIX}{job_id}",
            self.owner.id,
            BULK_UPLOAD_OWNER_CACHE_TTL_SECONDS,
        )
        AsyncResult(job_id).backend.store_result(
            job_id,
            {
                "success": True,
                "completed": True,
                "total_files": 7,
                "document_ids": ["secret-doc-1", "secret-doc-2"],
                "errors": [],
            },
            "SUCCESS",
        )

        # Owner sees the real result.
        owner_response = self._execute_status_query(self.owner, job_id)
        owner_payload = owner_response["data"]["bulkDocumentUploadStatus"]
        self.assertEqual(owner_payload["totalFiles"], 7)
        self.assertEqual(owner_payload["documentIds"], ["secret-doc-1", "secret-doc-2"])

        # Attacker gets the opaque not-found response — no document IDs leak.
        attacker_response = self._execute_status_query(self.attacker, job_id)
        attacker_payload = attacker_response["data"]["bulkDocumentUploadStatus"]
        self.assertFalse(attacker_payload["completed"])
        self.assertFalse(attacker_payload["success"])
        # documentIds is unset on the not-found path → serializes to None.
        self.assertIsNone(attacker_payload["documentIds"])
        self.assertIn("not found", " ".join(attacker_payload["errors"]).lower())

    @override_settings(
        CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_STORE_EAGER_RESULT=True
    )
    def test_unknown_job_id_returns_uniform_not_found(self):
        """Cache miss must fail closed — no celery state should be revealed."""
        job_id = f"unknown-{uuid.uuid4()}"
        response = self._execute_status_query(self.attacker, job_id)
        payload = response["data"]["bulkDocumentUploadStatus"]
        self.assertFalse(payload["completed"])
        self.assertFalse(payload["success"])
        self.assertIn("not found", " ".join(payload["errors"]).lower())
