package com.yanzhu.module.ai.framework.ai.core.model.midjourney.api;

import cn.hutool.core.util.StrUtil;
import com.yanzhu.framework.common.util.json.JsonUtils;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpRequest;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.web.client.RestTemplate;
import org.springframework.util.StringUtils;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Predicate;

/**
 * Midjourney API
 *
 * @author fansili
 * @since 1.0
 */
@Slf4j
public class MidjourneyApi {

    private final RestTemplate restTemplate;
    private final String baseUrl;
    private final String apiKey;

    /**
     * 回调地址
     */
    private final String notifyUrl;

    public MidjourneyApi(String baseUrl, String apiKey, String notifyUrl) {
        this.restTemplate = new RestTemplate();
        this.baseUrl = baseUrl;
        this.apiKey = apiKey;
        this.notifyUrl = notifyUrl;
    }

    /**
     * imagine - 根据提示词提交绘画任务
     *
     * @param request 请求
     * @return 提交结果
     */
    public SubmitResponse imagine(ImagineRequest request) {
        if (StrUtil.isEmpty(request.getNotifyHook())) {
            request.setNotifyHook(notifyUrl);
        }
        String response = post("/submit/imagine", request);
        return JsonUtils.parseObject(response, SubmitResponse.class);
    }

    /**
     * action - 放大、缩小、U1、U2...
     *
     * @param request 请求
     * @return 提交结果
     */
    public SubmitResponse action(ActionRequest request) {
        if (StrUtil.isEmpty(request.getNotifyHook())) {
            request.setNotifyHook(notifyUrl);
        }
        String response = post("/submit/action", request);
        return JsonUtils.parseObject(response, SubmitResponse.class);
    }

    /**
     * 批量查询 task 任务
     *
     * @param ids 任务编号数组
     * @return task 任务
     */
    public List<Notify> getTaskList(Collection<String> ids) {
        String res = post("/task/list-by-condition", ImmutableMap.of("ids", ids));
        return JsonUtils.parseArray(res, Notify.class);
    }

    private String post(String uri, Object body) {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(apiKey);
        
        HttpEntity<String> entity = new HttpEntity<>(JsonUtils.toJsonString(body), headers);
        
        try {
            ResponseEntity<String> response = restTemplate.exchange(
                    baseUrl + uri,
                    HttpMethod.POST,
                    entity,
                    String.class
            );
            
            if (response.getStatusCode().series() != HttpStatus.Series.SUCCESSFUL) {
                log.error("[midjourney-api] 调用失败！请求地址:[{}]，请求参数:[{}]，响应数据: [{}]",
                        baseUrl + uri, body, response.getBody());
                throw new IllegalStateException("[midjourney-api] 调用失败！");
            }
            
            return response.getBody();
        } catch (Exception e) {
            log.error("[midjourney-api] 调用异常！请求地址:[{}]，请求参数:[{}]", baseUrl + uri, body, e);
            throw new RuntimeException("[midjourney-api] 调用失败！", e);
        }
    }

    // ========== record 结构 ==========

    /**
     * Imagine 请求（生成图片）
     */
    @Data
    public static final class ImagineRequest {

        /**
         * 垫图(参考图) base64 数组
         */
        private List<String> base64Array;
        /**
         * 提示词
         */
        private String prompt;
        /**
         * 通知地址
         */
        private String notifyHook;
        /**
         * 自定义参数
         */
        private String state;



        public static String buildState(Integer width, Integer height, String version, String model) {
            StringBuilder params = new StringBuilder();
            //  --ar 来设置尺寸
            params.append(String.format(" --ar %s:%s ", width, height));
            // --niji 模型
            if (ModelEnum.NIJI.getModel().equals(model)) {
                params.append(String.format(" --niji %s ", version));
            } else {
                params.append(String.format(" --v %s ", version));
            }
            return params.toString();
        }

    }

    /**
     * Action 请求
     */
    @Data
    public static final class ActionRequest {

        private String customId;
        private String taskId;
        private String notifyHook;



    }

    /**
     * 提交返回
     *
     * @param code 状态码
     * @param description 描述
     * @param properties 扩展字段
     * @param result 任务ID
     */
    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public static class SubmitResponse {
        private String code;
        private String description;
        private Map<String, Object> properties;
        private String result;
    }

    /**
     * 通知 request
     *
     * @param id job id
     * @param action 任务类型 {@link TaskActionEnum}
     * @param status 任务状态 {@link TaskStatusEnum}
     * @param prompt 提示词
     * @param promptEn 提示词-英文
     * @param description 任务描述
     * @param state 自定义参数
     * @param submitTime 提交时间
     * @param startTime 开始执行时间
     * @param finishTime 结束时间
     * @param imageUrl 图片url
     * @param progress 任务进度
     * @param failReason 失败原因
     * @param buttons 任务完成后的可执行按钮
     */
    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public static class Notify {
        private String id;
        private String action;
        private String status;
        private String prompt;
        private String promptEn;
        private String description;
        private String state;
        private Long submitTime;
        private Long startTime;
        private Long finishTime;
        private String imageUrl;
        private String progress;
        private String failReason;
        private List<Button> buttons;
    }

    /**
     * button
     *
     * @param customId MJ::JOB::upsample::1::85a4b4c1-8835-46c5-a15c-aea34fad1862 动作标识
     * @param emoji 图标 emoji
     * @param label Make Variations 文本
     * @param type 类型，系统内部使用
     * @param style 样式: 2（Primary）、3（Green）
     */
    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public static class Button {
        private String customId;
        private String emoji;
        private String label;
        private String type;
        private String style;
    }

    // ============ enums ============

    /**
     * 模型枚举
     */
    @Getter
    public enum ModelEnum {

        MIDJOURNEY("midjourney", "midjourney"),
        NIJI("niji", "niji");

        private final String model;
        private final String name;

        ModelEnum(String model, String name) {
            this.model = model;
            this.name = name;
        }

    }

    /**
     * 提交返回的状态码的枚举
     */
    @Getter
    public enum SubmitCodeEnum {

        SUBMIT_SUCCESS("1", "提交成功"),
        ALREADY_EXISTS("21", "已存在"),
        QUEUING("22", "排队中");

        public static final List<String> SUCCESS_CODES = Lists.newArrayList(
                SUBMIT_SUCCESS.code,
                ALREADY_EXISTS.code,
                QUEUING.code
        );

        private final String code;
        private final String name;

        SubmitCodeEnum(String code, String name) {
            this.code = code;
            this.name = name;
        }

    }

    /**
     * Action 枚举
     */
    @Getter
    @AllArgsConstructor
    public enum TaskActionEnum {

        /**
         * 生成图片
         */
        IMAGINE,
        /**
         * 选中放大
         */
        UPSCALE,
        /**
         * 选中其中的一张图，生成四张相似的
         */
        VARIATION,
        /**
         * 重新执行
         */
        REROLL,
        /**
         * 图转 prompt
         */
        DESCRIBE,
        /**
         * 多图混合
         */
        BLEND

    }

    /**
     * 任务状态枚举
     */
    @Getter
    public enum TaskStatusEnum {

        /**
         * 未启动
         */
        NOT_START(0),
        /**
         * 已提交
         */
        SUBMITTED(1),
        /**
         * 执行中
         */
        IN_PROGRESS(3),
        /**
         * 失败
         */
        FAILURE(4),
        /**
         * 成功
         */
        SUCCESS(5);

        private final int order;

        TaskStatusEnum(int order) {
            this.order = order;
        }

    }

}
