package com.smyhub.infogenie.infogeniebackendjava.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.smyhub.infogenie.infogeniebackendjava.entity.AIConfig;
import com.smyhub.infogenie.infogeniebackendjava.entity.SiteAIRuntime;
import com.smyhub.infogenie.infogeniebackendjava.repository.AIConfigRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.*;

/**
 * Proxies chat and AI-tool requests to DeepSeek or Kimi endpoints.
 * Endpoint resolution mirrors the Go ai.go service:
 *   - DeepSeek: use SiteAIRuntime if api_base+api_key are set, else ai_configs table.
 *   - Kimi: always from ai_configs table (provider=kimi).
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class AIModelService {

    public static final String PROVIDER_DEEPSEEK = "deepseek";
    public static final String PROVIDER_KIMI = "kimi";
    public static final String DEFAULT_MODEL_DEEPSEEK = "deepseek-chat";

    private static final Set<String> ALLOWED_DEEPSEEK = Set.of(
            "deepseek-chat", "deepseek-reasoner");
    private static final Set<String> ALLOWED_KIMI = Set.of(
            "kimi-k2-0905-preview", "kimi-k2-0711-preview");

    private final AIConfigRepository aiConfigRepo;
    private final AIRuntimeService aiRuntimeService;
    @Qualifier("aiRestTemplate")
    private final RestTemplate aiRestTemplate;
    private final ObjectMapper objectMapper;

    public Map<String, List<String>> getAllowedModels() {
        return Map.of(
                PROVIDER_DEEPSEEK, new ArrayList<>(ALLOWED_DEEPSEEK),
                PROVIDER_KIMI, new ArrayList<>(ALLOWED_KIMI)
        );
    }

    /**
     * Calls the AI provider and returns the assistant message content.
     * Retries up to {@code maxRetries} times for DeepSeek; once for Kimi.
     */
    public String callAI(String provider, String model,
                         List<Map<String, String>> messages) throws Exception {
        String resolvedProvider = provider == null ? PROVIDER_DEEPSEEK : provider;
        String resolvedModel = resolveModel(resolvedProvider, model);

        EndpointConfig ep = resolveEndpoint(resolvedProvider);
        int maxAttempts = PROVIDER_DEEPSEEK.equals(resolvedProvider) ? 3 : 1;

        Exception lastEx = null;
        for (int i = 0; i < maxAttempts; i++) {
            try {
                return doChat(ep, resolvedModel, messages);
            } catch (Exception e) {
                lastEx = e;
                log.warn("AI call attempt {}/{} failed: {}", i + 1, maxAttempts, e.getMessage());
                if (i < maxAttempts - 1) Thread.sleep(1000L * (i + 1));
            }
        }
        throw lastEx != null ? lastEx : new RuntimeException("AI call failed");
    }

    /**
     * Returns the raw SSE stream URL and headers needed for streaming chat.
     * The controller is responsible for proxying the response stream.
     */
    public StreamConfig buildStreamConfig(String provider, String model,
                                          List<Map<String, String>> messages) {
        String resolvedProvider = provider == null ? PROVIDER_DEEPSEEK : provider;
        String resolvedModel = resolveModel(resolvedProvider, model);
        EndpointConfig ep = resolveEndpoint(resolvedProvider);

        Map<String, Object> body = buildRequestBody(resolvedModel, messages, true);
        return new StreamConfig(ep.url(), ep.apiKey(), body);
    }

    private String doChat(EndpointConfig ep, String model,
                          List<Map<String, String>> messages) throws Exception {
        Map<String, Object> body = buildRequestBody(model, messages, false);

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(ep.apiKey());

        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(body, headers);
        ResponseEntity<String> response =
                aiRestTemplate.postForEntity(ep.url(), entity, String.class);

        if (!response.getStatusCode().is2xxSuccessful() || response.getBody() == null) {
            throw new RuntimeException("AI API returned " + response.getStatusCode());
        }

        JsonNode root = objectMapper.readTree(response.getBody());
        return root.path("choices").path(0).path("message").path("content").asText();
    }

    private Map<String, Object> buildRequestBody(String model,
                                                  List<Map<String, String>> messages,
                                                  boolean stream) {
        Map<String, Object> body = new LinkedHashMap<>();
        body.put("model", model);
        body.put("messages", messages);
        if (stream) body.put("stream", true);
        return body;
    }

    private String resolveModel(String provider, String model) {
        if (model == null || model.isBlank()) {
            return DEFAULT_MODEL_DEEPSEEK;
        }
        Set<String> allowed = PROVIDER_KIMI.equals(provider) ? ALLOWED_KIMI : ALLOWED_DEEPSEEK;
        return allowed.contains(model) ? model : DEFAULT_MODEL_DEEPSEEK;
    }

    private EndpointConfig resolveEndpoint(String provider) {
        if (PROVIDER_DEEPSEEK.equals(provider)) {
            SiteAIRuntime runtime = aiRuntimeService.getRuntime().orElse(null);
            if (runtime != null
                    && runtime.getApiBase() != null && !runtime.getApiBase().isBlank()
                    && runtime.getApiKey() != null && !runtime.getApiKey().isBlank()) {
                return new EndpointConfig(
                        runtime.getApiBase().stripTrailing() + "/chat/completions",
                        runtime.getApiKey());
            }
        }

        String prov = PROVIDER_KIMI.equals(provider) ? PROVIDER_KIMI : PROVIDER_DEEPSEEK;
        AIConfig cfg = aiConfigRepo.findByProviderAndIsEnabledTrue(prov)
                .orElseThrow(() -> new RuntimeException(
                        "No enabled AI config for provider: " + prov));

        String base = cfg.getApiBase().stripTrailing();
        String url = PROVIDER_KIMI.equals(prov)
                ? base + "/v1/chat/completions"
                : base + "/chat/completions";

        return new EndpointConfig(url, cfg.getApiKey());
    }

    public record EndpointConfig(String url, String apiKey) {}

    public record StreamConfig(String url, String apiKey, Map<String, Object> body) {}
}
