package org.dromara.ai.queue;

import lombok.Data;
import org.springframework.stereotype.Service;

import java.time.Instant;
import java.util.Comparator;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.*;

@Service
public class TaskQueueService {
    private final PriorityBlockingQueue<TaskItem> queue = new PriorityBlockingQueue<>(100, Comparator.comparingInt(TaskItem::getPriority).reversed());
    private final ExecutorService workers = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
    private final Map<String, TaskItem> tasks = new ConcurrentHashMap<>();

    public String submit(String type, int priority, Runnable r) {
        String id = UUID.randomUUID().toString();
        TaskItem item = new TaskItem();
        item.id = id;
        item.type = type;
        item.priority = priority;
        item.createdAt = Instant.now();
        item.status = "PENDING";
        tasks.put(id, item);
        queue.offer(item);
        workers.submit(() -> {
            TaskItem t = tasks.get(id);
            t.status = "RUNNING";
            try {
                r.run();
                t.status = "SUCCESS";
            } catch (Exception e) {
                t.status = "FAILED";
                t.message = e.getMessage();
            } finally {
                t.updatedAt = Instant.now();
            }
        });
        return id;
    }

    public Map<String, Object> metrics() {
        return Map.of("depth", queue.size(), "active", ((ThreadPoolExecutor) workers).getActiveCount());
    }

    public TaskItem get(String id) {
        return tasks.get(id);
    }

    @Data
    public static class TaskItem {
        private String id;
        private String type;
        private int priority;
        private String status;
        private String message;
        private Instant createdAt;
        private Instant updatedAt;

        public int getPriority() {
            return priority;
        }
    }
}
