package com.iailab.framework.ai.core.model.midjourney.api;

import cn.hutool.core.util.StrUtil;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.iailab.framework.common.util.json.JsonUtils;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpRequest;
import org.springframework.http.HttpStatusCode;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

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 Predicate<HttpStatusCode> STATUS_PREDICATE = status -> !status.is2xxSuccessful();

    private final Function<Object, Function<ClientResponse, Mono<? extends Throwable>>> EXCEPTION_FUNCTION =
            reqParam -> response -> response.bodyToMono(String.class).handle((responseBody, sink) -> {
                HttpRequest request = response.request();
                log.error("[midjourney-api] 调用失败！请求方式:[{}]，请求地址:[{}]，请求参数:[{}]，响应数据: [{}]",
                        request.getMethod(), request.getURI(), reqParam, responseBody);
                sink.error(new IllegalStateException("[midjourney-api] 调用失败！"));
            });

    private final WebClient webClient;

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

    public MidjourneyApi(String baseUrl, String apiKey, String notifyUrl) {
        this.webClient = WebClient.builder()
                .baseUrl(baseUrl)
                .defaultHeaders(httpHeaders -> {
                    httpHeaders.setContentType(MediaType.APPLICATION_JSON);
                    httpHeaders.setBearerAuth(apiKey);
                })
                .build();
        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) {
        return webClient.post()
                .uri(uri)
                .body(Mono.just(JsonUtils.toJsonString(body)), String.class)
                .retrieve()
                .onStatus(STATUS_PREDICATE, EXCEPTION_FUNCTION.apply(body))
                .bodyToMono(String.class)
                .block();
    }

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

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

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

        public ImagineRequest(List<String> base64Array, String prompt, String notifyHook, String state) {
            this.base64Array = base64Array;
            this.prompt = prompt;
            this.notifyHook = notifyHook;
            this.state = 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;

        public ActionRequest(String taskId, String customId, String notifyHook) {
            this.customId = customId;
            this.taskId = taskId;
            this.notifyHook = notifyHook;
        }

    }

    /**
     * Submit 统一返回
     *
     * @param code 状态码: 1(提交成功), 21(已存在), 22(排队中), other(错误)
     * @param description 描述
     * @param properties 扩展字段
     * @param result 任务ID
     */
    public record SubmitResponse(String code,
                                 String description,
                                 Map<String, Object> properties,
                                 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 任务完成后的可执行按钮
     */
    public record Notify(String id,
                         String action,
                         String status,

                         String prompt,
                         String promptEn,

                         String description,
                         String state,

                         Long submitTime,
                         Long startTime,
                         Long finishTime,

                         String imageUrl,
                         String progress,
                         String failReason,
                         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）
     */
    public record Button(String customId,
                         String emoji,
                         String label,
                         String type,
                         String style) {
    }

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

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

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

        private final String model;
        private final String name;

    }

    /**
     * 提交返回的状态码的枚举
     */
    @Getter
    @AllArgsConstructor
    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;

    }

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

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

    }

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

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

        private final int order;

    }

}