antidrift/src/main.rs

use anyhow::Result;
use std::collections::HashMap;
use std::io::stdout;
use std::rc::Rc;
use std::time::{Duration, Instant};
mod window;

use ratatui::{
    crossterm::{
        event::{self, Event, KeyCode},
        execute,
        terminal::{
            disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen, SetTitle,
        },
        ExecutableCommand,
    },
    prelude::*,
    widgets::*,
};

#[derive(PartialEq)]
enum State {
    InputIntention,
    InputDuration,
    InProgress,
    ShouldQuit,
}

struct App {
    state: State,
    user_intention: String,
    user_duration_str: String,
    user_duration: Duration,
    current_window_title: Rc<String>,
    current_window_time: Instant,
    session_start: Instant,
    session_stats: HashMap<Rc<String>, Duration>,
    session_remaining: Duration,
    last_tick_50ms: Instant,
    last_tick_1s: Instant,
}

impl App {
    fn new() -> Self {
        let window_title = window::get_title_clean();

        App {
            state: State::InputIntention,
            user_intention: String::new(),
            user_duration_str: String::new(),
            user_duration: Duration::new(0, 0),
            current_window_title: window_title.into(),
            current_window_time: Instant::now(),
            session_start: Instant::now(),
            session_stats: HashMap::new(),
            session_remaining: Duration::new(0, 0),
            last_tick_50ms: Instant::now(),
            last_tick_1s: Instant::now(),
        }
    }

    fn handle_ticks(&mut self) {
        if self.last_tick_50ms.elapsed() >= Duration::from_millis(50) {
            self.tick_50ms();
        }
        if self.last_tick_1s.elapsed() >= Duration::from_secs(1) {
            self.tick_1s();
        }
    }

    fn to_in_progress(&mut self) {
        let Ok(user_duration) = self.user_duration_str.parse::<u64>() else {
            // TODO: Print error to the user
            return;
        };

        self.user_duration = Duration::from_secs(user_duration * 60);
        self.state = State::InProgress;
        self.current_window_time = Instant::now();
        self.session_start = self.current_window_time;
        self.session_stats = HashMap::new();
    }

    fn tick_50ms(&mut self) {
        match self.state {
            State::InputIntention | State::InputDuration => {
                window::minimize_other("AntiDrift");
            }
            State::InProgress => {
                let elapsed = self.session_start.elapsed();
                self.session_remaining = self.user_duration.saturating_sub(elapsed);
                update_session_stats(self);

                if self.session_remaining == Duration::new(0, 0) {
                    self.state = State::InputIntention;
                    self.user_intention = "Session complete. New session?".to_string();
                }
            }
            State::ShouldQuit => {}
        }

        self.last_tick_50ms = Instant::now();
    }

    fn tick_1s(&mut self) {
        self.last_tick_1s = Instant::now();
    }

    fn timeout(&self) -> Duration {
        Duration::from_millis(50).saturating_sub(self.last_tick_50ms.elapsed())
    }
}

fn duration_as_str(duration: &Duration) -> String {
    format!(
        "{:3}:{:02}",
        duration.as_secs() / 60,
        duration.as_secs() % 60
    )
}

fn session_stats_to_lines(session_stats: &HashMap<Rc<String>, Duration>) -> Vec<Line> {
    let mut stats: Vec<_> = session_stats
        .iter()
        .map(|(title, duration)| (title.clone(), *duration))
        .collect();
    stats.sort_by(|a, b| b.1.cmp(&a.1));
    stats
        .iter()
        .map(|(title, duration)| {
            Line::from(Span::raw(format!(
                "{}: {}",
                duration_as_str(&duration),
                title
            )))
        })
        .collect()
}

