import asyncio
import os
import shutil
import tempfile
from uuid import UUID
import pytest
from itertools import count
from typing import (
    Generator,
    List,
    Callable,
    Optional,
    Dict,
    Union,
    Iterator,
    Sequence,
    Tuple,
)
from chromadb.errors import BatchSizeExceededError
from chromadb.ingest import Producer, Consumer
from chromadb.db.impl.sqlite import SqliteDB
from chromadb.test.conftest import ProducerFn
from chromadb.types import (
    OperationRecord,
    Operation,
    LogRecord,
    ScalarEncoding,
)
from chromadb.config import System, Settings
from pytest import FixtureRequest, approx
from asyncio import Event, wait_for, TimeoutError
import numpy as np


def sqlite() -> Generator[Tuple[Producer, Consumer], None, None]:
    """Fixture generator for sqlite Producer + Consumer"""
    system = System(Settings(allow_reset=True))
    db = system.require(SqliteDB)
    system.start()
    yield db, db
    system.stop()


def sqlite_persistent() -> Generator[Tuple[Producer, Consumer], None, None]:
    """Fixture generator for sqlite_persistent Producer + Consumer"""
    save_path = tempfile.mkdtemp()
    system = System(
        Settings(allow_reset=True, is_persistent=True, persist_directory=save_path)
    )
    db = system.require(SqliteDB)
    system.start()
    yield db, db
    system.stop()
    if os.path.exists(save_path):
        shutil.rmtree(save_path)


def fixtures() -> List[Callable[[], Generator[Tuple[Producer, Consumer], None, None]]]:
    fixtures = [sqlite, sqlite_persistent]
    if "CHROMA_CLUSTER_TEST_ONLY" in os.environ:
        # TODO: We should add the new log service here
        fixtures = []

    return fixtures


@pytest.fixture(scope="module", params=fixtures())
def producer_consumer(
    request: FixtureRequest,
) -> Generator[Tuple[Producer, Consumer], None, None]:
    yield next(request.param())


@pytest.fixture(scope="module")
def sample_embeddings() -> Iterator[OperationRecord]:
    def create_record(i: int) -> OperationRecord:
        vector = np.array([i + i * 0.1, i + 1 + i * 0.1])
        metadata: Optional[Dict[str, Union[str, int, float]]]
        if i % 2 == 0:
            metadata = None
        else:
            metadata = {"str_key": f"value_{i}", "int_key": i, "float_key": i + i * 0.1}

        record = OperationRecord(
            id=f"embedding_{i}",
            embedding=vector,
            encoding=ScalarEncoding.FLOAT32,
            metadata=metadata,
            operation=Operation.ADD,
        )
        return record

    return (create_record(i) for i in count())


class CapturingConsumeFn:
    embeddings: List[LogRecord]
    waiters: List[Tuple[int, Event]]

    def __init__(self) -> None:
        """A function that captures embeddings and allows you to wait for a certain
        number of embeddings to be available. It must be constructed in the thread with
        the main event loop
        """
        self.embeddings = []
        self.waiters = []
        self._loop = asyncio.get_event_loop()

    def __call__(self, embeddings: Sequence[LogRecord]) -> None:
        self.embeddings.extend(embeddings)
        for n, event in self.waiters:
            if len(self.embeddings) >= n:
                # event.set() is not thread safe, so we need to call it in the main event loop
                self._loop.call_soon_threadsafe(event.set)

    async def get(self, n: int, timeout_secs: int = 10) -> Sequence[LogRecord]:
        "Wait until at least N embeddings are available, then return all embeddings"
        if len(self.embeddings) >= n:
            return self.embeddings[:n]
        else:
            event = Event()
            self.waiters.append((n, event))
            # timeout so we don't hang forever on failure
            await wait_for(event.wait(), timeout_secs)
            return self.embeddings[:n]


def assert_approx_equal(a: Sequence[float], b: Sequence[float]) -> None:
    for i, j in zip(a, b):
        assert approx(i) == approx(j)


def assert_records_match(
    inserted_records: Sequence[OperationRecord],
    consumed_records: Sequence[LogRecord],
) -> None:
    """Given a list of inserted and consumed records, make sure they match"""
    assert len(consumed_records) == len(inserted_records)
    for inserted, consumed in zip(inserted_records, consumed_records):
        assert inserted["id"] == consumed["record"]["id"]
        assert inserted["operation"] == consumed["record"]["operation"]
        assert inserted["encoding"] == consumed["record"]["encoding"]
        assert inserted["metadata"] == consumed["record"]["metadata"]

        if inserted["embedding"] is not None:
            assert consumed["record"]["embedding"] is not None
            assert_approx_equal(inserted["embedding"], consumed["record"]["embedding"])


