package com.pharmacopoeia.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.pharmacopoeia.config.QwenProperties;
import com.pharmacopoeia.exception.DashScopeRateLimitExceededException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.util.UriBuilder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

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

@Service
public class LLMService {

    private static final Logger log = LoggerFactory.getLogger(LLMService.class);

    private final WebClient chatClient;
    private final WebClient embedClient;
    private final WebClient visionClient;
    private final QwenProperties props;
    private final DashScopeRateLimiter rateLimiter;
    private final ObjectMapper mapper = new ObjectMapper();

    public LLMService(QwenProperties props, DashScopeRateLimiter rateLimiter) {
        this.props = props;
        this.rateLimiter = rateLimiter;
        this.chatClient = WebClient.builder()
                .baseUrl(props.getBaseUrl())
                .defaultHeader("Authorization", "Bearer " + props.getApiKey())
                .build();
        this.embedClient = WebClient.builder()
                .defaultHeader("Authorization", "Bearer " + props.getApiKey())
                .build();
        // 原生多模态 API（非兼容模式），endpoint 与 chat/embed 不同
        this.visionClient = WebClient.builder()
                .baseUrl("https://dashscope.aliyuncs.com/api/v1")
                .defaultHeader("Authorization", "Bearer " + props.getApiKey())
                .build();
    }

    /** 检查百炼限流，超限时抛专用异常 */
    private void checkRateLimit() {
        if (!rateLimiter.tryAcquire()) {
            throw new DashScopeRateLimitExceededException("百炼 API 调用频率超限，请稍后再试");
        }
    }

    // ============================================================
    // 纯文本对话
    // ============================================================

    public Flux<String> chatStream(List<Map<String, String>> messages) {
        return chatStream(messages, false);
    }

    public Flux<String> chatStream(List<Map<String, String>> messages, boolean enableSearch) {
        checkRateLimit();
        long t0 = System.currentTimeMillis();
        Map<String, Object> body = new java.util.HashMap<>(Map.of(
                "model", props.getModel(),
                "messages", messages,
                "max_tokens", props.getMaxTokens(),
                "temperature", props.getTemperature(),
                "stream", true
        ));
        if (enableSearch) {
            body.put("enable_search", true);
        }
        log.info("[chatStream] POST /chat/completions, model={}, msgs={}, bodyBuilt={}ms",
                props.getModel(), messages.size(), System.currentTimeMillis() - t0);

        final long t1 = System.currentTimeMillis();
        return chatClient.post()
                .uri("/chat/completions")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToFlux(String.class)
                .doFirst(() -> {
                    long elapsed = System.currentTimeMillis() - t1;
                    log.info("[chatStream] 百炼已响应, latency={}ms (从POST到收到第一个字节)", elapsed);
                })
                .filter(data -> !"[DONE]".equals(data.trim()))
                .map(this::extractDeltaContent)
                .doOnComplete(() -> {
                    long total = System.currentTimeMillis() - t0;
                    log.info("[chatStream] 流结束, total={}ms", total);
                })
                .doOnError(e -> {
                    long total = System.currentTimeMillis() - t0;
                    log.error("[chatStream] 流出错, total={}ms, error={}", total, e.getMessage());
                });
    }

    public String chat(List<Map<String, String>> messages) {
        return chat(messages, false);
    }

    public String chat(List<Map<String, String>> messages, boolean enableSearch) {
        checkRateLimit();
        long t0 = System.currentTimeMillis();
        Map<String, Object> body = new java.util.HashMap<>(Map.of(
                "model", props.getModel(),
                "messages", messages,
                "max_tokens", props.getMaxTokens(),
                "temperature", props.getTemperature()
        ));
        if (enableSearch) {
            body.put("enable_search", true);
        }
        long t1 = System.currentTimeMillis();
        String response = chatClient.post()
                .uri("/chat/completions")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToMono(String.class)
                .block();
        long t2 = System.currentTimeMillis();
        log.info("[chat] POST /chat/completions completed, latency={}ms, bodyBuilt={}ms, total={}ms",
                t2 - t1, t1 - t0, t2 - t0);

        try {
            JsonNode node = mapper.readTree(response);
            return node.get("choices").get(0).get("message").get("content").asText();
        } catch (Exception e) {
            log.error("chat() 解析 LLM 响应失败，response={}", response, e);
            return "";
        }
    }

