package com.xmzs.midjourney.controller;

import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.json.JSONUtil;
import com.xmzs.common.chat.constant.OpenAIConst;
import com.xmzs.common.core.domain.model.LoginUser;
import com.xmzs.common.core.exception.base.BaseException;
import com.xmzs.common.satoken.utils.LoginHelper;
import com.xmzs.midjourney.Constants;
import com.xmzs.midjourney.ProxyProperties;
import com.xmzs.midjourney.ReturnCode;
import com.xmzs.midjourney.dto.*;
import com.xmzs.midjourney.enums.TaskAction;
import com.xmzs.midjourney.enums.TaskStatus;
import com.xmzs.midjourney.enums.TranslateWay;
import com.xmzs.midjourney.exception.BannedPromptException;
import com.xmzs.midjourney.result.SubmitResultVO;
import com.xmzs.midjourney.service.TaskService;
import com.xmzs.midjourney.service.TaskStoreService;
import com.xmzs.midjourney.service.TranslateService;
import com.xmzs.midjourney.support.Task;
import com.xmzs.midjourney.support.TaskCondition;
import com.xmzs.midjourney.util.BannedPromptUtils;
import com.xmzs.midjourney.util.ConvertUtils;
import com.xmzs.midjourney.util.MimeTypeUtils;
import com.xmzs.midjourney.util.SnowFlake;
import com.xmzs.midjourney.util.TaskChangeParams;
import com.xmzs.system.domain.bo.ChatMessageBo;
import com.xmzs.system.service.IChatMessageService;
import com.xmzs.system.service.IChatService;
import com.xmzs.system.service.ISseService;
import eu.maxschuster.dataurl.DataUrl;
import eu.maxschuster.dataurl.DataUrlSerializer;
import eu.maxschuster.dataurl.IDataUrlSerializer;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.OkHttpClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.net.MalformedURLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import okhttp3.*;

@Api(tags = "任务提交")
@RestController
@RequestMapping("/mj/submit")
@RequiredArgsConstructor
@Slf4j
public class SubmitController {
    private final TranslateService translateService;
    private final ProxyProperties properties;
    private final TaskService taskService;
    private final TaskStoreService taskStoreService;
    @Value("${chat.apiKey}")
    private String apiKey;
    @Value("${chat.apiHost}")
    private String apiHost;
    @Autowired
    private IChatService chatService;
    @Autowired
    private IChatMessageService chatMessageService;
    @Autowired
    private ISseService sseService;

