# Copyright 2025 The RLinf Authors.
#
# 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
#
#     https://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 argparse
import asyncio
import json
import logging
import os
import socket
import time
from concurrent.futures import ProcessPoolExecutor
from multiprocessing import Queue, set_start_method
from typing import Optional

import numpy as np
import torch
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from qdrant_client import AsyncQdrantClient
from qdrant_client.models import (
    QuantizationSearchParams,
    SearchParams,
)
from qdrant_encoder import Encoder
from tqdm import tqdm


class AsyncEncoderPool:
    @staticmethod
    def set_global_encoder(init_queue: Queue):
        args = init_queue.get()
        assert "global_encoder" not in globals()
        globals()["global_encoder"] = Encoder(*args)

    @staticmethod
    def global_encode(*args, **kwargs):
        assert "global_encoder" in globals()
        encoder: Encoder = globals()["global_encoder"]
        return encoder.encode(*args, **kwargs)

    def __init__(
        self, model_name, model_path, pooling_method, max_length, use_fp16, devices
    ):
        init_queue = Queue()
        for device in devices:
            init_queue.put(
                [model_name, model_path, pooling_method, max_length, use_fp16, device]
            )

        self.encoders = ProcessPoolExecutor(
            max_workers=len(devices),
            initializer=AsyncEncoderPool.set_global_encoder,
            initargs=(init_queue,),
        )

    async def encode(self, query_list: list[str], is_query=True) -> np.ndarray:
        return await loop.run_in_executor(
            self.encoders, AsyncEncoderPool.global_encode, (query_list, is_query)
        )


class AsyncBaseRetriever:
    def __init__(self, config):
        self.config = config
        self.retrieval_method = config.retrieval_method
        self.topk = config.retrieval_topk

    async def _asearch(self, query: str, num: int, return_score: bool):
        raise NotImplementedError

    async def _abatch_search(self, query_list: list[str], num: int, return_score: bool):
        raise NotImplementedError

    async def asearch(
        self, query: str, num: int | None = None, return_score: bool = False
    ):
        return await self._asearch(query, num, return_score)

    async def abatch_search(
        self, query_list: list[str], num: int | None = None, return_score: bool = False
    ):
        return await self._abatch_search(query_list, num, return_score)


class AsyncDenseRetriever(AsyncBaseRetriever):
    @staticmethod
    async def wait_qdrant_load(url, connect_timeout):
        client = AsyncQdrantClient(url=url, prefer_grpc=True, timeout=60)
        wait_collection_time = 0
        while True:
            if wait_collection_time >= connect_timeout:
                assert False, f"wait longer than {connect_timeout}s, exit"
            logging.info(f"wait {wait_collection_time}s for qdrant load")
            time.sleep(5)
            wait_collection_time += 5
            try:
                await client.info()
                logging.info("qdrant loaded and connected")
                break
            except Exception:
                pass
        return client

    def __init__(self, config: "Config"):
        super().__init__(config)

    async def ainit(self, config: "Config"):
        self.client = await self.wait_qdrant_load(
            url=config.qdrant_url, connect_timeout=300
        )

        self.collection_name = config.qdrant_collection_name
        collections = (await self.client.get_collections()).collections
        collection_names = [col.name for col in collections]
        assert self.collection_name in collection_names, (
            f"no collection! exists: [{collection_names}], need: {self.collection_name}"
        )

        from qdrant_client.http.models import Batch
        from qdrant_client.models import CollectionStatus

        coll_status = (await self.client.get_collection(self.collection_name)).status
        if coll_status != CollectionStatus.GREEN:
            wait_collection_time = 0
            optimize_timeout = 3000
            await self.client.upsert(
                collection_name=self.collection_name,
                points=Batch(ids=[], vectors=[]),
            )
            logging.info(
                f"Optimizers triggered for collection '{self.collection_name}' with an empty update operation."
            )
            while coll_status != CollectionStatus.GREEN:
                time.sleep(5)
                if wait_collection_time >= optimize_timeout:
                    assert False, f"wait longer than {optimize_timeout}s, exit"
                logging.info(f"wait {wait_collection_time}s for qdrant optimize")
            logging.info("qdrant optimized")
        else:
            logging.info("collection status is green now")

        # Initialize encoder first (needed for building collection)
        devices = [
            torch.device(f"cuda:{i}") for i in range(0, torch.cuda.device_count())
        ]
        self.encoder = AsyncEncoderPool(
            model_name=self.retrieval_method,
            model_path=config.retrieval_model_path,
            pooling_method=config.retrieval_pooling_method,
            max_length=config.retrieval_query_max_length,
            use_fp16=config.retrieval_use_fp16,
            devices=devices,
        )
        self.topk = config.retrieval_topk
        if config.qdrant_search_quant_param is not None:
            self.search_params = SearchParams(
                **json.loads(config.qdrant_search_param),
                quantization=QuantizationSearchParams(
                    **json.loads(config.qdrant_search_quant_param)
                ),
            )
        else:
            self.search_params = SearchParams(
                **json.loads(config.qdrant_search_param),
            )
        logging.info(f"qdrant search_params: {self.search_params}")

    async def _asearch(
        self, query: str, num: int | None = None, return_score: bool = False
    ):
        time_start = time.time()
        if num is None:
            num = self.topk
        query_emb = await self.encoder.encode(query)
        query_vector = query_emb[0].tolist()
        time_embed = time.time()

        # Search in Qdrant
        search_results = (
            await self.client.query_points(
                collection_name=self.collection_name,
                query=query_vector,
                limit=num,
                search_params=self.search_params,
            )
        ).points
        time_search = time.time()
        time_elapse_search = time_search - time_embed
        time_elapse_embed = time_embed - time_start
        logging.info(
            f"time elapse: search: {time_elapse_search}; embed: {time_elapse_embed}"
        )

        # Extract IDs and scores
        payloads = [result.payload for result in search_results]
        scores = [result.score for result in search_results]
        if return_score:
            return payloads, scores
        else:
            return payloads

    async def _abatch_search(
        self, query_list: list[str], num: int | None = None, return_score: bool = False
    ):
        if return_score:
            all_payloads, all_scores = [], []
            for query in query_list:
                payloads, scores = await self._asearch(query, num, return_score)
                all_payloads.append(payloads)
                all_scores.append(scores)
            return all_payloads, all_scores
        else:
            all_payloads = []
            for query in query_list:
                payloads = await self._asearch(query, num, return_score)
                all_payloads.append(payloads)
            return all_payloads


