# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

from dataclasses import dataclass
from typing import Mapping, Optional

import ctypes
import numpy as np

from . import _ffi
from ._ffi import lib

INDEX_TYPES = {
    0: "ivf_flat",
    1: "ivf_pq",
    2: "ivf_hnsw_flat",
    3: "ivf_hnsw_sq",
    4: "ivf_rq",
}

METRICS = {
    0: "l2",
    1: "inner_product",
    2: "cosine",
}


@dataclass(frozen=True)
class VectorIndexMetadata:
    index_type: str
    dimension: int
    nlist: int
    metric: str
    total_vectors: int
    pq_m: Optional[int] = None
    hnsw_m: Optional[int] = None
    hnsw_ef_construction: Optional[int] = None
    hnsw_max_level: Optional[int] = None


@dataclass(frozen=True)
class SearchParams:
    top_k: int
    nprobe: int
    ef_search: int = 0
    query_bits: int = 0

    def to_ffi(self):
        return _ffi.PaimonVindexSearchParams(
            self.top_k,
            self.nprobe,
            self.ef_search,
            self.query_bits,
        )


def _check_error(message="operation failed"):
    err = lib.paimon_vindex_last_error()
    if err:
        raise RuntimeError(err.decode("utf-8", errors="replace"))
    raise RuntimeError(message)


def _metadata_from_ffi(raw):
    return VectorIndexMetadata(
        index_type=INDEX_TYPES.get(raw.index_type, f"unknown_{raw.index_type}"),
        dimension=raw.dimension,
        nlist=raw.nlist,
        metric=METRICS.get(raw.metric, f"unknown_{raw.metric}"),
        total_vectors=raw.total_vectors,
        pq_m=raw.pq_m or None,
        hnsw_m=raw.hnsw_m or None,
        hnsw_ef_construction=raw.hnsw_ef_construction or None,
        hnsw_max_level=raw.hnsw_max_level or None,
    )


def _float32_matrix(value, name):
    array = np.asarray(value, dtype=np.float32)
    if array.ndim != 2:
        raise ValueError(f"{name} must be a two-dimensional float32 array")
    return np.ascontiguousarray(array)


def _float32_vector(value, name):
    array = np.asarray(value, dtype=np.float32)
    if array.ndim != 1:
        raise ValueError(f"{name} must be a one-dimensional float32 array")
    return np.ascontiguousarray(array)


def _int64_vector(value, name):
    array = np.asarray(value, dtype=np.int64)
    if array.ndim != 1:
        raise ValueError(f"{name} must be a one-dimensional int64 array")
    return np.ascontiguousarray(array)


def _bytes_buffer(value, name):
    if isinstance(value, memoryview):
        value = value.tobytes()
    if not isinstance(value, (bytes, bytearray)):
        raise ValueError(f"{name} must be bytes")
    data = bytes(value)
    if not data:
        return None, 0, data
    buf = (ctypes.c_uint8 * len(data)).from_buffer_copy(data)
    return buf, len(data), data


def _option_arrays(options: Mapping[str, str]):
    option_items = list(options.items())
    key_bytes = []
    value_bytes = []
    for key, value in option_items:
        if not isinstance(key, str) or not isinstance(value, str):
            raise ValueError("options must be a mapping of str to str")
        key_bytes.append(key.encode("utf-8"))
        value_bytes.append(value.encode("utf-8"))
    keys = (ctypes.c_char_p * len(key_bytes))(*key_bytes)
    values = (ctypes.c_char_p * len(value_bytes))(*value_bytes)
    return option_items, key_bytes, value_bytes, keys, values


class VectorIndexTraining:
    def __init__(self, handle):
        self._closed = False
        self._handle = handle

    def _require_open(self):
        if self._closed or not self._handle:
            raise RuntimeError("VectorIndexTraining is closed")

    def _take_handle(self):
        self._require_open()
        handle = self._handle
        self._handle = None
        self._closed = True
        return handle

    def close(self):
        if self._handle:
            lib.paimon_vindex_training_free(self._handle)
            self._handle = None
        self._closed = True

    def __enter__(self):
        self._require_open()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


