package com.rj.scheduler;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.rj.config.AliyunConfig;
import com.rj.entity.ImageModel;
import com.rj.service.IImageModelService;
import com.rj.service.MinIOService;
import com.rj.utils.MinIOUrlGenerator;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.time.LocalDateTime;
import java.util.List;

/**
 * 图像模型任务状态检查定时器
 * 每5分钟检查一次PENDING或RUNNING状态的任务，调用阿里云API查询任务状态并更新数据库
 * 
 * @author rj
 * @date 2025-01-30
 */
@Slf4j
@Component
public class ImageModelStatusScheduler {

    @Autowired
    private IImageModelService imageModelService;

    @Autowired
    private AliyunConfig aliyunConfig;

    @Autowired
    private RestTemplate restTemplate;

    @Autowired
    private MinIOService minIOService;

    @Autowired
    private MinIOUrlGenerator urlGenerator;

    /**
     * 每5分钟执行一次任务状态检查
     */
//    @Scheduled(fixedRate = 3 * 60 * 1000) // 5分钟 = 5 * 60 * 1000毫秒
    public void checkImageModelStatus() {
        try {
            log.info("开始执行图像模型任务状态检查...");
            
            // 查询所有PENDING或RUNNING状态的任务
            List<ImageModel> pendingTasks = getPendingTasks();
            
            if (pendingTasks.isEmpty()) {
                log.info("没有找到PENDING或RUNNING状态的图像任务");
                return;
            }
            
            log.info("找到{}个PENDING或RUNNING状态的图像任务，开始检查状态", pendingTasks.size());
            
            // 遍历每个任务，检查状态
            for (ImageModel task : pendingTasks) {
                try {
                    checkAndUpdateTaskStatus(task);
                } catch (Exception e) {
                    log.error("检查图像任务状态失败，任务ID: {}, 错误: {}", task.getTaskId(), e.getMessage(), e);
                }
            }
            
            log.info("图像模型任务状态检查完成");
            
        } catch (Exception e) {
            log.error("图像模型任务状态检查执行失败", e);
        }
    }

    /**
     * 查询所有PENDING或RUNNING状态的任务
     */
    private List<ImageModel> getPendingTasks() {
        QueryWrapper<ImageModel> queryWrapper = new QueryWrapper<>();
        queryWrapper.in("task_status", "PENDING", "RUNNING");
        
        return imageModelService.list(queryWrapper);
    }

    /**
     * 检查并更新单个任务状态
     */
    private void checkAndUpdateTaskStatus(ImageModel task) {
        try {
            log.info("检查图像任务状态，任务ID: {}", task.getTaskId());
            if (task == null || task.getTaskId() == null || task.getTaskId().isEmpty()) return;
            
            // 调用阿里云API检查任务状态
            String apiUrl = "https://dashscope.aliyuncs.com/api/v1/tasks/" + task.getTaskId();
            
            HttpHeaders headers = new HttpHeaders();
            headers.set("Authorization", "Bearer " + aliyunConfig.getApiKey());
            headers.set("Content-Type", "application/json");
            
            HttpEntity<String> entity = new HttpEntity<>(headers);
            
            ResponseEntity<String> response = restTemplate.exchange(
                apiUrl, 
                HttpMethod.GET, 
                entity, 
                String.class
            );
            
            if (response.getStatusCode().is2xxSuccessful()) {
                String responseBody = response.getBody();
                log.info("阿里云API响应: {}", responseBody);
                
                // 解析响应并更新数据库
                updateTaskFromResponse(task, responseBody);
            } else {
                log.error("阿里云API调用失败，状态码: {}, 任务ID: {}", response.getStatusCode(), task.getTaskId());
            }
            
        } catch (Exception e) {
            log.error("检查图像任务状态异常，任务ID: {}, 错误: {}", task.getTaskId(), e.getMessage(), e);
        }
    }

    /**
     * 根据API响应更新任务状态
     */
    private void updateTaskFromResponse(ImageModel task, String responseBody) {
        try {
            JSONObject responseJson = JSON.parseObject(responseBody);
            
            if (responseJson.containsKey("output")) {
                JSONObject output = responseJson.getJSONObject("output");
                
                // 更新任务状态
                String taskStatus = output.getString("task_status");
                task.setTaskStatus(taskStatus);
                task.setUpdateTime(LocalDateTime.now());
                
                // 如果任务成功完成，处理生成的图片
                if ("SUCCEEDED".equals(taskStatus)) {
                    log.info("图像任务成功完成，output: {}", output);
                    
                    // 处理生成的图片结果
                    processSuccessfulTask(task, output);
                    
                } else if ("FAILED".equals(taskStatus)) {
                    log.error("图像任务失败，任务ID: {}, 状态: {}", task.getTaskId(), taskStatus);
                    task.setUpdateTime(LocalDateTime.now());
                }
                
                // 更新数据库
                imageModelService.updateById(task);
                log.info("图像任务状态已更新，任务ID: {}, 状态: {}", task.getTaskId(), taskStatus);
                
            } else {
                log.warn("API响应格式异常，任务ID: {}, 响应: {}", task.getTaskId(), responseBody);
            }
            
        } catch (Exception e) {
            log.error("解析API响应失败，任务ID: {}, 错误: {}", task.getTaskId(), e.getMessage(), e);
        }
    }