async def get_retriever(config):
    retriever = AsyncDenseRetriever(config)
    await retriever.ainit(config)
    return retriever


class PageAccess:
    def __init__(self, pages_path):
        pages = []
        with open(pages_path, "r") as f:
            for line in tqdm(f, desc="Loading pages"):
                pages.append(json.loads(line))
        self.pages = {page["url"]: page for page in pages}

    def access(self, url):
        # php parsing
        if "index.php/" in url:
            url = url.replace("index.php/", "index.php?title=")
        if url not in self.pages:
            return None
        return self.pages[url]


#####################################
# FastAPI server below
#####################################


class Config:
    """
    Minimal config class for local Qdrant retrieval server.(simulating your argparse)
    Replace this with your real arguments or load them dynamically.
    """

    def __init__(
        self,
        retrieval_method: str = "bm25",
        retrieval_topk: int = 10,
        dataset_path: str = "./data",
        data_split: str = "train",
        qdrant_url: Optional[str] = None,
        qdrant_collection_name: str = "default_collection",
        qdrant_search_param: Optional[str] = None,
        qdrant_search_quant_param: Optional[str] = None,
        retrieval_model_path: str = "./model",
        retrieval_pooling_method: str = "mean",
        retrieval_query_max_length: int = 256,
        retrieval_use_fp16: bool = False,
    ):
        self.retrieval_method = retrieval_method
        self.retrieval_topk = retrieval_topk
        self.dataset_path = dataset_path
        self.data_split = data_split
        self.qdrant_url = qdrant_url
        self.qdrant_collection_name = qdrant_collection_name
        self.qdrant_search_param = qdrant_search_param
        self.qdrant_search_quant_param = qdrant_search_quant_param
        self.retrieval_model_path = retrieval_model_path
        self.retrieval_pooling_method = retrieval_pooling_method
        self.retrieval_query_max_length = retrieval_query_max_length
        self.retrieval_use_fp16 = retrieval_use_fp16


class QueryRequest(BaseModel):
    queries: list[str]
    topk: Optional[int] = None
    return_scores: bool = False


class AccessRequest(BaseModel):
    urls: list[str]


app = FastAPI()


@app.post("/retrieve")
async def retrieve_endpoint(request: QueryRequest):
    """
    Endpoint that accepts queries and performs retrieval.
    Input format:
    {
      "queries": ["What is Python?", "Tell me about neural networks."],
      "topk": 3,
      "return_scores": true
    }
    """
    time_start = time.time()
    if not request.topk:
        request.topk = config.retrieval_topk  # fallback to default

    # Perform batch retrieval
    if request.return_scores:
        results, scores = await retriever.abatch_search(
            query_list=request.queries,
            num=request.topk,
            return_score=request.return_scores,
        )
    else:
        results = await retriever.abatch_search(
            query_list=request.queries,
            num=request.topk,
            return_score=request.return_scores,
        )

    # Format response
    resp = []
    for i, single_result in enumerate(results):
        if request.return_scores:
            # If scores are returned, combine them with results
            combined = []
            for doc, score in zip(single_result, scores[i]):
                combined.append({"document": doc, "score": score})
            resp.append(combined)
        else:
            resp.append(single_result)
    time_elapse = time.time() - time_start
    logging.info(f"request: {request}, time_elapse: {time_elapse}")
    return {"result": resp}


