import logging
import time

import docker
from opensearchpy import OpenSearch, helpers

from hseb.core.config import Config, IndexArgs, QuantDatatype, SearchArgs
from hseb.core.dataset import Doc, Query
from hseb.core.response import DocScore, Response
from hseb.engine.base import EngineBase

logger = logging.getLogger()

OS_DATATYPES = {
    QuantDatatype.FLOAT32: "float",
    QuantDatatype.INT8: "byte",
    QuantDatatype.INT1: "binary",
}


class OpenSearchEngine(EngineBase):
    def __init__(self, config: Config):
        self.config = config

    def start(self, index_args: IndexArgs):
        self.index_args = index_args
        docker_client = docker.from_env()
        heap_size = index_args.kwargs.get("heap_size", "8g")
        self.container = docker_client.containers.run(
            image=self.config.image,
            ports={"9200/tcp": 9200, "9300/tcp": 9300},
            detach=True,
            environment={
                "discovery.type": "single-node",
                "OPENSEARCH_JAVA_OPTS": f"-Xms{heap_size} -Xmx{heap_size}",
                "DISABLE_SECURITY_PLUGIN": "true",
                "bootstrap.memory_lock": "true",
            },
        )
        self._wait_for_logs(self.container, "Cluster health status changed from [YELLOW] to [GREEN]")
        self.client = OpenSearch(
            hosts=[{"host": "localhost", "port": 9200}],
            http_compress=True,
            use_ssl=False,
            verify_certs=False,
            ssl_assert_hostname=False,
            ssl_show_warn=False,
            timeout=30,
        )
        self.client.indices.create(
            index="test",
            body={
                "settings": {
                    "index": {
                        "knn": True,
                        "refresh_interval": "1h",
                        "merge": {
                            "policy": {"max_merged_segment": index_args.kwargs.get("max_merged_segment", "128mb")}
                        },
                    },
                },
                "mappings": {
                    "properties": {
                        "text": {
                            "type": "knn_vector",
                            "dimension": self.config.dataset.dim,
                            "space_type": "innerproduct",
                            "data_type": OS_DATATYPES[index_args.quant],
                            "method": {
                                "name": "hnsw",
                                "engine": "lucene",
                                "parameters": {
                                    "m": index_args.m,
                                    "ef_construction": index_args.ef_construction,
                                },
                            },
                        },
                        "tag": {"type": "integer"},
                    }
                },
            },
        )
        self.docs_in_segment = 0

    def stop(self):
        self.container.stop()

    def commit(self):
        self.client.indices.refresh(index="test")

    def index_batch(self, batch: list[Doc]):
        actions = []
        for doc in batch:
            actions.append(
                {
                    "_op_type": "index",
                    "_index": "test",
                    "_source": {
                        "text": doc.embedding.tolist(),
                        "tag": doc.tag,
                    },
                    "_id": doc.id,
                }
            )
        helpers.bulk(self.client, actions)
        self.docs_in_segment += len(batch)
        if self.docs_in_segment >= self.index_args.kwargs.get("docs_per_segment", 1024):
            self.client.indices.refresh(index="test")
            self.docs_in_segment = 0

    def search(self, search_params: SearchArgs, query: Query, top_k: int) -> Response:
        knn_query = {
            "text": {
                "vector": query.embedding.tolist(),
                "k": top_k,
                "method_parameters": {
                    "ef_search": search_params.ef_search,
                },
            }
        }

        if search_params.filter_selectivity != 100:
            knn_query["text"]["filter"] = {"terms": {"tag": [search_params.filter_selectivity]}}

        search_body = {
            "size": top_k,
            "query": {"knn": knn_query},
            "_source": False,
        }

        start = time.time_ns()
        response = self.client.search(index="test", body=search_body)
        end = time.time_ns()
        return Response(
            results=[DocScore(doc=int(doc["_id"]), score=doc["_score"]) for doc in response["hits"]["hits"]],
            client_latency=(end - start) / 1000000000.0,
        )