    // ============================================================
    // 图片分析 + OCR（原生多模态 API）
    // ============================================================

    /**
     * 调用 Qwen 多模态模型分析图片，提取文字（OCR）
     */
    public String analyzeImage(String imageBase64, String mimeType, String prompt) {
        checkRateLimit();
        if (prompt == null || prompt.isBlank()) {
            prompt = """
                请仔细分析这张图片，完成以下任务：
                1. 描述图片内容（药品包装、说明书、处方单、症状照片等）
                2. 提取图片中所有可见的文字（OCR），特别是药品名称、成分、用法用量等关键信息
                3. 如果图片是药品包装/说明书，提取：药品通用名、规格、生产企业、批准文号
                请按以下格式输出：
                【图片描述】简要描述。
                【OCR 提取文字】逐条列出提取到的文字内容。
                """;
        }

        // 原生 DashScope 多模态格式
        List<Map<String, Object>> contentParts = List.of(
                Map.of("image", "data:" + mimeType + ";base64," + imageBase64),
                Map.of("text", prompt)
        );

        Map<String, Object> userMessage = Map.of("role", "user", "content", contentParts);

        Map<String, Object> body = new java.util.HashMap<>();
        body.put("model", props.getVlModel());
        body.put("input", Map.of("messages", List.of(userMessage)));
        body.put("parameters", Map.of("max_tokens", 2048));

        String response = visionClient.post()
                .uri(props.getMultimodalUrl(), UriBuilder::build)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToMono(String.class)
                .block();

        try {
            JsonNode node = mapper.readTree(response);
            // 原生 API 响应：output.choices[0].message.content[0].text
            JsonNode content = node.get("output").get("choices").get(0)
                    .get("message").get("content");
            if (content.isArray() && !content.isEmpty()) {
                return content.get(0).get("text").asText();
            }
            return content.asText();
        } catch (Exception e) {
            log.error("analyzeImage() 解析响应失败，response={}", response, e);
            return "";
        }
    }

    /**
     * 流式版本：Qwen 多模态分析图片
     */
    public Flux<String> analyzeImageStream(String imageBase64, String mimeType, String prompt) {
        checkRateLimit();
        if (prompt == null || prompt.isBlank()) {
            prompt = "请分析这张图片，提取所有文字信息（OCR），特别是药品名称、成分、用法用量等。简要输出。";
        }

        List<Map<String, Object>> contentParts = List.of(
                Map.of("image", "data:" + mimeType + ";base64," + imageBase64),
                Map.of("text", prompt)
        );

        Map<String, Object> userMessage = Map.of("role", "user", "content", contentParts);

        Map<String, Object> body = new java.util.HashMap<>();
        body.put("model", props.getVlModel());
        body.put("input", Map.of("messages", List.of(userMessage)));
        body.put("parameters", Map.of("max_tokens", 2048, "stream", true, "incremental_output", true));

        return visionClient.post()
                .uri(props.getMultimodalUrl(), UriBuilder::build)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToFlux(String.class)
                .filter(data -> !"[DONE]".equals(data.trim()))
                .map(this::extractVisionDelta);
    }

    // ============================================================
    // 视频分析 + OCR（原生多模态 API）
    // ============================================================