class VectorIndexTrainer:
    def __init__(self, options: Mapping[str, str]):
        self._closed = False
        (
            option_items,
            self._key_bytes,
            self._value_bytes,
            self._keys,
            self._values,
        ) = _option_arrays(options)
        self._handle = lib.paimon_vindex_trainer_open(
            self._keys,
            self._values,
            len(option_items),
        )
        if not self._handle:
            _check_error("failed to open trainer")
        self._dimension = self._read_dimension()

    @classmethod
    def create(cls, options: Mapping[str, str]):
        return cls(options)

    @classmethod
    def train(cls, options: Mapping[str, str], data):
        with cls(options) as trainer:
            return trainer.add_training_vectors(data).finish_training()

    def _require_open(self):
        if self._closed or not self._handle:
            raise RuntimeError("VectorIndexTrainer is closed")

    def _read_dimension(self):
        out = ctypes.c_size_t(0)
        rc = lib.paimon_vindex_trainer_dimension(self._handle, ctypes.byref(out))
        if rc != 0:
            _check_error("trainer dimension failed")
        return out.value

    @property
    def dimension(self):
        self._require_open()
        return self._dimension

    def add_training_vectors(self, data):
        self._require_open()
        data = _float32_matrix(data, "data")
        if data.shape[1] != self._dimension:
            raise RuntimeError(
                f"training data length {data.size} does not match vector count "
                f"* dimension {data.shape[0] * self._dimension}"
            )
        rc = lib.paimon_vindex_trainer_add_training_vectors(
            self._handle,
            data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
            data.shape[0],
        )
        if rc != 0:
            _check_error("add training vectors failed")
        return self

    def finish_training(self):
        self._require_open()
        handle = self._handle
        training = lib.paimon_vindex_trainer_finish(handle)
        lib.paimon_vindex_trainer_free(handle)
        self._handle = None
        self._closed = True
        if not training:
            _check_error("finish training failed")
        return VectorIndexTraining(training)

    def close(self):
        if self._handle:
            lib.paimon_vindex_trainer_free(self._handle)
            self._handle = None
        self._closed = True

    def __enter__(self):
        self._require_open()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


class VectorIndexWriter:
    def __init__(self, training: VectorIndexTraining):
        if not isinstance(training, VectorIndexTraining):
            raise TypeError("training must be a VectorIndexTraining")
        self._closed = False
        training_handle = training._take_handle()
        self._handle = lib.paimon_vindex_writer_open(training_handle)
        lib.paimon_vindex_training_free(training_handle)
        if not self._handle:
            _check_error("failed to open writer")
        self._dimension = self._read_dimension()

    def _require_open(self):
        if self._closed or not self._handle:
            raise RuntimeError("VectorIndexWriter is closed")

    def _read_dimension(self):
        out = ctypes.c_size_t(0)
        rc = lib.paimon_vindex_writer_dimension(self._handle, ctypes.byref(out))
        if rc != 0:
            _check_error("writer dimension failed")
        return out.value

    @property
    def dimension(self):
        self._require_open()
        return self._dimension

    def add_vectors(self, ids, data):
        self._require_open()
        data = _float32_matrix(data, "data")
        ids = _int64_vector(ids, "ids")
        if data.shape[1] != self._dimension:
            raise RuntimeError(
                f"vector data length {data.size} does not match vector count "
                f"* dimension {data.shape[0] * self._dimension}"
            )
        if ids.shape[0] != data.shape[0]:
            raise RuntimeError(
                f"ids length {ids.shape[0]} does not match vector count {data.shape[0]}"
            )
        rc = lib.paimon_vindex_writer_add_vectors(
            self._handle,
            ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)),
            data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
            data.shape[0],
        )
        if rc != 0:
            _check_error("add_vectors failed")

    def write(self, file):
        self._require_open()
        pos = 0

        @_ffi.WRITE_FN
        def write_callback(ctx, data, length):
            nonlocal pos
            try:
                payload = ctypes.string_at(data, length)
                written = file.write(payload)
                if written is not None and written != length:
                    return -1
                pos += length
                return 0
            except Exception:
                return -1

        @_ffi.FLUSH_FN
        def flush_callback(ctx):
            try:
                flush = getattr(file, "flush", None)
                if flush is not None:
                    flush()
                return 0
            except Exception:
                return -1

        @_ffi.GET_POS_FN
        def pos_callback(ctx):
            return pos

        output = _ffi.PaimonVindexOutputFile()
        output.ctx = None
        output.write_fn = write_callback
        output.flush_fn = flush_callback
        output.get_pos_fn = pos_callback

        rc = lib.paimon_vindex_writer_write_index(self._handle, output)
        if rc != 0:
            _check_error("write index failed")

    def close(self):
        if self._handle:
            lib.paimon_vindex_writer_free(self._handle)
            self._handle = None
        self._closed = True

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


