package com.wmyun.module.system.controller.admin.forward;

import com.volcengine.ark.runtime.exception.ArkHttpException;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionContentPart;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import com.wmyun.framework.common.pojo.CommonResult;
import com.wmyun.module.system.controller.admin.forward.vo.OCRReqVO;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.validation.Valid;
import jodd.util.StringUtil;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

import static com.wmyun.framework.common.pojo.CommonResult.error;
import static com.wmyun.framework.common.pojo.CommonResult.success;

/**
 * @author zhouc
 * @date 2025/3/3 21:01
 * @className OCRController
 * @description
 */
@Tag(name = "管理后台 -OCR识别")
@RestController
@RequestMapping("/system/ocr")
@Validated
@Transactional
public class OCRController {
    private static final String API_KEY = "83e8d382-1f6f-47f3-aa8e-1ea1f135ca2b";
    private static final String DEEPSEEK_OCR_MODEL = "ep-20250224214646-xxb4g";
    private static final String baseUrl = "https://ark.cn-beijing.volces.com/api/v3/chat/completions";

    @PostMapping(value = "/OCRinterface", consumes = "multipart/form-data")
    @Operation(summary = "OCR识别")
    @PreAuthorize("@ss.hasPermission('system:forward:OCRinterface')")
    public CommonResult<String> OCRinterface(@RequestParam("file") MultipartFile file,
                                             @RequestParam("mag") String mag) {
        try {
            String imageurl = "";
            // 校验文件非空
            if (file.isEmpty()) {
                throw new IllegalArgumentException("文件为空");
            }
            String contentType = file.getContentType();
            // 读取字节流并编码
            try (InputStream inputStream = file.getInputStream()) {
                byte[] bytes = inputStream.readAllBytes();
                imageurl = Base64.getEncoder().encodeToString(bytes);
            }
            imageurl = "data:" + contentType + ";base64," + imageurl;
            Object getmsg = getmsg(mag, imageurl);
            return success(getmsg.toString());
        } catch (Exception exception) {
            String message = exception.getMessage();
            return error(411, message);
        }
    }

    public Object getmsg(String mag, String imageurl) {
        AtomicReference<Object> msg = new AtomicReference<>("");
        try {
            ArkService service = ArkService.builder().apiKey(API_KEY).baseUrl(baseUrl).build();
//        ArkService service = new ArkService(ARK_API_KEY);
            System.out.println("----- image input -----");
            final List<ChatMessage> messages = new ArrayList<>();
            final List<ChatCompletionContentPart> multiParts = new ArrayList<>();
            multiParts.add(ChatCompletionContentPart.builder().type("text").text(mag).build());
            multiParts.add(ChatCompletionContentPart.builder().type("image_url").imageUrl(
                    new ChatCompletionContentPart.ChatCompletionContentPartImageURL(imageurl)
            ).build());
            final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER)
                    .multiContent(multiParts).build();
            messages.add(userMessage);

            ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model(DEEPSEEK_OCR_MODEL)
                    .messages(messages)
                    .build();

            service.createChatCompletion(chatCompletionRequest).getChoices().forEach(
                    choice ->
                            msg.set(choice.getMessage().getContent()));

            service.shutdownExecutor();
        } catch (ArkHttpException e) {
            System.out.print(e.toString());
        }
        return msg;
    }


}
