src-tauri/src/queue.rs

use std::collections::BinaryHeap;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use llama_cpp_2::token::LlamaToken;

#[derive(Clone)]
pub struct ChatRequest {
    pub task_id: String,
    pub prompt: String,
    pub device_name: String,
    pub tokens: Vec<LlamaToken>,
}

pub struct QueueItem {
    pub request: ChatRequest,
    pub usage_count: u64,
    pub timestamp: Instant,
    pub response_tx: tokio::sync::mpsc::Sender<Result<String, String>>,
}

// We want a Min-Heap based on usage_count, and then FIFO for timestamp.
// BinaryHeap is a Max-Heap, so we reverse the ordering.
impl PartialEq for QueueItem {
    fn eq(&self, other: &Self) -> bool {
        self.usage_count == other.usage_count && self.timestamp == other.timestamp
    }
}

impl Eq for QueueItem {}

impl PartialOrd for QueueItem {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for QueueItem {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        // Lower usage_count = higher priority (greater in Max-Heap)
        match other.usage_count.cmp(&self.usage_count) {
            std::cmp::Ordering::Equal => {
                // Older timestamp = higher priority
                other.timestamp.cmp(&self.timestamp)
            }
            ord => ord,
        }
    }
}

pub struct QueueManager {
    heap: Arc<Mutex<BinaryHeap<QueueItem>>>,
    notifier: Arc<tokio::sync::Notify>,
}

impl QueueManager {
    pub fn new() -> Self {
        Self {
            heap: Arc::new(Mutex::new(BinaryHeap::new())),
            notifier: Arc::new(tokio::sync::Notify::new()),
        }
    }

    pub fn enqueue(&self, item: QueueItem) {
        self.heap.lock().unwrap().push(item);
        self.notifier.notify_one();
    }

    pub fn get_position_by_device(&self, device_name: &str) -> Option<usize> {
        let heap = self.heap.lock().unwrap();
        let mut items: Vec<_> = heap.iter().collect();
        items.sort();
        let position = items.iter().rev().position(|item| item.request.device_name == device_name);
        position.map(|p| p + 1)
    }

    pub async fn wait_for_next(&self) -> QueueItem {
        loop {
            if let Some(item) = self.heap.lock().unwrap().pop() {
                return item;
            }
            self.notifier.notified().await;
        }
    }
}