# Copyright 2025-present the zvec project
#
# 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.

import logging
import pytest
import threading
import numpy as np
import zvec

from zvec import (
    CollectionOption,
    InvertIndexParam,
    HnswIndexParam,
    Collection,
    Doc,
    DataType,
    FieldSchema,
    VectorSchema,
)


class TestCollectionConcurrency:
    @pytest.fixture(scope="function")
    def test_collection(self, tmp_path_factory):
        """Fixture to create a test collection"""
        collection_schema = zvec.CollectionSchema(
            name="test_collection",
            fields=[
                FieldSchema(
                    "id",
                    DataType.INT64,
                    nullable=False,
                    index_param=InvertIndexParam(enable_range_optimization=True),
                ),
                FieldSchema(
                    "name",
                    DataType.STRING,
                    nullable=False,
                    index_param=InvertIndexParam(),
                ),
                FieldSchema("weight", DataType.FLOAT, nullable=True),
            ],
            vectors=[
                VectorSchema(
                    "dense",
                    DataType.VECTOR_FP32,
                    dimension=128,
                    index_param=HnswIndexParam(),
                ),
                VectorSchema(
                    "sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
                ),
            ],
        )

        collection_option = CollectionOption(read_only=False, enable_mmap=True)

        temp_dir = tmp_path_factory.mktemp("zvec")
        collection_path = temp_dir / "test_collection"

        coll = zvec.create_and_open(
            path=str(collection_path),
            schema=collection_schema,
            option=collection_option,
        )

        assert coll is not None, "Failed to create and open collection"

        yield coll

        # Clean up
        if hasattr(coll, "destroy") and coll is not None:
            try:
                coll.destroy()
            except Exception as e:
                print(f"Warning: failed to destroy collection: {e}")

    def test_concurrent_read_write(self, test_collection: Collection):
        results = []

        def insert_docs(thread_id):
            try:
                docs = [
                    Doc(
                        id=f"{thread_id}_{i}",
                        fields={
                            "id": int(f"{thread_id}{i}"),
                            "name": f"thread_{thread_id}_doc_{i}",
                            "weight": float(i),
                        },
                        vectors={
                            "dense": np.random.random(128).tolist(),
                            "sparse": {1: float(i), 2: float(i * 2)},
                        },
                    )
                    for i in range(5)
                ]

                result = test_collection.insert(docs)
                results.append((thread_id, "insert", len(result)))
            except Exception as e:
                results.append((thread_id, "insert_exception", str(e)))

        def query_docs(thread_id):
            try:
                result = test_collection.query(filter="id > 0", topk=10)
                results.append((thread_id, "query", len(result)))
            except Exception as e:
                results.append((thread_id, "query_exception", str(e)))

        # Create threads for concurrent operations
        threads = []

        # Start insert threads
        for i in range(3):
            thread = threading.Thread(target=insert_docs, args=(i,))
            threads.append(thread)
            thread.start()

        # Start query threads
        for i in range(3):
            thread = threading.Thread(target=query_docs, args=(i,))
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Analyze results
        insert_results = [r for r in results if r[1] == "insert"]
        query_results = [r for r in results if r[1] == "query"]

        logging.info(
            f"Concurrent read/write results - Inserts: {len(insert_results)}, Queries: {len(query_results)}"
        )

        # At least some operations should succeed
        assert len(insert_results) + len(query_results) > 0

    def test_concurrent_query(self, test_collection: Collection):
        # First insert some data
        docs = [
            Doc(
                id=f"{i}",
                fields={"id": i, "name": f"test_{i}", "weight": float(i)},
                vectors={
                    "dense": np.random.random(128).tolist(),
                    "sparse": {1: float(i), 2: float(i * 2)},
                },
            )
            for i in range(20)
        ]

        insert_result = test_collection.insert(docs)
        assert len(insert_result) == 20

        results = []

        def query_operation(thread_id):
            """Perform query operation from a thread"""
            try:
                result = test_collection.query(filter=f"id > {thread_id}", topk=5)
                results.append((thread_id, "query", len(result)))
            except Exception as e:
                results.append((thread_id, "query_exception", str(e)))

        # Create multiple threads for concurrent queries
        threads = []
        for i in range(5):
            thread = threading.Thread(target=query_operation, args=(i,))
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Analyze results
        query_results = [r for r in results if r[1] == "query"]
        logging.info(f"Concurrent query results - Queries: {len(query_results)}")

        # All query operations should succeed
        assert len(query_results) == 5

    def test_concurrent_modifications(self, test_collection: Collection):
        # First insert some data
        docs = [
            Doc(
                id=f"{i}",
                fields={"id": i, "name": f"test_{i}", "weight": float(i)},
                vectors={
                    "dense": np.random.random(128).tolist(),
                    "sparse": {1: float(i), 2: float(i * 2)},
                },
            )
            for i in range(10)
        ]

        insert_result = test_collection.insert(docs)
        assert len(insert_result) == 10

        results = []

        def update_operation(thread_id):
            """Perform update operation from a thread"""
            try:
                # Each thread updates different documents
                update_docs = [
                    Doc(
                        id=f"{i}",
                        fields={
                            "id": i,
                            "name": f"updated_by_thread_{thread_id}",
                            "weight": float(i + thread_id),
                        },
                        vectors={
                            "dense": np.random.random(128).tolist(),
                            "sparse": {1: float(i) + 0.5, 2: float(i * 2) + 0.5},
                        },
                    )
                    for i in range(thread_id * 2, thread_id * 2 + 2)
                ]

                result = test_collection.update(update_docs)
                results.append((thread_id, "update", len(result)))
            except Exception as e:
                results.append((thread_id, "update_exception", str(e)))

        def delete_operation(thread_id):
            """Perform delete operation from a thread"""
            try:
                # Each thread deletes different documents
                delete_ids = [f"{thread_id * 2 + 2}", f"{thread_id * 2 + 3}"]
                result = test_collection.delete(delete_ids)
                results.append((thread_id, "delete", len(result)))
            except Exception as e:
                results.append((thread_id, "delete_exception", str(e)))

        # Create threads for concurrent operations
        threads = []

        # Start update threads
        for i in range(3):
            thread = threading.Thread(target=update_operation, args=(i,))
            threads.append(thread)
            thread.start()

        # Start delete threads
        for i in range(2):
            thread = threading.Thread(target=delete_operation, args=(i,))
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Analyze results
        update_results = [r for r in results if r[1] == "update"]
        delete_results = [r for r in results if r[1] == "delete"]

        logging.info(
            f"Concurrent modification results - Updates: {len(update_results)}, Deletes: {len(delete_results)}"
        )

        # At least some operations should succeed
        assert len(update_results) + len(delete_results) > 0

    def test_read_write_locking(self, test_collection: Collection):
        # Perform operations that should be thread-safe
        docs = [
            Doc(
                id=f"{i}",
                fields={"id": i, "name": f"test_{i}", "weight": float(i)},
                vectors={
                    "dense": np.random.random(128).tolist(),
                    "sparse": {1: float(i), 2: float(i * 2)},
                },
            )
            for i in range(5)
        ]

        # Insert data
        insert_result = test_collection.insert(docs)
        assert len(insert_result) == 5

        # Concurrent operations should not cause data corruption
        results = []

        def mixed_operation(thread_id):
            """Perform mixed operations from a thread"""
            try:
                # Mix of read and write operations
                if thread_id % 2 == 0:
                    # Read operation
                    result = test_collection.fetch([f"{thread_id % 5}"])
                    results.append((thread_id, "read", len(result)))
                else:
                    # Write operation
                    doc = Doc(
                        id=f"{thread_id % 5}",
                        fields={
                            "id": thread_id % 5,
                            "name": f"mixed_op_{thread_id}",
                            "weight": float(thread_id),
                        },
                        vectors={
                            "dense": np.random.random(128).tolist(),
                            "sparse": {1: float(thread_id), 2: float(thread_id * 2)},
                        },
                    )
                    result = test_collection.upsert(doc)
                    results.append((thread_id, "write", len(result)))
            except Exception as e:
                results.append((thread_id, "exception", str(e)))

        # Create multiple threads
        threads = []
        for i in range(10):
            thread = threading.Thread(target=mixed_operation, args=(i,))
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Verify that the collection is still in a consistent state
        final_result = test_collection.query()
        assert len(final_result) >= 0  # Should not crash or return corrupted data

    def test_race_condition_detection(self, test_collection: Collection):
        # Insert initial data
        docs = [
            Doc(
                id=f"{i}",
                fields={"id": i, "name": f"initial_{i}", "weight": float(i)},
                vectors={
                    "dense": np.random.random(128).tolist(),
                    "sparse": {1: float(i), 2: float(i * 2)},
                },
            )
            for i in range(10)
        ]

        insert_result = test_collection.insert(docs)
        assert len(insert_result) == 10

        # Perform many rapid concurrent operations
        operation_count = 100
        results = []

        def rapid_operation(op_id):
            """Perform rapid operations"""
            try:
                # Alternate between different types of operations
                if op_id % 4 == 0:
                    # Insert
                    doc = Doc(
                        id=f"rapid_{op_id}",
                        fields={
                            "id": op_id,
                            "name": f"rapid_{op_id}",
                            "weight": float(op_id),
                        },
                        vectors={
                            "dense": np.random.random(128).tolist(),
                            "sparse": {1: float(op_id), 2: float(op_id * 2)},
                        },
                    )
                    result = test_collection.insert(doc)
                    results.append(("insert", len(result)))
                elif op_id % 4 == 1:
                    # Update
                    doc = Doc(
                        id=f"{op_id % 10}",
                        fields={
                            "id": op_id % 10,
                            "name": f"rapid_update_{op_id}",
                            "weight": float(op_id),
                        },
                        vectors={
                            "dense": np.random.random(128).tolist(),
                            "sparse": {1: float(op_id), 2: float(op_id * 2)},
                        },
                    )
                    result = test_collection.update(doc)
                    results.append(("update", len(result)))
                elif op_id % 4 == 2:
                    # Query
                    result = test_collection.query(filter=f"id > {op_id % 5}", topk=3)
                    results.append(("query", len(result)))
                else:
                    # Fetch
                    result = test_collection.fetch([f"{op_id % 10}"])
                    results.append(("fetch", len(result)))
            except Exception as e:
                results.append(("exception", str(e)))

        # Create many threads for rapid concurrent operations
        threads = []
        for i in range(operation_count):
            thread = threading.Thread(target=rapid_operation, args=(i,))
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Verify collection is still functional
        final_query = test_collection.query()
        assert len(final_query) >= 0  # Should not be corrupted

        logging.info(
            f"Rapid concurrent operations completed - Total operations: {len(results)}"
        )