    /**
     * 处理成功完成的任务
     */
    private void processSuccessfulTask(ImageModel task, JSONObject output) {
        try {
            // 获取生成的图片结果
            JSONArray results = output.getJSONArray("results");
            results.iterator().forEachRemaining(result -> {
                JSONObject resultObj = (JSONObject) result;
                String url = resultObj.getString("url");
                log.info("获取到生成的图片URL: {}", url);
                String[] uploadResult = downloadAndUploadToMinIO(url, task);
                if (uploadResult != null && uploadResult.length >= 2) {
                    String minioImageUrl = uploadResult[0];
                    String objectName = uploadResult[1];
                    task.setResultImageUrl(minioImageUrl);
                    // 生成7天有效期的临时URL
                    MinIOUrlGenerator.UrlInfo tempUrlInfo = urlGenerator.generateTempUrl(objectName, 7 * 24 * 3600);
                    if (tempUrlInfo != null && tempUrlInfo.isSuccess()) {
                        task.setResultImageTempUrl(tempUrlInfo.getUrl());
                        log.info("生成7天有效期临时URL: {}", tempUrlInfo.getUrl());
                    }

                    log.info("图片已成功转存到MinIO: {}", minioImageUrl);
                } else {
                    // 如果转存失败，保留原始URL
                    task.setResultImageUrl(url);
                    log.warn("图片转存到MinIO失败，保留原始URL: {}", url);
                }
            });
        } catch (Exception e) {
            log.error("处理成功任务时发生异常，任务ID: {}, 错误: {}", task.getTaskId(), e.getMessage(), e);
        }
    }

    /**
     * 从阿里云OSS下载图片并上传到MinIO
     */
    private String[] downloadAndUploadToMinIO(String imageUrl, ImageModel task) {
        try {
            log.info("开始从阿里云下载图片: {}", imageUrl);
            
            // 下载图片数据
            byte[] imageData = downloadImageFromUrl(imageUrl);
            if (imageData == null || imageData.length == 0) {
                log.error("下载图片失败，图片URL: {}", imageUrl);
                return null;
            }
            
            // 生成MinIO存储路径
            String fileName = "generated_image_" + task.getTaskId() + "_" + System.currentTimeMillis() + ".jpg";

            
            // 上传到MinIO
            String minioUrl = minIOService.uploadFile(imageData, fileName, "image/jpeg");
            
            if (minioUrl != null) {
                log.info("图片已成功上传到MinIO: {}", minioUrl);
                return new String[]{minioUrl, fileName};
            } else {
                log.error("上传图片到MinIO失败");
                return null;
            }
            
        } catch (Exception e) {
            log.error("下载并上传图片到MinIO失败，图片URL: {}, 错误: {}", imageUrl, e.getMessage(), e);
            return null;
        }
    }

    /**
     * 从URL下载图片数据
     */
    private byte[] downloadImageFromUrl(String imageUrl) {
        try {
            URL url = new URL(imageUrl);
            HttpURLConnection connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("GET");
            connection.setConnectTimeout(30000); // 30秒连接超时
            connection.setReadTimeout(60000);   // 60秒读取超时
            
            int responseCode = connection.getResponseCode();
            if (responseCode == HttpURLConnection.HTTP_OK) {
                try (InputStream inputStream = connection.getInputStream();
                     ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
                    
                    byte[] buffer = new byte[4096];
                    int bytesRead;
                    while ((bytesRead = inputStream.read(buffer)) != -1) {
                        outputStream.write(buffer, 0, bytesRead);
                    }
                    
                    return outputStream.toByteArray();
                }
            } else {
                log.error("下载图片失败，HTTP状态码: {}, URL: {}", responseCode, imageUrl);
                return null;
            }
            
        } catch (Exception e) {
            log.error("下载图片异常，URL: {}, 错误: {}", imageUrl, e.getMessage(), e);
            return null;
        }
    }
}
