/*
 *  Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
 *  <p>
 *  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
 *  <p>
 *  http://www.apache.org/licenses/LICENSE-2.0
 *  <p>
 *  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.
 */
package com.easyagents.rerank;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.fastjson2.JSONPath;
import com.easyagents.core.document.Document;
import com.easyagents.core.model.client.HttpClient;
import com.easyagents.core.model.rerank.BaseRerankModel;
import com.easyagents.core.model.rerank.RerankException;
import com.easyagents.core.model.rerank.RerankOptions;
import com.easyagents.core.util.Maps;
import com.easyagents.core.util.StringUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class DefaultRerankModel extends BaseRerankModel<DefaultRerankModelConfig> {

    private static final Logger LOG = LoggerFactory.getLogger(DefaultRerankModel.class);
    private static final int MAX_RESPONSE_LOG_LENGTH = 512;

    private HttpClient httpClient = new HttpClient();

    public DefaultRerankModel(DefaultRerankModelConfig config) {
        super(config);
    }

    public HttpClient getHttpClient() {
        return httpClient;
    }

    public void setHttpClient(HttpClient httpClient) {
        this.httpClient = httpClient;
    }

    @Override
    public List<Document> rerank(String query, List<Document> documents, RerankOptions options) {

        DefaultRerankModelConfig config = getConfig();
        String url = config.getEndpoint() + config.getRequestPath();

        Map<String, String> headers = new HashMap<>(2);
        headers.put("Content-Type", "application/json");
        headers.put("Authorization", "Bearer " + config.getApiKey());

        List<String> payloadDocuments = new ArrayList<>(documents.size());
        List<Integer> documentIndexMapping = new ArrayList<>(documents.size());
        for (int i = 0; i < documents.size(); i++) {
            Document document = documents.get(i);
            if (document == null || StringUtil.noText(document.getContent())) {
                continue;
            }
            payloadDocuments.add(document.getContent());
            documentIndexMapping.add(i);
        }

        if (payloadDocuments.isEmpty()) {
            throw new RerankException("empty input documents");
        }

        String payload = Maps.of("model", options.getModelOrDefault(config.getModel()))
            .set("query", query)
            .set("documents", payloadDocuments)
            .toJSON();

        String response = httpClient.post(url, headers, payload);
        if (StringUtil.noText(response)) {
            throw new RerankException("empty response");
        }

        //{
        //  "model": "Qwen3-Reranker-4B",
        //  "usage": {
        //    "totalTokens": 0,
        //    "promptTokens": 0
        //  },
        //  "results": [
        //    {
        //      "index": 0,
        //      "document": {
        //        "text": "Use pandas: `import pandas as pd; df = pd.read_csv('data.csv')`"
        //      },
        //      "relevance_score": 0.95654296875
        //    },
        //    {
        //      "index": 3,
        //      "document": {
        //        "text": "CSV means Comma Separated Values. Python files can be opened using read() method."
        //      },
        //      "relevance_score": 0.822265625
        //    },
        //    {
        //      "index": 1,
        //      "document": {
        //        "text": "You can read CSV files with numpy.loadtxt()"
        //      },
        //      "relevance_score": 0.310791015625
        //    },
        //    {
        //      "index": 2,
        //      "document": {
        //        "text": "To write JSON files, use json.dump() in Python"
        //      },
        //      "relevance_score": 0.00009608268737792969
        //    }
        //  ]
        //}
        JSONObject jsonObject = JSON.parseObject(response);
        JSONArray results = (JSONArray) JSONPath.eval(jsonObject, config.getResultsJsonPath());

        if (results == null || results.isEmpty()) {
            String error = extractErrorMessage(jsonObject);
            String detail = "empty results";
            if (StringUtil.hasText(error)) {
                detail = detail + ", error=" + error;
            }
            detail = detail + ", response=" + truncate(response, MAX_RESPONSE_LOG_LENGTH);
            LOG.warn("Rerank response has no results. query={}, response={}", query, truncate(response, MAX_RESPONSE_LOG_LENGTH));
            throw new RerankException(detail);
        }


        for (int i = 0; i < results.size(); i++) {
            JSONObject result = results.getJSONObject(i);
            int index = result.getIntValue(config.getIndexJsonKey());
            if (index < 0 || index >= documentIndexMapping.size()) {
                continue;
            }
            int originalIndex = documentIndexMapping.get(index);
            Document document = documents.get(originalIndex);
            if (document != null) {
                document.setScore(result.getDoubleValue(config.getScoreJsonKey()));
            }
        }

        // 对 documents 排序， score 越大的越靠前
        documents.sort((d1, d2) -> Double.compare(scoreOrMin(d2), scoreOrMin(d1)));

        return documents;
    }

    private double scoreOrMin(Document document) {
        if (document == null || document.getScore() == null) {
            return Double.NEGATIVE_INFINITY;
        }
        return document.getScore();
    }

    private String extractErrorMessage(JSONObject jsonObject) {
        if (jsonObject == null || jsonObject.isEmpty()) {
            return null;
        }
        Object nestedMessage = JSONPath.eval(jsonObject, "$.error.message");
        if (nestedMessage != null && StringUtil.hasText(String.valueOf(nestedMessage))) {
            return String.valueOf(nestedMessage);
        }
        Object message = jsonObject.get("message");
        if (message != null && StringUtil.hasText(String.valueOf(message))) {
            return String.valueOf(message);
        }
        Object msg = jsonObject.get("msg");
        if (msg != null && StringUtil.hasText(String.valueOf(msg))) {
            return String.valueOf(msg);
        }
        return null;
    }

    private String truncate(String value, int maxLength) {
        if (value == null) {
            return null;
        }
        if (value.length() <= maxLength) {
            return value;
        }
        return value.substring(0, maxLength) + "...";
    }
}
