"""End-to-end tests for the Weaviate adaptor."""

from __future__ import annotations

import time
import uuid

import docker
import numpy as np
import pytest
import requests
from datasets import Dataset
from docker.errors import APIError, DockerException

from inatinqperf.adaptors.base import DataPoint, Query
from inatinqperf.adaptors.enums import Metric
from inatinqperf.adaptors.weaviate_adaptor import Weaviate

# Mark these as integration tests and skip if Docker is unavailable.


def _docker_available() -> bool:
    """Return True when the Docker daemon is reachable."""

    try:
        client = docker.from_env()
        try:
            client.ping()
        finally:
            client.close()
    except DockerException:
        return False
    return True


pytestmark = pytest.mark.skipif(
    not _docker_available(), reason="Docker daemon unavailable for integration tests"
)

# These tests talk to a real Weaviate instance listening on localhost.
BASE_URL = "http://localhost:8080"
WEAVIATE_IMAGE = "semitechnologies/weaviate:1.31.16-ab5cb66.arm64"
# Minimal command to expose the HTTP interface on localhost:8080.
WEAVIATE_COMMAND = ["--host", "0.0.0.0", "--port", "8080", "--scheme", "http"]
# Environment mirrors the configuration used during benchmark runs.
WEAVIATE_ENVIRONMENT = {
    "AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED": "true",
    "ENABLE_MODULES": "",
    "DEFAULT_VECTORIZER_MODULE": "none",
    "PERSISTENCE_DATA_PATH": "/var/lib/weaviate",
    "QUERY_DEFAULTS_LIMIT": "20",
}


def wait_for_weaviate(timeout_seconds: int = 120) -> None:
    """Poll the readiness endpoint until the container reports healthy."""

    deadline = time.time() + timeout_seconds
    endpoint = f"{BASE_URL}/v1/.well-known/ready"
    while time.time() < deadline:
        try:
            response = requests.get(endpoint, timeout=2)
            if response.status_code == 200:
                return
        except requests.RequestException:
            pass  # container still booting -> retry after a short sleep
        time.sleep(1)
    msg = "Weaviate service did not become ready within the allotted timeout."
    raise RuntimeError(msg)


@pytest.fixture(scope="module", autouse=True)
def weaviate_container():
    """Ensure a Weaviate container is available for the duration of the tests."""
    # Guard against environments where the daemon/socket is not accessible.
    try:
        client = docker.from_env()
    except DockerException as exc:
        pytest.skip(f"docker daemon unavailable: {exc}")
    container = None
    started = False
    try:
        # Reuse an already-running instance if one is reachable.
        try:
            response = requests.get(f"{BASE_URL}/v1/.well-known/ready", timeout=2)
            if response.status_code == 200:
                yield
                return
        except requests.RequestException:
            pass

        try:
            container = client.containers.run(  # spin up the test instance
                WEAVIATE_IMAGE,
                ports={"8080": "8080", "8081": "8081"},
                environment=WEAVIATE_ENVIRONMENT,
                command=WEAVIATE_COMMAND,
                remove=True,
                detach=True,
            )
            started = True
        except APIError as exc:  # container already running or port busy
            if "port is already allocated" not in str(exc).lower():
                pytest.skip(f"unable to start weaviate container: {exc}")
        wait_for_weaviate()
        yield
    finally:
        if started and container is not None:
            container.stop()


@pytest.fixture(name="class_name")
def class_name_fixture() -> str:
    """Generate a valid Weaviate class name for isolation."""
    prefix = "Perf"
    suffix = uuid.uuid4().hex[:10].capitalize()
    return f"{prefix}{suffix}"


def make_dataset(dim: int) -> Dataset:
    """Create a simple dataset placeholder with the desired dimensionality."""

    return Dataset.from_dict({"embedding": [[0.0] * dim]})


def to_datapoints(ids: np.ndarray, vectors: np.ndarray) -> list[DataPoint]:
    """Convert arrays of ids and vectors into DataPoint instances."""

    return [
        DataPoint(id=int(identifier), vector=vector.tolist(), metadata={})
        for identifier, vector in zip(ids, vectors, strict=False)
    ]