fn main() -> Result<()> {
    enable_raw_mode()?;
    stdout().execute(EnterAlternateScreen)?;
    execute!(stdout(), SetTitle("AntiDrift"))?;
    let mut terminal = Terminal::new(CrosstermBackend::new(stdout()))?;
    let mut app = App::new();

    while app.state != State::ShouldQuit {
        terminal.draw(|frame| ui(frame, &app))?;
        handle_events(&mut app)?;
        app.handle_ticks();
    }

    disable_raw_mode()?;
    stdout().execute(LeaveAlternateScreen)?;
    Ok(())
}

fn handle_events(app: &mut App) -> Result<()> {
    if !event::poll(app.timeout())? {
        return Ok(());
    }

    let Event::Key(key) = event::read()? else {
        return Ok(());
    };

    if key.kind != event::KeyEventKind::Press {
        return Ok(());
    }

    if key.code == KeyCode::Esc {
        app.state = State::ShouldQuit;
    }

    if app.state == State::InputIntention {
        match key.code {
            KeyCode::Enter => app.state = State::InputDuration,
            KeyCode::Tab => app.state = State::InputDuration,
            KeyCode::Backspace => {
                let _ = app.user_intention.pop();
            }
            KeyCode::Char(c) => {
                app.user_intention.push(c);
            }
            _ => {}
        }
    } else if app.state == State::InputDuration {
        match key.code {
            KeyCode::Enter => app.to_in_progress(),
            KeyCode::Tab => app.state = State::InputIntention,
            KeyCode::Backspace => {
                let _ = app.user_duration_str.pop();
            }
            KeyCode::Char(c) if c.is_ascii_digit() => {
                app.user_duration_str.push(c);
            }
            _ => {}
        }
    } else if app.state == State::InProgress {
        match key.code {
            KeyCode::Char('q') => {
                app.state = State::InputIntention;
            }
            KeyCode::Char('a') => {
                app.user_duration = app.user_duration.saturating_add(Duration::from_secs(60));
            }
            KeyCode::Char('x') => {
                app.user_duration = app.user_duration.saturating_sub(Duration::from_secs(60));
            }
            _ => {}
        }
    }

    Ok(())
}

fn update_session_stats(app: &mut App) {
    let window_title = window::get_title_clean().into();
    let delta = app.current_window_time.elapsed();
    if app.current_window_title != window_title || (delta > Duration::from_secs(1)) {
        let entry = app
            .session_stats
            .entry(app.current_window_title.clone())
            .or_insert_with(|| Duration::default());
        *entry += app.current_window_time.elapsed();
        app.current_window_time = Instant::now();
        app.current_window_title = window_title;
    }
}

fn ui(frame: &mut Frame, app: &App) {
    let layout = Layout::vertical([
        Constraint::Min(3),
        Constraint::Min(3),
        Constraint::Percentage(100),
    ]);
    let [layout_intention, layout_duration, layout_titles] = layout.areas(frame.size());

    let border_type_intention = if app.state == State::InputIntention {
        BorderType::Thick
    } else {
        BorderType::Plain
    };
    let input_intention: Vec<Line> = vec![Line::from(Span::raw(&app.user_intention))];
    frame.render_widget(
        Paragraph::new(input_intention).block(
            Block::bordered()
                .border_type(border_type_intention)
                .title("Intention"),
        ),
        layout_intention,
    );

    let input_duration: Vec<Line> = if app.state == State::InProgress {
        let s = duration_as_str(&app.session_remaining);
        vec![Line::from(Span::raw(s))]
    } else {
        vec![Line::from(Span::raw(&app.user_duration_str))]
    };
    let border_type_duration = if app.state == State::InputDuration {
        BorderType::Thick
    } else {
        BorderType::Plain
    };

    frame.render_widget(
        Paragraph::new(input_duration).block(
            Block::bordered()
                .border_type(border_type_duration)
                .title("Duration"),
        ),
        layout_duration,
    );

    let stats = session_stats_to_lines(&app.session_stats);
    frame.render_widget(
        Paragraph::new(stats).block(Block::bordered().title("Session")),
        layout_titles,
    );
}