    @ApiOperation(value = "提交Imagine任务")
    @PostMapping("/imagine")
    public SubmitResultVO imagine(@RequestBody SubmitImagineDTO imagineDTO) {
        String prompt = imagineDTO.getPrompt();
        if (CharSequenceUtil.isBlank(prompt)) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "prompt不能为空");
        }
        prompt = prompt.trim();
        Task task = newTask(imagineDTO);
        task.setAction(TaskAction.IMAGINE);
        task.setPrompt(prompt);
        String promptEn = translatePrompt(prompt);
        try {
            BannedPromptUtils.checkBanned(promptEn);
        } catch (BannedPromptException e) {
            return SubmitResultVO.fail(ReturnCode.BANNED_PROMPT, "可能包含敏感词")
                .setProperty("promptEn", promptEn).setProperty("bannedWord", e.getMessage());
        }
        List<String> base64Array = Optional.ofNullable(imagineDTO.getBase64Array()).orElse(new ArrayList<>());
        if (CharSequenceUtil.isNotBlank(imagineDTO.getBase64())) {
            base64Array.add(imagineDTO.getBase64());
        }
        List<DataUrl> dataUrls;
        try {
            dataUrls = ConvertUtils.convertBase64Array(base64Array);
        } catch (MalformedURLException e) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64格式错误");
        }
        task.setPromptEn(promptEn);
        task.setDescription("/imagine " + prompt);
        return this.taskService.submitImagine(task, dataUrls);
    }

    @ApiOperation(value = "绘图变化-simple")
    @PostMapping("/simple-change")
    public SubmitResultVO simpleChange(@RequestBody SubmitSimpleChangeDTO simpleChangeDTO) {
        TaskChangeParams changeParams = ConvertUtils.convertChangeParams(simpleChangeDTO.getContent());
        if (changeParams == null) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "content参数错误");
        }
        SubmitChangeDTO changeDTO = new SubmitChangeDTO();
        changeDTO.setAction(changeParams.getAction());
        changeDTO.setTaskId(changeParams.getId());
        changeDTO.setIndex(changeParams.getIndex());
        changeDTO.setState(simpleChangeDTO.getState());
        changeDTO.setNotifyHook(simpleChangeDTO.getNotifyHook());
        return change(changeDTO);
    }

    @ApiOperation(value = "绘图变化")
    @PostMapping("/change")
    public SubmitResultVO change(@RequestBody SubmitChangeDTO changeDTO) {
        if (CharSequenceUtil.isBlank(changeDTO.getTaskId())) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "taskId不能为空");
        }
        if (!Set.of(TaskAction.UPSCALE, TaskAction.VARIATION, TaskAction.REROLL).contains(changeDTO.getAction())) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "action参数错误");
        }
        String description = "/up " + changeDTO.getTaskId();
        if (TaskAction.REROLL.equals(changeDTO.getAction())) {
            description += " R";
        } else {
            description += " " + changeDTO.getAction().name().charAt(0) + changeDTO.getIndex();
        }
        if (TaskAction.UPSCALE.equals(changeDTO.getAction())) {
            TaskCondition condition = new TaskCondition().setDescription(description);
            Task existTask = this.taskStoreService.findOne(condition);
            if (existTask != null) {
                return SubmitResultVO.of(ReturnCode.EXISTED, "任务已存在", existTask.getId())
                    .setProperty("status", existTask.getStatus())
                    .setProperty("imageUrl", existTask.getImageUrl());
            }
        }
        Task targetTask = this.taskStoreService.get(changeDTO.getTaskId());
        if (targetTask == null) {
            return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "关联任务不存在或已失效");
        }
        if (!TaskStatus.SUCCESS.equals(targetTask.getStatus())) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "关联任务状态错误");
        }
        if (!Set.of(TaskAction.IMAGINE, TaskAction.VARIATION, TaskAction.REROLL, TaskAction.BLEND).contains(targetTask.getAction())) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "关联任务不允许执行变化");
        }
        Task task = newTask(changeDTO);
        task.setAction(changeDTO.getAction());
        task.setPrompt(targetTask.getPrompt());
        task.setPromptEn(targetTask.getPromptEn());
        task.setProperty(Constants.TASK_PROPERTY_FINAL_PROMPT, targetTask.getProperty(Constants.TASK_PROPERTY_FINAL_PROMPT));
        task.setProperty(Constants.TASK_PROPERTY_PROGRESS_MESSAGE_ID, targetTask.getProperty(Constants.TASK_PROPERTY_MESSAGE_ID));
        task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, targetTask.getProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID));
        task.setDescription(description);
        int messageFlags = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_FLAGS);
        String messageId = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_MESSAGE_ID);
        String messageHash = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_MESSAGE_HASH);
        if (TaskAction.UPSCALE.equals(changeDTO.getAction())) {
            return this.taskService.submitUpscale(task, messageId, messageHash, changeDTO.getIndex(), messageFlags);
            } else if (TaskAction.VARIATION.equals(changeDTO.getAction())) {
            return this.taskService.submitVariation(task, messageId, messageHash, changeDTO.getIndex(), messageFlags);
        } else {
            return this.taskService.submitReroll(task, messageId, messageHash, messageFlags);
        }
    }

    @ApiOperation(value = "绘图变化")
    @PostMapping("/action")
    public String action(@RequestBody SubmitActionDTO changeDTO) {
        // 查询是否是付费用户
        sseService.checkUserGrade();
        // 扣除接口费用
        if ("upsample".equals(getAction(changeDTO.getCustomId()))) {
            mjTaskDeduct("放大", OpenAIConst.MJ_COST_TYPE2);
        } else {
            // Inpaint: 局部重绘
            // reroll 重绘
            // upsample 放大
            // zoom 变焦
            // upscale 高清放大
            // variation 变化
            if (!"Inpaint".equals(getAction(changeDTO.getCustomId()))) {
                mjTaskDeduct("变化", OpenAIConst.MJ_COST_TYPE1);
            }
        }
        OkHttpClient client = new OkHttpClient.Builder()
            .connectTimeout(30, TimeUnit.SECONDS) // 连接超时时间
            .writeTimeout(30, TimeUnit.SECONDS)   // 写入超时时间
            .readTimeout(30, TimeUnit.SECONDS)    // 读取超时时间
            .build();

        String jsonStr = JSONUtil.toJsonStr(changeDTO);

        MediaType mediaType = MediaType.parse("application/json");
        okhttp3.RequestBody body = okhttp3.RequestBody.create(jsonStr, mediaType);
        Request request = new Request.Builder()
            .url(apiHost + "mj/submit/action")
            .method("POST", body)
            .header("mj-api-secret", apiKey) // 设置Authorization header
            .build();
        try {
            Response response = client.newCall(request).execute();
            return response.body().string();
        } catch (IOException e) {
            log.error("绘图变化失败:{}", e.getMessage());
        }
        return null;
    }

    public String getAction(String customId) {
        // 检查 customId 是否为空
        if(customId == null || customId.isEmpty()) {
            return null;
        }
        // 使用 "::" 分割字符串
        String[] parts = customId.split("::");
        // "MJ", "Inpaint", "1", "4fca7c14-181c-4...", "SOLO"
        if(customId.endsWith("SOLO")) {
            return parts[1];
        }
        // 返回 "upsample" 值，假设它总是在第三个位置
        return parts[2];
    }

    public void mjTaskDeduct(String prompt, double cost) {
        //扣除费用
        chatService.deductUserBalance(getUserId(), cost);
        // 保存消息记录
        ChatMessageBo chatMessageBo = new ChatMessageBo();
        chatMessageBo.setUserId(getUserId());
        chatMessageBo.setModelName("mj");
        chatMessageBo.setContent(prompt);
        chatMessageBo.setDeductCost(cost);
        chatMessageBo.setTotalTokens(0);
        chatMessageService.insertByBo(chatMessageBo);
    }

    /**
     * 获取用户Id
     *
     * @return
     */
    public Long getUserId() {
        LoginUser loginUser = LoginHelper.getLoginUser();
        if (loginUser == null) {
            throw new BaseException("用户未登录！");
        }
        return loginUser.getUserId();
    }

    @ApiOperation(value = "提交Describe任务")
    @PostMapping("/describe")
    public SubmitResultVO describe(@RequestBody SubmitDescribeDTO describeDTO) {
        if (CharSequenceUtil.isBlank(describeDTO.getBase64())) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64不能为空");
        }
        IDataUrlSerializer serializer = new DataUrlSerializer();
        DataUrl dataUrl;
        try {
            dataUrl = serializer.unserialize(describeDTO.getBase64());
        } catch (MalformedURLException e) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64格式错误");
        }
        Task task = newTask(describeDTO);
        task.setAction(TaskAction.DESCRIBE);
        String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
        task.setDescription("/describe " + taskFileName);
        return this.taskService.submitDescribe(task, dataUrl);
    }

    @ApiOperation(value = "提交Blend任务")
    @PostMapping("/blend")
    public SubmitResultVO blend(@RequestBody SubmitBlendDTO blendDTO) {
        List<String> base64Array = blendDTO.getBase64Array();
        if (base64Array == null || base64Array.size() < 2 || base64Array.size() > 5) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64List参数错误");
        }
        if (blendDTO.getDimensions() == null) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "dimensions参数错误");
        }
        IDataUrlSerializer serializer = new DataUrlSerializer();
        List<DataUrl> dataUrlList = new ArrayList<>();
        try {
            for (String base64 : base64Array) {
                DataUrl dataUrl = serializer.unserialize(base64);
                dataUrlList.add(dataUrl);
            }
        } catch (MalformedURLException e) {
            return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64格式错误");
        }
        Task task = newTask(blendDTO);
        task.setAction(TaskAction.BLEND);
        task.setDescription("/blend " + task.getId() + " " + dataUrlList.size());
        return this.taskService.submitBlend(task, dataUrlList, blendDTO.getDimensions());
    }

    private Task newTask(BaseSubmitDTO base) {
        Task task = new Task();
        task.setId(System.currentTimeMillis() + RandomUtil.randomNumbers(3));
        task.setSubmitTime(System.currentTimeMillis());
        task.setState(base.getState());
        String notifyHook = CharSequenceUtil.isBlank(base.getNotifyHook()) ? this.properties.getNotifyHook() : base.getNotifyHook();
        task.setProperty(Constants.TASK_PROPERTY_NOTIFY_HOOK, notifyHook);
        task.setProperty(Constants.TASK_PROPERTY_NONCE, SnowFlake.INSTANCE.nextId());
        return task;
    }

    private String translatePrompt(String prompt) {
        if (TranslateWay.NULL.equals(this.properties.getTranslateWay()) || CharSequenceUtil.isBlank(prompt)) {
            return prompt;
        }
        List<String> imageUrls = new ArrayList<>();
        Matcher imageMatcher = Pattern.compile("https?://[a-z0-9-_:@&?=+,.!/~*'%$]+\\x20+", Pattern.CASE_INSENSITIVE).matcher(prompt);
        while (imageMatcher.find()) {
            imageUrls.add(imageMatcher.group(0));
        }
        String paramStr = "";
        Matcher paramMatcher = Pattern.compile("\\x20+-{1,2}[a-z]+.*$", Pattern.CASE_INSENSITIVE).matcher(prompt);
        if (paramMatcher.find()) {
            paramStr = paramMatcher.group(0);
        }
        String imageStr = CharSequenceUtil.join("", imageUrls);
        String text = prompt.substring(imageStr.length(), prompt.length() - paramStr.length());
        if (CharSequenceUtil.isNotBlank(text)) {
            text = this.translateService.translateToEnglish(text).trim();
        }
        return imageStr + text + paramStr;
    }

}