@app.post("/access")
async def access_endpoint(request: AccessRequest):
    resp = []
    for url in request.urls:
        resp.append(page_access.access(url))

    return {"result": resp}


if __name__ == "__main__":
    set_start_method("spawn")
    loop = asyncio.get_event_loop()

    parser = argparse.ArgumentParser(description="Launch the local qdrant retriever.")
    parser.add_argument("--pages_path", type=str, default=None, help="Local page file.")
    parser.add_argument(
        "--topk",
        type=int,
        default=3,
        help="Number of retrieved passages for one query.",
    )
    parser.add_argument(
        "--retriever_name", type=str, default="e5", help="Name of the retriever model."
    )
    parser.add_argument(
        "--retriever_model",
        type=str,
        required=True,
        help="Path of the retriever model.",
    )
    parser.add_argument(
        "--qdrant_url",
        type=str,
        default="http://localhost:6333",
        help="Qdrant server URL",
    )
    parser.add_argument(
        "--qdrant_collection_name",
        type=str,
        required=True,
        help="Name of the Qdrant collection.",
    )
    parser.add_argument(
        "--qdrant_search_param",
        type=str,
        default="{}",
        help="HNSW search parameters as JSON string (e.g., '{\"hnsw_ef\":256}')",
    )
    parser.add_argument(
        "--qdrant_search_quant_param",
        type=str,
        default=None,
        help="Quantization search parameters as JSON string (optional)",
    )
    parser.add_argument(
        "--port", type=int, default=8000, help="Port to run the server on"
    )
    parser.add_argument(
        "--save-address-to",
        type=str,
        default=None,
        help="Directory to save server address file (optional)",
    )

    args = parser.parse_args()
    logging.getLogger().setLevel(logging.INFO)
    host_name = socket.gethostname()
    host_ip = socket.gethostbyname(socket.gethostname())
    port = args.port

    host_addr = f"{host_ip}:{port}"

    logging.info(f"Server address: {host_addr}")

    if args.save_address_to:
        os.makedirs(args.save_address_to, exist_ok=True)
        with open(
            os.path.join(
                args.save_address_to, "Host" + host_ip + "_" + "IP" + str(port) + ".txt"
            ),
            "w",
        ) as f:
            f.write(host_addr)

    # 1) Build a config (could also parse from arguments).
    #    In real usage, you'd parse your CLI arguments or environment variables.
    config = Config(
        retrieval_method=args.retriever_name,
        retrieval_topk=args.topk,
        qdrant_url=args.qdrant_url,
        qdrant_collection_name=args.qdrant_collection_name,
        qdrant_search_param=args.qdrant_search_param,
        qdrant_search_quant_param=args.qdrant_search_quant_param,
        retrieval_model_path=args.retriever_model,
        retrieval_pooling_method="mean",
        retrieval_query_max_length=256,
        retrieval_use_fp16=True,
    )

    # 2) Instantiate a global retriever so it is loaded once and reused.
    retriever = loop.run_until_complete(get_retriever(config))

    async def test():
        query1 = "Tell me about Red Bull"
        result1 = await retriever.asearch(query1, 1, return_score=False)
        logging.info(f"test1: query: {query1}, result: {result1}")
        query2 = "Tell me about Ljubljana"
        result2 = await retriever.asearch(query2, 2, return_score=True)
        logging.info(f"test2: query: {query2}, result: {result2}")
        query3 = ["Tell me about Mars", "Tell me about Mercury"]
        result3 = await retriever.abatch_search(query3, 3, return_score=True)
        logging.info(f"test3: query: {query3}, result: {result3}")
        logging.info("Retriever is ready.")

    loop.run_until_complete(test())

    # 3) Load pages
    if not args.pages_path:
        logging.info("Page Access is off.")
    elif not os.path.exists(args.pages_path):
        logging.info(
            f"Page Access is not load because pages_path({args.pages_path}) not exist."
        )
    else:
        page_access = PageAccess(args.pages_path)
        logging.info("Page Access is ready.")

    # 4) Launch the server.
    config = uvicorn.Config(
        app,
        host="0.0.0.0",
        port=args.port,
        log_level="warning",
        loop=loop,
    )
    server = uvicorn.Server(config)
    logging.info(f"Server is ready at port {args.port}")

    loop.run_until_complete(server.serve())