    public String analyzeVideo(String videoBase64, String mimeType, String prompt) {
        checkRateLimit();
        if (prompt == null || prompt.isBlank()) {
            prompt = """
                请仔细分析这段视频，完成以下任务：
                1. 描述视频内容（药品展示、用药指导、症状表现等）
                2. 提取视频中所有可见的文字（OCR），特别是药品名称、成分、用法用量等
                3. 如果视频中有药品包装/说明书，提取全部文字信息
                请按格式输出：【视频描述】→【OCR 提取文字】→【关键信息总结】
                """;
        }

        List<Map<String, Object>> contentParts = List.of(
                Map.of("video", "data:" + mimeType + ";base64," + videoBase64),
                Map.of("text", prompt)
        );

        Map<String, Object> userMessage = Map.of("role", "user", "content", contentParts);

        Map<String, Object> body = new java.util.HashMap<>();
        body.put("model", props.getVlModel());
        body.put("input", Map.of("messages", List.of(userMessage)));
        body.put("parameters", Map.of("max_tokens", 2048));

        String response = visionClient.post()
                .uri(props.getMultimodalUrl(), UriBuilder::build)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToMono(String.class)
                .block();

        try {
            JsonNode node = mapper.readTree(response);
            JsonNode content = node.get("output").get("choices").get(0)
                    .get("message").get("content");
            if (content.isArray() && !content.isEmpty()) {
                return content.get(0).get("text").asText();
            }
            return content.asText();
        } catch (Exception e) {
            log.error("analyzeVideo() 解析响应失败，response={}", response, e);
            return "";
        }
    }

    public Flux<String> analyzeVideoStream(String videoBase64, String mimeType, String prompt) {
        checkRateLimit();
        if (prompt == null || prompt.isBlank()) {
            prompt = "请分析这段视频，提取关键帧中的文字信息（OCR），特别是药品名称、成分、用法用量等。简要输出。";
        }

        List<Map<String, Object>> contentParts = List.of(
                Map.of("video", "data:" + mimeType + ";base64," + videoBase64),
                Map.of("text", prompt)
        );

        Map<String, Object> userMessage = Map.of("role", "user", "content", contentParts);

        Map<String, Object> body = new java.util.HashMap<>();
        body.put("model", props.getVlModel());
        body.put("input", Map.of("messages", List.of(userMessage)));
        body.put("parameters", Map.of("max_tokens", 2048, "stream", true, "incremental_output", true));

        return visionClient.post()
                .uri(props.getMultimodalUrl(), UriBuilder::build)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToFlux(String.class)
                .filter(data -> !"[DONE]".equals(data.trim()))
                .map(this::extractVisionDelta);
    }

    // ============================================================
    // 统一多模态分析入口
    // ============================================================

    public String analyzeMedia(String mediaBase64, String mediaType, String mimeType, String prompt) {
        if ("video".equals(mediaType)) {
            return analyzeVideo(mediaBase64, mimeType != null ? mimeType : "video/mp4", prompt);
        }
        return analyzeImage(mediaBase64, mimeType != null ? mimeType : "image/jpeg", prompt);
    }

    public Flux<String> analyzeMediaStream(String mediaBase64, String mediaType, String mimeType, String prompt) {
        if ("video".equals(mediaType)) {
            return analyzeVideoStream(mediaBase64, mimeType != null ? mimeType : "video/mp4", prompt);
        }
        return analyzeImageStream(mediaBase64, mimeType != null ? mimeType : "image/jpeg", prompt);
    }

    // ============================================================
    // Embedding（同步版：供非 reactor 线程调用）
    // ============================================================

    public List<Float> embed(String text) {
        List<List<Float>> result = embedBatch(List.of(text));
        if (result.isEmpty()) {
            return List.of();
        }
        return result.getFirst();
    }

    private List<List<Float>> embedBatch(List<String> texts) {
        checkRateLimit();
        long t0 = System.currentTimeMillis();
        int maxRetries = 3;
        for (int attempt = 1; attempt <= maxRetries; attempt++) {
            try {
                long t1 = System.currentTimeMillis();
                String response = embedClient.post()
                        .uri(props.getEmbeddingUrl())
                        .contentType(MediaType.APPLICATION_JSON)
                        .bodyValue(Map.of(
                                "model", props.getEmbeddingModel(),
                                "input", Map.of("texts", texts),
                                "parameters", Map.of("text_type", "query")
                        ))
                        .retrieve()
                        .bodyToMono(String.class)
                        .block();
                long t2 = System.currentTimeMillis();
                log.info("[embed] POST embedding latency={}ms, texts={}", t2 - t1, texts.size());

                JsonNode node = mapper.readTree(response);
                var embeddings = node.get("output").get("embeddings");
                var result = Flux.fromIterable(embeddings)
                        .map(e -> Flux.fromIterable(e.get("embedding"))
                                .map(JsonNode::asDouble)
                                .map(Double::floatValue)
                                .collectList()
                                .block())
                        .collectList()
                        .block();
                log.info("[embed] total={}ms (含解析), attempt={}", System.currentTimeMillis() - t0, attempt);
                return result;
            } catch (Exception e) {
                if (attempt < maxRetries) {
                    log.warn("Embedding API失败，第 {} 次重试...", attempt);
                    try { Thread.sleep(attempt * 1000L); } catch (InterruptedException ignored) {}
                } else {
                    log.error("Embedding API 重试 {} 次后仍失败，降级返回空", maxRetries, e);
                    return List.of();
                }
            }
        }
        return List.of();
    }