@pytest.mark.asyncio
async def test_backfill(
    producer_consumer: Tuple[Producer, Consumer],
    sample_embeddings: Iterator[OperationRecord],
    produce_fns: ProducerFn,
) -> None:
    producer, consumer = producer_consumer
    producer.reset_state()
    consumer.reset_state()
    collection_id = UUID("00000000-0000-0000-0000-000000000000")
    embeddings = produce_fns(producer, collection_id, sample_embeddings, 3)[0]

    consume_fn = CapturingConsumeFn()
    consumer.subscribe(collection_id, consume_fn, start=consumer.min_seqid())

    recieved = await consume_fn.get(3)
    assert_records_match(embeddings, recieved)


@pytest.mark.asyncio
async def test_notifications(
    producer_consumer: Tuple[Producer, Consumer],
    sample_embeddings: Iterator[OperationRecord],
) -> None:
    producer, consumer = producer_consumer
    producer.reset_state()
    consumer.reset_state()
    collection_id = UUID("00000000-0000-0000-0000-000000000000")

    embeddings: List[OperationRecord] = []

    consume_fn = CapturingConsumeFn()

    consumer.subscribe(collection_id, consume_fn, start=consumer.min_seqid())

    for i in range(10):
        e = next(sample_embeddings)
        embeddings.append(e)
        producer.submit_embedding(collection_id, e)
        received = await consume_fn.get(i + 1)
        assert_records_match(embeddings, received)


@pytest.mark.asyncio
async def test_multiple_collections(
    producer_consumer: Tuple[Producer, Consumer],
    sample_embeddings: Iterator[OperationRecord],
) -> None:
    producer, consumer = producer_consumer
    producer.reset_state()
    consumer.reset_state()
    collection_1 = UUID("00000000-0000-0000-0000-000000000001")
    collection_2 = UUID("00000000-0000-0000-0000-000000000002")

    embeddings_1: List[OperationRecord] = []
    embeddings_2: List[OperationRecord] = []

    consume_fn_1 = CapturingConsumeFn()
    consume_fn_2 = CapturingConsumeFn()

    consumer.subscribe(collection_1, consume_fn_1, start=consumer.min_seqid())
    consumer.subscribe(collection_2, consume_fn_2, start=consumer.min_seqid())

    for i in range(10):
        e_1 = next(sample_embeddings)
        embeddings_1.append(e_1)
        producer.submit_embedding(collection_1, e_1)
        results_2 = await consume_fn_1.get(i + 1)
        assert_records_match(embeddings_1, results_2)

        e_2 = next(sample_embeddings)
        embeddings_2.append(e_2)
        producer.submit_embedding(collection_2, e_2)
        results_2 = await consume_fn_2.get(i + 1)
        assert_records_match(embeddings_2, results_2)


@pytest.mark.asyncio
async def test_start_seq_id(
    producer_consumer: Tuple[Producer, Consumer],
    sample_embeddings: Iterator[OperationRecord],
    produce_fns: ProducerFn,
) -> None:
    producer, consumer = producer_consumer
    producer.reset_state()
    consumer.reset_state()
    collection = UUID("00000000-0000-0000-0000-000000000000")

    consume_fn_1 = CapturingConsumeFn()
    consume_fn_2 = CapturingConsumeFn()

    consumer.subscribe(collection, consume_fn_1, start=consumer.min_seqid())

    embeddings = produce_fns(producer, collection, sample_embeddings, 5)[0]

    results_1 = await consume_fn_1.get(5)
    assert_records_match(embeddings, results_1)

    start = consume_fn_1.embeddings[-1]["log_offset"]
    consumer.subscribe(collection, consume_fn_2, start=start)
    second_embeddings = produce_fns(producer, collection, sample_embeddings, 5)[0]
    assert isinstance(embeddings, list)
    embeddings.extend(second_embeddings)
    results_2 = await consume_fn_2.get(5)
    assert_records_match(embeddings[-5:], results_2)


