supper/src/control_loop.rs

use std::{
    collections::HashMap,
    time::{Duration, Instant}, thread,
};

use anyhow::Result;
use egui_extras::RetainedImage;
use image::RgbImage;
use scrap::{Capturer, Display};


use crate::{
    capture,
    image_processing::{self, hash_image},
    ocr,
    state::{AppState, DebugOcrFrame, ParsedFrame, RaceState, SharedAppState},
};

fn is_finished_lap(state: &AppState, frame: &ParsedFrame) -> bool {
    if let Some(race) = &state.current_race {
        if let Some(last_finish) = &race.last_lap_record_time {
            let diff = Instant::now().duration_since(*last_finish);
            if diff < Duration::from_secs(5) {
                return false;
            }
        }
    }
    if let Some(prev_frame) = state.buffered_frames.back() {
        prev_frame.lap_time.is_some()
            && prev_frame.lap_time == frame.lap_time
            && prev_frame.lap_time.unwrap() > Duration::from_secs(5)
    } else {
        false
    }
}

fn merge_with_max<T: Ord + Clone>(a: &Option<T>, b: &Option<T>) -> Option<T> {
    match (a, b) {
        (Some(a), None) => Some(a.clone()),
        (None, Some(b)) => Some(b.clone()),
        (None, None) => None,
        (Some(a), Some(b)) => {
            if a > b {
                Some(a.clone())
            } else {
                Some(b.clone())
            }
        }
    }
}

fn merge_frames(prev: &ParsedFrame, next: &ParsedFrame) -> ParsedFrame {
    ParsedFrame {
        lap: merge_with_max(&prev.lap, &next.lap),
        health: merge_with_max(&prev.health, &next.health),
        gas: merge_with_max(&prev.gas, &next.gas),
        tyres: merge_with_max(&prev.tyres, &next.tyres),
        lap_time: merge_with_max(&prev.lap_time, &next.lap_time),
        best_time: merge_with_max(&prev.best_time, &next.best_time),
    }
}

fn handle_new_frame(state: &mut AppState, frame: ParsedFrame, image: &RgbImage) {
    if frame.lap_time.is_some() {
        state.last_frame = Some(frame.clone());
        state.frames_without_lap = 0;

        if state.current_race.is_none() {
            let mut race = RaceState::default();
            race.screencap = Some(
                RetainedImage::from_image_bytes(
                    "screencap",
                    &image_processing::to_png_bytes(image),
                )
                .expect("failed to save screenshot"),
            );
            state.current_race = Some(race);
        }
    } else {
        state.frames_without_lap += 1;
    }

    if state.frames_without_lap >= 10 {
        if let Some(race) = state.current_race.take() {
            state.past_races.push_front(race);
        }
    }

    if is_finished_lap(state, &frame) {
        let mut merged = merge_frames(state.buffered_frames.back().unwrap(), &frame);
        if let Some(lap) = &merged.lap {
            merged.lap = Some(lap - 1);
        }

        if let Some(race) = state.current_race.as_mut() {
            race.laps.push(merged);
            race.last_lap_record_time = Some(Instant::now());
        }
    }

    state.buffered_frames.push_back(frame);
    if state.buffered_frames.len() >= 20 {
        state.buffered_frames.pop_front();
    }
}

fn run_loop_once(state: &SharedAppState) -> Result<()> {
    let mut capturer = Capturer::new(Display::primary()?)?;
    let config = state.lock().unwrap().config.clone();
    let learned_config = state.lock().unwrap().learned.clone();
    let frame = capture::get_frame(&mut capturer)?;
    let ocr_results = ocr::ocr_all_regions(&frame, config.clone(), learned_config.clone());

    let mut saved_frames = HashMap::new();

    if state.lock().unwrap().debug_frames {
        let hasher = img_hash::HasherConfig::new().to_hasher();
        for region in &config.ocr_regions {
            let mut extracted = image_processing::extract_region(&frame, region);
            image_processing::filter_to_white(&mut extracted);
            let retained = RetainedImage::from_image_bytes(
                &region.name,
                &image_processing::to_png_bytes(&extracted),
            )
            .unwrap();
            let hash = hash_image(&extracted);
            saved_frames.insert(
                region.name.clone(),
                DebugOcrFrame {
                    image: retained,
                    rgb_image: extracted,
                    img_hash: hash,
                },
            );
        }
    }
    {
        let mut state = state.lock().unwrap();
        let parsed = ParsedFrame::parse(&ocr_results);
        handle_new_frame(&mut state, parsed, &frame);
        state.raw_data = ocr_results;
        state.saved_frames = saved_frames;
    }
    Ok(())
}

pub fn run_control_loop(state: SharedAppState) {
    loop {
        if let Err(e) = run_loop_once(&state) {
            eprintln!("Error in control loop: {:?}", e)
        }
        thread::sleep(Duration::from_millis(500));
    }
}