def test_weaviate_adaptor_lifecycle(class_name: str):
    """Exercise the full lifecycle against a live Weaviate instance."""
    adaptor = Weaviate(
        dataset=make_dataset(4), metric=Metric.COSINE, url=BASE_URL, collection_name=class_name
    )
    adaptor.drop_index()  # ensure clean start

    adaptor._ensure_schema_exists()

    ids = np.arange(100, 104, dtype=np.int64)
    vectors = np.array(
        [
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 0.0],
            [0.0, 0.0, 0.0, 1.0],
        ],
        dtype=np.float32,
    )
    adaptor.upsert(to_datapoints(ids, vectors))

    stats = adaptor.stats()  # GraphQL Aggregate meta-count
    assert stats["ntotal"] == len(ids)
    assert stats["class_name"] == class_name

    query = Query(vector=[1.0, 0.1, 0.0, 0.0])
    results = adaptor.search(query, topk=3)  # GraphQL nearVector search
    assert len(results) == 3
    assert results[0].id == 100

    adaptor.delete([100, 999])  # data_object.delete
    expected_total = len(ids) - 1
    deadline = time.time() + 10.0
    stats_after_delete = adaptor.stats()
    while stats_after_delete["ntotal"] != expected_total and time.time() < deadline:
        time.sleep(0.5)
        stats_after_delete = adaptor.stats()
    assert stats_after_delete["ntotal"] == expected_total

    adaptor.drop_index()


def test_weaviate_upsert_replaces_vectors(class_name: str):
    """Verify upsert overwrites existing vectors rather than duplicating them."""
    adaptor = Weaviate(
        dataset=make_dataset(3), metric=Metric.COSINE, url=BASE_URL, collection_name=class_name
    )
    # Drop any leftover schema, then create a fresh one for the test.
    adaptor.drop_index()
    adaptor._ensure_schema_exists()

    ids = np.array([1, 2], dtype=np.int64)
    first = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
    adaptor.upsert(to_datapoints(ids, first))
    second = first + 0.5
    adaptor.upsert(to_datapoints(ids, second))

    stats = adaptor.stats()  # Aggregate meta count after replacement
    assert stats["ntotal"] == len(ids)

    # Ensure subsequent search surfaces the updated vectors (id=2 closer to query)
    query = Query(vector=[0.0, 1.0, 0.5])
    results = adaptor.search(query, topk=2)
    assert results[0].id == 2

    adaptor.drop_index()


def test_weaviate_metric_mapping(class_name: str):
    """Confirm metric names map to Weaviate's expected distance types."""

    ids = np.array([10, 11], dtype=np.int64)
    vectors = np.array([[1.0, 0.0], [0.5, 0.5]], dtype=np.float32)

    for metric, expected in (
        (Metric.INNER_PRODUCT, "dot"),
        (Metric.COSINE, "cosine"),
        (Metric.L2, "l2-squared"),
    ):
        adaptor = Weaviate(
            dataset=make_dataset(2),
            metric=metric,
            url=BASE_URL,
            collection_name=f"{class_name}{metric.value.title()}",
            index_type="hnsw",
        )
        adaptor.drop_index()
        adaptor._ensure_schema_exists()

        schema = adaptor._client.schema.get()
        classes = {entry["class"]: entry for entry in schema.get("classes", [])}
        current = classes.get(adaptor.collection_name, {})
        assert current.get("vectorIndexType") == "hnsw"

        adaptor.upsert(to_datapoints(ids, vectors))

        query = Query(vector=[0.6, 0.4])
        results = adaptor.search(query, topk=2)
        assert len(results) == 2
        assert set(result.id for result in results) == {10, 11}

        adaptor.drop_index()


@pytest.mark.parametrize("index_type, expected", [(None, "hnsw"), ("flat", "flat")])
def test_weaviate_index_types(class_name: str, index_type: str | None, expected: str) -> None:
    """Ensure requested index types are persisted when the schema is created."""

    adaptor = Weaviate(
        dataset=make_dataset(2),
        metric=Metric.COSINE,
        url=BASE_URL,
        collection_name=f"{class_name}Idx",
        index_type=index_type,
    )
    adaptor.drop_index()
    adaptor._ensure_schema_exists()

    schema = adaptor._client.schema.get()
    classes = {entry["class"]: entry for entry in schema.get("classes", [])}
    payload = classes.get(adaptor.collection_name, {})
    assert payload.get("vectorIndexType") == expected

    adaptor.drop_index()