class VectorIndexReader:
    def __init__(self, input):
        self._input = input
        self._closed = False

        @_ffi.READ_AT_FN
        def read_at_callback(ctx, offset, buf, length):
            try:
                chunks = self._input.pread_many([(offset, length)])
                if len(chunks) != 1:
                    return -1
                data = bytes(chunks[0])
                if len(data) != length:
                    return -1
                ctypes.memmove(buf, data, length)
                return 0
            except Exception:
                return -1

        self._read_at_callback = read_at_callback
        input_file = _ffi.PaimonVindexInputFile()
        input_file.ctx = None
        input_file.read_at_fn = self._read_at_callback
        self._handle = lib.paimon_vindex_reader_open(input_file)
        if not self._handle:
            _check_error("failed to open reader")
        self._metadata = self.metadata()

    def _require_open(self):
        if self._closed or not self._handle:
            raise RuntimeError("VectorIndexReader is closed")

    @property
    def index_type(self):
        return self.metadata().index_type

    @property
    def dimension(self):
        return self.metadata().dimension

    @property
    def nlist(self):
        return self.metadata().nlist

    @property
    def total_vectors(self):
        return self.metadata().total_vectors

    def metadata(self):
        self._require_open()
        raw = _ffi.PaimonVindexMetadata()
        rc = lib.paimon_vindex_reader_metadata(self._handle, ctypes.byref(raw))
        if rc != 0:
            _check_error("metadata failed")
        return _metadata_from_ffi(raw)

    def optimize_for_search(self):
        self._require_open()
        rc = lib.paimon_vindex_reader_optimize_for_search(self._handle)
        if rc != 0:
            _check_error("optimize_for_search failed")

    def _filter_args(self, filter_bytes):
        if filter_bytes is None:
            return None, 0, None
        return _bytes_buffer(filter_bytes, "filter_bytes")

    def search(self, query, params: SearchParams, filter_bytes=None):
        self._require_open()
        query = _float32_vector(query, "query")
        if query.shape[0] != self._metadata.dimension:
            raise RuntimeError(
                f"query length {query.shape[0]} does not match index dimension "
                f"{self._metadata.dimension}"
            )
        ffi_params = params.to_ffi()
        ids = np.empty(params.top_k, dtype=np.int64)
        distances = np.empty(params.top_k, dtype=np.float32)

        if filter_bytes is None:
            rc = lib.paimon_vindex_reader_search(
                self._handle,
                query.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
                ffi_params,
                ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)),
                distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
                params.top_k,
            )
        else:
            filter_buf, filter_len, _ = self._filter_args(filter_bytes)
            rc = lib.paimon_vindex_reader_search_with_roaring_filter(
                self._handle,
                query.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
                ffi_params,
                filter_buf,
                filter_len,
                ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)),
                distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
                params.top_k,
            )
        if rc != 0:
            _check_error("search failed")
        return ids, distances

    def search_batch(self, queries, params: SearchParams, filter_bytes=None):
        self._require_open()
        queries = _float32_matrix(queries, "queries")
        if queries.shape[1] != self._metadata.dimension:
            raise RuntimeError(
                f"queries length {queries.size} does not match nq * dimension "
                f"{queries.shape[0] * self._metadata.dimension}"
            )
        ffi_params = params.to_ffi()
        result_len = queries.shape[0] * params.top_k
        ids = np.empty((queries.shape[0], params.top_k), dtype=np.int64)
        distances = np.empty((queries.shape[0], params.top_k), dtype=np.float32)

        if filter_bytes is None:
            rc = lib.paimon_vindex_reader_search_batch(
                self._handle,
                queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
                queries.shape[0],
                ffi_params,
                ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)),
                distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
                result_len,
            )
        else:
            filter_buf, filter_len, _ = self._filter_args(filter_bytes)
            rc = lib.paimon_vindex_reader_search_batch_with_roaring_filter(
                self._handle,
                queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
                queries.shape[0],
                ffi_params,
                filter_buf,
                filter_len,
                ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)),
                distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
                result_len,
            )
        if rc != 0:
            _check_error("batch search failed")
        return ids, distances

    def close(self):
        if self._handle:
            lib.paimon_vindex_reader_free(self._handle)
            self._handle = None
        self._closed = True

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


__all__ = [
    "VectorIndexMetadata",
    "VectorIndexReader",
    "VectorIndexTrainer",
    "VectorIndexTraining",
    "VectorIndexWriter",
]