package org.jabref.logic.ai.models;

import java.net.MalformedURLException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

import org.jabref.logic.importer.FetcherException;
import org.jabref.logic.net.URLDownload;
import org.jabref.model.ai.llm.AiProvider;

import kong.unirest.core.JsonNode;
import kong.unirest.core.json.JSONArray;
import kong.unirest.core.json.JSONException;
import kong.unirest.core.json.JSONObject;
import org.jspecify.annotations.NullMarked;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/// Model provider for OpenAI-compatible APIs.
/// Fetches available models from the /v1/models endpoint.
/// Mistral provides an OpenAI-compatible API, so this works for Mistral as well.
@NullMarked
public class OpenAiCompatibleModelProvider implements AiModelProvider {

    private static final Logger LOGGER = LoggerFactory.getLogger(OpenAiCompatibleModelProvider.class);

    private static final Duration CONNECT_TIMEOUT = Duration.ofSeconds(5);

    @Override
    public List<String> fetchModels(AiProvider aiProvider, String apiBaseUrl, @Nullable String apiKey) {
        if (apiKey == null || apiKey.isBlank()) {
            LOGGER.debug("API key is not provided for {}, skipping model fetch", aiProvider.name());
            return List.of();
        }

        List<String> models = List.of();

        String modelsEndpoint = buildModelsEndpoint(apiBaseUrl);

        try {
            URLDownload urlDownload = new URLDownload(modelsEndpoint);
            urlDownload.setConnectTimeout(CONNECT_TIMEOUT);
            urlDownload.addHeader("Authorization", "Bearer " + apiKey);
            urlDownload.addHeader("accept", "application/json");

            String response = urlDownload.asString();
            models = parseModelsFromResponse(new JsonNode(response));

            LOGGER.debug("Successfully fetched {} models from {}", models.size(), aiProvider.name());
        } catch (FetcherException | MalformedURLException | JSONException e) {
            LOGGER.error("Failed to fetch models from {}", aiProvider.name(), e);
        }

        return models;
    }

    @Override
    public boolean supports(AiProvider aiProvider) {
        return aiProvider == AiProvider.OPEN_AI || aiProvider == AiProvider.MISTRAL_AI;
    }

    /// Builds the URL for the models endpoint from the given API base URL.
    ///
    /// The OpenAI API specification defines the models endpoint at /v1/models.
    /// This method handles various URL formats:
    ///
    /// - If the URL already ends with /v1, appends /models
    /// - If the URL doesn't end with /v1, appends /v1/models
    /// - Removes trailing slashes before building the path
    ///
    /// @param apiBaseUrl the base URL of the API (e.g., "https://api.openai.com" or "https://api.openai.com/v1")
    /// @return the complete URL for the models endpoint
    private String buildModelsEndpoint(String apiBaseUrl) {
        String baseUrl = apiBaseUrl.trim();
        if (baseUrl.endsWith("/")) {
            baseUrl = baseUrl.substring(0, baseUrl.length() - 1);
        }

        if (baseUrl.endsWith("/v1")) {
            return baseUrl + "/models";
        } else {
            return baseUrl + "/v1/models";
        }
    }

    private List<String> parseModelsFromResponse(JsonNode jsonNode) {
        List<String> models = new ArrayList<>();

        try {
            JSONObject jsonResponse = jsonNode.getObject();

            if (jsonResponse.has("data")) {
                JSONArray modelsArray = jsonResponse.getJSONArray("data");

                for (int i = 0; i < modelsArray.length(); i++) {
                    JSONObject modelObject = modelsArray.getJSONObject(i);
                    if (modelObject.has("id")) {
                        String modelId = modelObject.getString("id");
                        models.add(modelId);
                    }
                }
            }
        } catch (Exception e) {
            LOGGER.warn("Failed to parse models response.", e);
        }

        return models;
    }
}