@pytest.mark.asyncio
async def test_end_seq_id(
    producer_consumer: Tuple[Producer, Consumer],
    sample_embeddings: Iterator[OperationRecord],
    produce_fns: ProducerFn,
) -> None:
    producer, consumer = producer_consumer
    producer.reset_state()
    consumer.reset_state()
    collection = UUID("00000000-0000-0000-0000-000000000000")

    consume_fn_1 = CapturingConsumeFn()
    consume_fn_2 = CapturingConsumeFn()

    consumer.subscribe(collection, consume_fn_1, start=consumer.min_seqid())

    embeddings = produce_fns(producer, collection, sample_embeddings, 10)[0]

    results_1 = await consume_fn_1.get(10)
    assert_records_match(embeddings, results_1)

    end = consume_fn_1.embeddings[-5]["log_offset"]
    consumer.subscribe(collection, consume_fn_2, start=consumer.min_seqid(), end=end)

    results_2 = await consume_fn_2.get(6)
    assert_records_match(embeddings[:6], results_2)

    # Should never produce a 7th
    with pytest.raises(TimeoutError):
        _ = await wait_for(consume_fn_2.get(7), timeout=1)


@pytest.mark.asyncio
async def test_submit_batch(
    producer_consumer: Tuple[Producer, Consumer],
    sample_embeddings: Iterator[OperationRecord],
) -> None:
    producer, consumer = producer_consumer
    producer.reset_state()
    consumer.reset_state()
    collection = UUID("00000000-0000-0000-0000-000000000000")

    embeddings = [next(sample_embeddings) for _ in range(100)]

    producer.submit_embeddings(collection, embeddings=embeddings)

    consume_fn = CapturingConsumeFn()
    consumer.subscribe(collection, consume_fn, start=consumer.min_seqid())

    recieved = await consume_fn.get(100)
    assert_records_match(embeddings, recieved)


@pytest.mark.asyncio
async def test_multiple_collections_batch(
    producer_consumer: Tuple[Producer, Consumer],
    sample_embeddings: Iterator[OperationRecord],
    produce_fns: ProducerFn,
) -> None:
    producer, consumer = producer_consumer
    producer.reset_state()
    consumer.reset_state()

    N_TOPICS = 2
    consume_fns = [CapturingConsumeFn() for _ in range(N_TOPICS)]
    for i in range(N_TOPICS):
        consumer.subscribe(
            UUID(f"00000000-0000-0000-0000-00000000000{i}"),
            consume_fns[i],
            start=consumer.min_seqid(),
        )

    embeddings_n: List[List[OperationRecord]] = [[] for _ in range(N_TOPICS)]

    PRODUCE_BATCH_SIZE = 10
    N_TO_PRODUCE = 100
    total_produced = 0
    for i in range(N_TO_PRODUCE // PRODUCE_BATCH_SIZE):
        for n in range(N_TOPICS):
            embeddings_n[n].extend(
                produce_fns(
                    producer,
                    UUID(f"00000000-0000-0000-0000-00000000000{n}"),
                    sample_embeddings,
                    PRODUCE_BATCH_SIZE,
                )[0]
            )
            recieved = await consume_fns[n].get(total_produced + PRODUCE_BATCH_SIZE)
            assert_records_match(embeddings_n[n], recieved)
        total_produced += PRODUCE_BATCH_SIZE


@pytest.mark.asyncio
async def test_max_batch_size(
    producer_consumer: Tuple[Producer, Consumer],
    sample_embeddings: Iterator[OperationRecord],
) -> None:
    producer, consumer = producer_consumer
    producer.reset_state()
    consumer.reset_state()
    collection = UUID("00000000-0000-0000-0000-000000000000")
    max_batch_size = producer.max_batch_size
    assert max_batch_size > 0

    # Make sure that we can produce a batch of size max_batch_size
    embeddings = [next(sample_embeddings) for _ in range(max_batch_size)]
    consume_fn = CapturingConsumeFn()
    consumer.subscribe(collection, consume_fn, start=consumer.min_seqid())
    producer.submit_embeddings(collection, embeddings=embeddings)
    received = await consume_fn.get(max_batch_size, timeout_secs=120)
    assert_records_match(embeddings, received)

    embeddings = [next(sample_embeddings) for _ in range(max_batch_size + 1)]
    # Make sure that we can't produce a batch of size > max_batch_size
    with pytest.raises(BatchSizeExceededError) as e:
        producer.submit_embeddings(collection, embeddings=embeddings)
    assert "Cannot submit more than" in str(e.value)