package org.dromara.ai.controller.app;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import jakarta.validation.constraints.NotBlank;
import jakarta.validation.constraints.Pattern;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.dromara.common.core.domain.R;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * AI图像控制器
 * 提供DALL·E-3图像生成、编辑和变体功能
 *
 * @author pengles
 */
@Slf4j
@Validated
@RequiredArgsConstructor
@RestController
@RequestMapping
public class AiImageController {

    @Value("${ruoyi.ai.external-ai.openai.base-url:https://api.openai.com}")
    private String openaiBaseUrl;
    @Value("${ruoyi.ai.external-ai.openai.api-key:}")
    private String openaiApiKey;

    private final OkHttpClient client = new OkHttpClient();

    /**
     * DALL·E-3 图像生成
     *
     * @param prompt 提示词
     * @param size   分辨率 (512x512 或 1024x1024)
     * @return 生成结果
     */
    @Operation(summary = "DALL·E-3 生成", description = "通过OpenAI DALL·E-3模型生成图像")
    @PostMapping("/dall3")
    public R<Map<String, Object>> dalle3(
            @Parameter(description = "提示词") @RequestParam @NotBlank String prompt,
            @Parameter(description = "分辨率") @RequestParam @Pattern(regexp = "^(512x512|1024x1024)$") String size) {
        try {
            log.info("DALL·E-3图像生成请求 - 提示词长度: {}, 分辨率: {}", prompt.length(), size);
            
            if (openaiApiKey == null || openaiApiKey.isBlank()) {
                log.warn("OpenAI API Key未配置");
                return R.fail("OpenAI API Key未配置");
            }
            
            String url = normalize(openaiBaseUrl) + "/v1/images/generations";
            String body = "{\"model\":\"dall-e-3\",\"prompt\":\"" + prompt.replace("\"", "\\\"") + "\",\"size\":\"" + size + "\"}";
            Request.Builder b = new Request.Builder().url(url)
                    .post(RequestBody.create(MediaType.parse("application/json"), body.getBytes(StandardCharsets.UTF_8)))
                    .header("Authorization", "Bearer " + openaiApiKey);
            try (Response resp = client.newCall(b.build()).execute()) {
                if (!resp.isSuccessful()) {
                    String errorBody = resp.body() != null ? resp.body().string() : "";
                    log.error("DALL·E-3 API调用失败: {} - {}", resp.code(), errorBody);
                    return R.fail("OpenAI调用失败: HTTP " + resp.code());
                }
                String respBody = resp.body() != null ? resp.body().string() : "{}";
                log.info("DALL·E-3图像生成成功");
                return R.ok(Map.of("raw", respBody));
            }
        } catch (Exception e) {
            log.error("DALL·E-3图像生成失败", e);
            return R.fail("OpenAI调用失败: " + e.getMessage());
        }
    }

