package org.dromara.ai.service;

import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.OkHttpClient;
import org.dromara.ai.painting.PaintingEngineType;
import org.dromara.ai.painting.PaintingProvider;
import org.dromara.ai.painting.PaintingStatus;
import org.dromara.ai.painting.PaintingTask;
import org.dromara.ai.painting.provider.DallE3Provider;
import org.dromara.ai.painting.provider.MidJourneyProvider;
import org.dromara.ai.painting.provider.StableDiffusionProvider;
import org.dromara.common.redis.utils.RedisUtils;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.time.Instant;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * AI绘画服务
 * 提供统一的AI绘画任务管理、状态跟踪和进度推送功能
 *
 * @author pengles
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class AiPaintingService {
    private final DallE3Provider dallE3Provider;
    private final MidJourneyProvider midJourneyProvider;
    private final StableDiffusionProvider stableDiffusionProvider;
    private final ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
    private final Map<String, SseEmitter> emitters = new ConcurrentHashMap<>();
    private final OkHttpClient client = new OkHttpClient();

    @PostConstruct
    public void init() {
    }

    public String createTask(PaintingEngineType engine, String prompt, Map<String, Object> params) {
        PaintingTask task = new PaintingTask();
        task.setId(UUID.randomUUID().toString());
        task.setEngine(engine);
        task.setPrompt(prompt);
        task.setParams(params != null ? params : new HashMap<>());
        task.setStatus(PaintingStatus.PENDING);
        task.setProgress(0);
        task.setUpdatedAt(Instant.now());
        saveTask(task);
        executor.submit(() -> runTask(task.getId()));
        return task.getId();
    }

    public PaintingTask getTask(String taskId) {
        PaintingTask t = RedisUtils.getCacheObject(key(taskId));
        return t;
    }

    public void addEmitter(String taskId, SseEmitter emitter) {
        emitters.put(taskId, emitter);
    }

    private void runTask(String taskId) {
        PaintingTask task = getTask(taskId);
        if (task == null) return;
        task.setStatus(PaintingStatus.RUNNING);
        task.setProgress(10);
        task.setUpdatedAt(Instant.now());
        saveTask(task);
        push(taskId, Map.of("status", "RUNNING", "progress", task.getProgress()));
        PaintingProvider provider = resolveProvider(task.getEngine());
        try {
            List<String> images = provider.create(task.getPrompt(), task.getParams());
            task.setImages(images);
            task.setProgress(100);
            task.setStatus(PaintingStatus.SUCCESS);
            task.setUpdatedAt(Instant.now());
            saveTask(task);
            push(taskId, Map.of("status", "SUCCESS", "progress", 100, "images", images));
            complete(taskId);
        } catch (Exception e) {
            task.setStatus(PaintingStatus.FAILED);
            task.setMessage(e.getMessage());
            task.setUpdatedAt(Instant.now());
            saveTask(task);
            push(taskId, Map.of("status", "FAILED", "message", e.getMessage()));
            complete(taskId);
        }
    }

    private PaintingProvider resolveProvider(PaintingEngineType type) {
        if (type == PaintingEngineType.DALL_E_3) return dallE3Provider;
        if (type == PaintingEngineType.MIDJOURNEY) return midJourneyProvider;
        return stableDiffusionProvider;
    }

    private void push(String taskId, Object data) {
        SseEmitter emitter = emitters.get(taskId);
        if (emitter == null) return;
        try {
            emitter.send(data);
        } catch (Exception ignored) {
        }
    }

    private void complete(String taskId) {
        SseEmitter emitter = emitters.remove(taskId);
        if (emitter == null) return;
        try {
            emitter.send("[DONE]");
            emitter.complete();
        } catch (Exception ignored) {
        }
    }

    private void saveTask(PaintingTask task) {
        RedisUtils.setCacheObject(key(task.getId()), task);
    }

    private String key(String taskId) {
        return "ai:painting:task:" + taskId;
    }
}

