activity-tracker/src/ai/classifier.rs

/// AI-powered activity classifier using NPU
use std::path::PathBuf;
use std::sync::Arc;
use ort::session::Session;

use crate::error::{Result, AppError};
use crate::analysis::{ActivityCategory, Entities};
use super::npu::NpuDevice;

pub struct NpuClassifier {
    npu: NpuDevice,
    session: Option<Arc<Session>>,
    model_path: Option<PathBuf>,
}

impl NpuClassifier {
    /// Create a new NPU classifier
    pub fn new() -> Self {
        let npu = NpuDevice::detect();

        Self {
            npu,
            session: None,
            model_path: None,
        }
    }

    /// Load a model from file
    pub fn load_model(&mut self, model_path: PathBuf) -> Result<()> {
        log::info!("Loading AI model from: {}", model_path.display());

        if !model_path.exists() {
            return Err(AppError::Analysis(format!(
                "Model file not found: {}",
                model_path.display()
            )));
        }

        let session = self.npu.create_session(
            model_path.to_str().ok_or_else(|| {
                AppError::Analysis("Invalid model path".to_string())
            })?
        )?;

        self.session = Some(Arc::new(session));
        self.model_path = Some(model_path);

        log::info!("Model loaded successfully on {}", self.npu.device_name());
        Ok(())
    }

    /// Classify activity using NPU-accelerated model
    pub fn classify(&self, window_title: &str, process_name: &str) -> Result<(ActivityCategory, f32, Entities)> {
        // If no model is loaded, fall back to rule-based classification
        if self.session.is_none() {
            log::debug!("No AI model loaded, using rule-based classification");
            return self.classify_rule_based(window_title, process_name);
        }

        // TODO: Implement real neural network inference
        // For now, use rule-based as fallback
        self.classify_rule_based(window_title, process_name)
    }

    /// Rule-based classification fallback
    fn classify_rule_based(&self, window_title: &str, process_name: &str) -> Result<(ActivityCategory, f32, Entities)> {
        use crate::analysis::Classifier;
        use crate::capture::WindowMetadata;

        // Create a temporary WindowMetadata for classification
        let metadata = WindowMetadata {
            title: window_title.to_string(),
            process_name: process_name.to_string(),
            process_id: 0,
            is_active: true,
        };

        let classifier = Classifier::new();
        let classification = classifier.classify(&metadata);

        Ok((classification.category, classification.confidence, classification.entities))
    }

    /// Check if NPU is available
    pub fn is_npu_available(&self) -> bool {
        self.npu.is_available()
    }

    /// Get device information
    pub fn device_info(&self) -> &str {
        self.npu.device_name()
    }

    /// Check if a model is loaded
    pub fn is_model_loaded(&self) -> bool {
        self.session.is_some()
    }
}

impl Default for NpuClassifier {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_npu_classifier_creation() {
        let classifier = NpuClassifier::new();
        assert!(!classifier.is_model_loaded());
    }

    #[test]
    fn test_rule_based_classification() {
        let classifier = NpuClassifier::new();
        let result = classifier.classify("VSCode - main.rs", "code.exe");
        assert!(result.is_ok());

        let (category, confidence, _entities) = result.unwrap();
        assert!(matches!(category, ActivityCategory::Development));
        assert!(confidence > 0.0);
    }
}