    /**
     * DALL·E-3 图像编辑
     *
     * @param prompt 提示词
     * @param size   分辨率 (512x512 或 1024x1024)
     * @param image  原图文件
     * @param mask   遮罩文件（可选）
     * @return 编辑结果
     */
    @Operation(summary = "DALL·E-3 编辑", description = "上传原图与遮罩进行图像编辑")
    @PostMapping("/dall3/edit")
    public R<Map<String, Object>> dalle3Edit(
            @Parameter(description = "提示词") @RequestParam @NotBlank String prompt,
            @Parameter(description = "分辨率") @RequestParam @Pattern(regexp = "^(512x512|1024x1024)$") String size,
            @Parameter(description = "原图") @RequestParam("image") MultipartFile image,
            @Parameter(description = "遮罩") @RequestParam(value = "mask", required = false) MultipartFile mask) {
        try {
            log.info("DALL·E-3图像编辑请求 - 提示词长度: {}, 分辨率: {}, 图片: {}", 
                    prompt.length(), size, image.getOriginalFilename());
            
            if (openaiApiKey == null || openaiApiKey.isBlank()) {
                log.warn("OpenAI API Key未配置");
                return R.fail("OpenAI API Key未配置");
            }
            
            String url = normalize(openaiBaseUrl) + "/v1/images/edits";
            MultipartBody.Builder mb = new MultipartBody.Builder().setType(MultipartBody.FORM)
                    .addFormDataPart("model", "dall-e-3")
                    .addFormDataPart("prompt", prompt)
                    .addFormDataPart("size", size)
                    .addFormDataPart("image", image.getOriginalFilename(), 
                            RequestBody.create(MediaType.parse("application/octet-stream"), toBytes(image)));
            if (mask != null && !mask.isEmpty()) {
                mb.addFormDataPart("mask", mask.getOriginalFilename(), 
                        RequestBody.create(MediaType.parse("application/octet-stream"), toBytes(mask)));
            }
            Request.Builder b = new Request.Builder().url(url)
                    .post(mb.build())
                    .header("Authorization", "Bearer " + openaiApiKey);
            try (Response resp = client.newCall(b.build()).execute()) {
                if (!resp.isSuccessful()) {
                    String errorBody = resp.body() != null ? resp.body().string() : "";
                    log.error("DALL·E-3编辑API调用失败: {} - {}", resp.code(), errorBody);
                    return R.fail("OpenAI调用失败: HTTP " + resp.code());
                }
                String respBody = resp.body() != null ? resp.body().string() : "{}";
                log.info("DALL·E-3图像编辑成功");
                return R.ok(Map.of("raw", respBody));
            }
        } catch (Exception e) {
            log.error("DALL·E-3图像编辑失败", e);
            return R.fail("OpenAI调用失败: " + e.getMessage());
        }
    }

    /**
     * DALL·E-3 图像变体生成
     *
     * @param size  分辨率 (512x512 或 1024x1024)
     * @param image 原图文件
     * @return 变体生成结果
     */
    @Operation(summary = "DALL·E-3 变体", description = "上传原图生成图像变体")
    @PostMapping("/dall3/variations")
    public R<Map<String, Object>> dalle3Variations(
            @Parameter(description = "分辨率") @RequestParam @Pattern(regexp = "^(512x512|1024x1024)$") String size,
            @Parameter(description = "原图") @RequestParam("image") MultipartFile image) {
        try {
            log.info("DALL·E-3图像变体生成请求 - 分辨率: {}, 图片: {}", size, image.getOriginalFilename());
            
            if (openaiApiKey == null || openaiApiKey.isBlank()) {
                log.warn("OpenAI API Key未配置");
                return R.fail("OpenAI API Key未配置");
            }
            
            String url = normalize(openaiBaseUrl) + "/v1/images/variations";
            MultipartBody.Builder mb = new MultipartBody.Builder().setType(MultipartBody.FORM)
                    .addFormDataPart("model", "dall-e-3")
                    .addFormDataPart("size", size)
                    .addFormDataPart("image", image.getOriginalFilename(), 
                            RequestBody.create(MediaType.parse("application/octet-stream"), toBytes(image)));
            Request.Builder b = new Request.Builder().url(url)
                    .post(mb.build())
                    .header("Authorization", "Bearer " + openaiApiKey);
            try (Response resp = client.newCall(b.build()).execute()) {
                if (!resp.isSuccessful()) {
                    String errorBody = resp.body() != null ? resp.body().string() : "";
                    log.error("DALL·E-3变体API调用失败: {} - {}", resp.code(), errorBody);
                    return R.fail("OpenAI调用失败: HTTP " + resp.code());
                }
                String respBody = resp.body() != null ? resp.body().string() : "{}";
                log.info("DALL·E-3图像变体生成成功");
                return R.ok(Map.of("raw", respBody));
            }
        } catch (Exception e) {
            log.error("DALL·E-3图像变体生成失败", e);
            return R.fail("OpenAI调用失败: " + e.getMessage());
        }
    }

    /**
     * URL标准化（去除尾部斜杠）
     */
    private String normalize(String base) {
        if (base == null || base.isBlank()) {
            return "";
        }
        if (base.endsWith("/")) {
            return base.substring(0, base.length() - 1);
        }
        return base;
    }

    /**
     * 将MultipartFile转换为字节数组
     */
    private byte[] toBytes(MultipartFile f) {
        try {
            return f.getBytes();
        } catch (Exception e) {
            log.error("转换文件为字节数组失败: {}", f.getOriginalFilename(), e);
            return new byte[0];
        }
    }
}