    // ============================================================
    // Embedding（响应式版：供 reactor 线程 / SSE 端点调用）
    // ============================================================

    public Mono<List<Float>> embedReactive(String text) {
        return embedReactiveBatch(List.of(text))
                .map(result -> result.isEmpty() ? List.of() : result.getFirst());
    }

    private Mono<List<List<Float>>> embedReactiveBatch(List<String> texts) {
        checkRateLimit();
        return embedClient.post()
                .uri(props.getEmbeddingUrl())
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(Map.of(
                        "model", props.getEmbeddingModel(),
                        "input", Map.of("texts", texts),
                        "parameters", Map.of("text_type", "query")
                ))
                .retrieve()
                .bodyToMono(String.class)
                .flatMap(response -> {
                    try {
                        JsonNode node = mapper.readTree(response);
                        var embeddings = node.get("output").get("embeddings");
                        List<List<Float>> result = new ArrayList<>();
                        for (JsonNode e : embeddings) {
                            List<Float> vec = new ArrayList<>();
                            for (JsonNode v : e.get("embedding")) {
                                vec.add((float) v.asDouble());
                            }
                            result.add(vec);
                        }
                        return Mono.just(result);
                    } catch (Exception ex) {
                        return Mono.error(ex);
                    }
                })
                .retryWhen(reactor.util.retry.Retry.backoff(3, java.time.Duration.ofSeconds(1))
                        .maxBackoff(java.time.Duration.ofSeconds(4))
                        .doBeforeRetry(rs -> log.warn("Embedding API 失败，第 {} 次重试...", rs.totalRetries() + 1))
                        .onRetryExhaustedThrow((rb, rs) -> {
                            log.error("Embedding API 重试 3 次后仍失败，降级返回空");
                            return new RuntimeException("Embedding exhausted");
                        }))
                .onErrorReturn(List.of());
    }

    // ============================================================
    // Helpers
    // ============================================================

    private String extractDeltaContent(String chunk) {
        try {
            String json = chunk.trim();
            if (json.startsWith("data:")) {
                json = json.substring(5).trim();
            }
            if (json.isEmpty() || "[DONE]".equals(json)) {
                return "";
            }
            JsonNode node = mapper.readTree(json);
            var choices = node.get("choices");
            if (choices != null && !choices.isEmpty()) {
                var delta = choices.get(0).get("delta");
                if (delta != null) {
                    var content = delta.get("content");
                    return content != null ? content.asText() : "";
                }
            }
        } catch (Exception e) { /* ignore parse errors on partial chunks */ }
        return "";
    }

    /** 原生多模态 API 流式响应解析：output.choices[0].message.content[0].text */
    private String extractVisionDelta(String chunk) {
        try {
            String json = chunk.trim();
            if (json.startsWith("data:")) {
                json = json.substring(5).trim();
            }
            if (json.isEmpty() || "[DONE]".equals(json)) {
                return "";
            }
            JsonNode node = mapper.readTree(json);
            JsonNode output = node.get("output");
            if (output != null) {
                var choices = output.get("choices");
                if (choices != null && !choices.isEmpty()) {
                    var msgContent = choices.get(0).get("message").get("content");
                    if (msgContent != null && msgContent.isArray() && !msgContent.isEmpty()) {
                        var text = msgContent.get(0).get("text");
                        return text != null ? text.asText() : "";
                    }
                }
            }
        } catch (Exception e) { /* ignore parse errors on partial chunks */ }
        return "";
    }
}
