//! YOLO detection loader with frame caching

use anyhow::{Context, Result};
use lru::LruCache;
use serde::Deserialize;
use std::collections::HashMap;
use std::fs::File;
use std::io::BufReader;
use std::num::NonZeroUsize;
use std::path::Path;

#[derive(Debug, Clone, Deserialize)]
pub struct Detection {
    pub class_id: u32,
    pub class_name: String,
    pub confidence: f64,
    pub x1: f64,
    pub y1: f64,
    pub x2: f64,
    pub y2: f64,
}

#[derive(Debug, Clone, Deserialize)]
pub struct FrameData {
    pub frame_number: u64,
    pub time_seconds: f64,
    pub time_formatted: String,
    pub detections: Vec<Detection>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YoloMetadata {
    pub video_path: String,
    pub model_path: Option<String>,
    pub width: u32,
    pub height: u32,
    pub fps: f64,
    pub total_frames: u64,
    pub total_duration: f64,
    pub processed_at: Option<String>,
    pub status: Option<String>,
    pub total_detections: u64,
    pub avg_detections_per_frame: f64,
    #[serde(default)]
    pub auto_save_interval: Option<u32>,
    #[serde(default)]
    pub processing_time: Option<f64>,
    #[serde(default)]
    pub avg_time_per_frame: Option<f64>,
    #[serde(default)]
    pub last_saved_at: Option<String>,
    #[serde(default)]
    pub completed_at: Option<String>,
    #[serde(default)]
    pub auto_save_count: Option<u32>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YoloData {
    pub metadata: YoloMetadata,
    pub frames: HashMap<String, FrameData>,
}

pub struct YoloLoader {
    data: YoloData,
    cache: LruCache<u64, Vec<Detection>>,
}

impl YoloLoader {
    const CACHE_SIZE: usize = 60;

    pub fn load(path: &Path) -> Result<Self> {
        let file =
            File::open(path).with_context(|| format!("Failed to open YOLO file: {:?}", path))?;
        let reader = BufReader::new(file);

        let data: YoloData =
            serde_json::from_reader(reader).with_context(|| "Failed to parse YOLO JSON")?;

        let cache = LruCache::new(NonZeroUsize::new(Self::CACHE_SIZE).unwrap());

        Ok(Self { data, cache })
    }

    pub fn get_detections(&mut self, frame: u64) -> Vec<Detection> {
        if let Some(dets) = self.cache.get(&frame) {
            return dets.clone();
        }

        if let Some(frame_data) = self.data.frames.get(&frame.to_string()) {
            let dets = frame_data.detections.clone();
            self.cache.put(frame, dets.clone());
            dets
        } else {
            Vec::new()
        }
    }

    pub fn get_detections_at_time(&mut self, time_ms: u64) -> Vec<Detection> {
        let fps = self.data.metadata.fps;
        let frame = ((time_ms as f64 / 1000.0) * fps) as u64;
        self.get_detections(frame)
    }

    pub fn metadata(&self) -> &YoloMetadata {
        &self.data.metadata
    }

    pub fn fps(&self) -> f64 {
        self.data.metadata.fps
    }

    pub fn total_frames(&self) -> u64 {
        self.data.metadata.total_frames
    }
}