
package com.rj.controller;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.*;
import org.springframework.web.client.RestTemplate;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class PosterGenerationTest {

    private static final Logger log = LoggerFactory.getLogger(PosterGenerationTest.class);
    private final RestTemplate restTemplate = new RestTemplate();
    private final ObjectMapper objectMapper = new ObjectMapper();

    public static void main(String[] args) {
        System.out.println("程序开始运行...");
        try {
            new PosterGenerationTest().testPosterGeneration();
        } catch (Exception e) {
            System.out.println("程序运行出错: " + e.getMessage());
            e.printStackTrace();
        }
        System.out.println("程序结束");
    }
    public void testPosterGeneration() {
        try {
            System.out.println("开始测试阿里云海报生成API");
            log.info("开始测试阿里云海报生成API");

            // 验证API Key
            String apiKey = System.getenv("DASHSCOPE_API_KEY");
            if (apiKey == null || apiKey.isEmpty()) {
                System.out.println("DASHSCOPE_API_KEY 环境变量未设置");
                log.error("DASHSCOPE_API_KEY 环境变量未设置");
                return;
            }
            System.out.println("API Key 已设置，开始构建请求...");
            System.out.println("API Key: " + apiKey.substring(0, 10) + "...");

            // 构建请求体 - 使用阿里云官方文档的标准格式
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", "wanx-poster-generation-v1");

            // 设置输入参数 - 使用完整的官方参数
            Map<String, Object> input = new HashMap<>();
            input.put("title", "春节快乐");
            input.put("sub_title", "家庭团聚，共享天伦之乐");
            input.put("body_text", "春节是中国最重要的传统节日之一，它象征着新的开始和希望");
            input.put("prompt_text_zh", "灯笼，小猫，梅花");
            input.put("wh_ratios", "竖版");
            input.put("lora_name", "童话油画");
            input.put("lora_weight", 0.8);
            input.put("ctrl_ratio", 0.7);
            input.put("ctrl_step", 0.7);
            input.put("generate_mode", "generate");
            // auxiliary_parameters 在首次调用时应该不传或者传空字符串
//            input.put("auxiliary_parameters", "[]");
            input.put("generate_num", 1);

            requestBody.put("input", input);
            requestBody.put("parameters", new HashMap<>());

            // 构建请求头
            HttpHeaders headers = new HttpHeaders();
            headers.set("X-DashScope-Async", "enable"); // 启用异步调用
            headers.set("Authorization", "Bearer " + apiKey);
            headers.setContentType(MediaType.APPLICATION_JSON);

            // 创建HTTP实体
            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

            // 发送请求
            String apiUrl = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis";
            System.out.println("发送请求到阿里云海报生成API: " + apiUrl);
            System.out.println("请求体: " + objectMapper.writeValueAsString(requestBody));
            log.info("发送请求到阿里云海报生成API: {}", apiUrl);
            log.info("请求体: {}", objectMapper.writeValueAsString(requestBody)); // 打印请求体

            ResponseEntity<String> response = restTemplate.exchange(
                    apiUrl,
                    HttpMethod.POST,
                    entity,
                    String.class
            );

            // 处理响应
            System.out.println("响应状态码: " + response.getStatusCode());
            System.out.println("响应体: " + response.getBody());
            log.info("响应状态码: {}", response.getStatusCode());
            log.info("响应体: {}", response.getBody());

            if (response.getStatusCode().is2xxSuccessful()) {
                System.out.println("✅ 阿里云海报生成API调用成功");
                log.info("✅ 阿里云海报生成API调用成功");

                // 解析响应获取task_id
                JsonNode responseNode = objectMapper.readTree(response.getBody());
                JsonNode outputNode = responseNode.path("output");

                if (!outputNode.isMissingNode() && outputNode.has("task_id")) {
                    String taskId = outputNode.path("task_id").asText();
                    String taskStatus = outputNode.path("task_status").asText();

                    log.info("任务ID: {}", taskId);
                    log.info("任务状态: {}", taskStatus);

                    // 轮询获取任务结果
                    getTaskResult(apiKey, taskId);
                } else {
                    log.error("❌ 响应中未找到task_id或output节点");
                }
            } else {
                log.error("❌ 阿里云海报生成API调用失败，状态码: {}", response.getStatusCode());
                log.error("❌ 响应详情: {}", response.getBody());
            }

        } catch (Exception e) {
            log.error("测试阿里云海报生成API失败: {}", e.getMessage(), e);
        }
    }

    /**
     * 轮询获取任务结果
     * @param apiKey API Key
     * @param taskId 任务ID
     */
    private void getTaskResult(String apiKey, String taskId) {
        String resultUrl = "https://dashscope.aliyuncs.com/api/v1/tasks/" + taskId;
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + apiKey);
        HttpEntity<String> entity = new HttpEntity<>(headers);

        int maxRetries = 30;
        int intervalSeconds = 5;

        for (int i = 0; i < maxRetries; i++) {
            try {
                ResponseEntity<String> response = restTemplate.exchange(resultUrl, HttpMethod.GET, entity, String.class);
                if (response.getStatusCode().is2xxSuccessful()) {
                    JsonNode resultNode = objectMapper.readTree(response.getBody());
                    JsonNode statusNode = resultNode.path("output").path("task_status");

                    if (!statusNode.isMissingNode()) {
                        String status = statusNode.asText();
                        log.info("当前任务状态: {}", status);

                        if ("SUCCEEDED".equalsIgnoreCase(status)) {
                            log.info("🎉 任务执行成功！");
                            log.info("任务结果详情: {}", response.getBody());
                            return;
                        } else if ("FAILED".equalsIgnoreCase(status)) {
                            log.error("❌ 任务执行失败！");
                            log.error("失败详情: {}", response.getBody());
                            // 如果失败，打印更详细的错误信息
                            JsonNode messageNode = resultNode.path("output").path("message");
                            if (!messageNode.isMissingNode()) {
                                log.error("失败具体原因: {}", messageNode.asText());
                            }
                            return;
                        }
                        // 如果是 RUNNING 或其他中间状态，继续轮询
                    } else {
                        log.error("❌ 获取任务状态失败，响应中缺少 task_status 字段。");
                        log.error("响应内容: {}", response.getBody());
                        return;
                    }
                } else {
                    log.error("❌ 查询任务状态失败，状态码: {}", response.getStatusCode());
                    log.error("响应内容: {}", response.getBody());
                }
            } catch (Exception e) {
                log.error("轮询任务结果时发生异常: {}", e.getMessage(), e);
            }

            log.info("等待 {} 秒后继续轮询...", intervalSeconds);
            try {
                TimeUnit.SECONDS.sleep(intervalSeconds);
            } catch (InterruptedException ie) {
                Thread.currentThread().interrupt();
                log.error("轮询被中断");
                return;
            }
        }

        log.error("❌ 超时，未能获取最终任务结果。");
    }
}



