ttypo 0.1.4

mod config;
mod test;
mod title;
mod ui;

use config::Config;
use test::{DisplayLine, Test, results::Results};

use clap::{CommandFactory, Parser, Subcommand};
use clap_complete::{Shell, generate};
use crossterm::{
    self, cursor,
    event::{self, Event, KeyCode, KeyEvent, KeyEventKind, KeyModifiers},
    execute, terminal,
};
use rand::seq::SliceRandom;
use ratatui::{Terminal, backend::CrosstermBackend};
use rust_embed::RustEmbed;
use std::{
    ffi::OsString,
    fs,
    io::{self, Read},
    num,
    path::PathBuf,
    str,
    time::Duration,
};

#[derive(RustEmbed)]
#[folder = "resources/runtime"]
struct Resources;

#[derive(Debug, Parser)]
#[command(about, version)]
struct Opt {
    /// Read test contents from the specified file, or "-" for stdin
    #[arg(value_name = "PATH")]
    contents: Option<PathBuf>,

    #[arg(short, long)]
    debug: bool,

    /// Specify word count
    #[arg(short, long, value_name = "N", default_value = "50")]
    words: num::NonZeroUsize,

    /// Use config file
    #[arg(short, long, value_name = "PATH")]
    config: Option<PathBuf>,

    /// Specify test language in file
    #[arg(long, value_name = "PATH")]
    language_file: Option<PathBuf>,

    /// Specify test language
    #[arg(short, long, value_name = "LANG")]
    language: Option<String>,

    /// List installed languages
    #[arg(long)]
    list_languages: bool,

    /// Disable backtracking to completed words
    #[arg(long)]
    no_backtrack: bool,

    /// Enable sudden death mode to restart on first error
    #[arg(long)]
    sudden_death: bool,

    /// Disable backspace
    #[arg(long)]
    no_backspace: bool,

    /// Display all but skip non-ASCII characters during typing
    #[arg(long)]
    ascii: bool,

    #[command(subcommand)]
    command: Option<Command>,
}

#[derive(Debug, Subcommand)]
enum Command {
    /// Generate shell completions
    Completions {
        /// Shell to generate completions for
        shell: Shell,
    },
}

impl Opt {
    /// Generate test contents.
    ///
    /// Returns `(words, lines)` where `lines` describes the original file
    /// layout (empty for language/word-list mode).
    fn gen_contents(&self) -> Option<(Vec<String>, Vec<DisplayLine>)> {
        match &self.contents {
            Some(path) => {
                let text = if path.as_os_str() == "-" {
                    let mut buf = String::new();
                    std::io::stdin()
                        .lock()
                        .read_to_string(&mut buf)
                        .expect("Error reading from stdin.");
                    buf
                } else {
                    fs::read_to_string(path).expect("Error reading file.")
                };

                let mut words = Vec::new();
                let mut lines = Vec::new();

                for line in text.lines() {
                    let indent: String = line
                        .chars()
                        .take_while(|c| c.is_whitespace())
                        .collect::<String>()
                        .replace('\t', "    ");

                    let word_start = words.len();
                    for token in line.split_whitespace() {
                        let word: String = token.chars().filter(|c| !c.is_control()).collect();
                        if !word.is_empty() {
                            words.push(word);
                        }
                    }
                    let word_count = words.len() - word_start;

                    lines.push(DisplayLine {
                        indent,
                        word_start,
                        word_count,
                    });
                }

                Some((words, lines))
            }
            None => {
                let lang_name = self
                    .language
                    .clone()
                    .unwrap_or_else(|| self.config().default_language);

                let bytes: Vec<u8> = self
                    .language_file
                    .as_ref()
                    .map(fs::read)
                    .and_then(Result::ok)
                    .or_else(|| fs::read(self.language_dir().join(&lang_name)).ok())
                    .or_else(|| {
                        Resources::get(&format!("language/{}", &lang_name))
                            .map(|f| f.data.into_owned())
                    })?;

                let mut rng = rand::rng();

                let mut language: Vec<&str> = str::from_utf8(&bytes)
                    .expect("Language file had non-utf8 encoding.")
                    .lines()
                    .collect();
                language.shuffle(&mut rng);

                let mut contents: Vec<_> = language
                    .into_iter()
                    .cycle()
                    .take(self.words.get())
                    .map(ToOwned::to_owned)
                    .collect();
                contents.shuffle(&mut rng);

                Some((contents, Vec::new()))
            }
        }
    }

    /// Configuration
    fn config(&self) -> Config {
        fs::read(
            self.config
                .clone()
                .unwrap_or_else(|| self.config_dir().join("config.toml")),
        )
        .map(|bytes| {
            toml::from_str(str::from_utf8(&bytes).unwrap_or_default())
                .expect("Configuration was ill-formed.")
        })
        .unwrap_or_default()
    }

    /// Installed languages under config directory
    fn languages(&self) -> io::Result<impl Iterator<Item = OsString> + use<>> {
        let builtin = Resources::iter().filter_map(|name| {
            name.strip_prefix("language/")
                .map(ToOwned::to_owned)
                .map(OsString::from)
        });

        let configured = self
            .language_dir()
            .read_dir()
            .into_iter()
            .flatten()
            .map_while(Result::ok)
            .map(|e| e.file_name());

        Ok(builtin.chain(configured))
    }

    /// Config directory
    fn config_dir(&self) -> PathBuf {
        dirs::config_dir()
            .expect("Failed to find config directory.")
            .join("ttypo")
    }

    /// Language directory under config directory
    fn language_dir(&self) -> PathBuf {
        self.config_dir().join("language")
    }

    /// Installed languages sorted and deduplicated.
    fn languages_sorted(&self) -> Vec<String> {
        let mut langs: Vec<String> = self
            .languages()
            .ok()
            .into_iter()
            .flatten()
            .filter_map(|os| os.into_string().ok())
            .collect();
        langs.sort();
        langs.dedup();
        langs
    }

    /// Validate that the language used in language mode resolves to a file.
    /// `--language-file` bypasses language-name lookup, so it's always ok.
    fn validate_language(&self, config: &Config) -> Result<(), String> {
        if self.language_file.is_some() {
            return Ok(());
        }
        let lang = self
            .language
            .clone()
            .unwrap_or_else(|| config.default_language.clone());
        let found = self.language_dir().join(&lang).is_file()
            || Resources::get(&format!("language/{}", &lang)).is_some();
        if found { Ok(()) } else { Err(lang) }
    }
}

fn teardown() -> io::Result<()> {
    terminal::disable_raw_mode()?;
    execute!(
        io::stdout(),
        cursor::RestorePosition,
        cursor::Show,
        terminal::LeaveAlternateScreen,
    )?;
    Ok(())
}

enum State {
    Test(Test),
    Results(Results),
}

impl State {
    fn render_into(
        &self,
        terminal: &mut Terminal<CrosstermBackend<io::Stdout>>,
        config: &Config,
    ) -> io::Result<()> {
        match self {
            State::Test(test) => {
                terminal.draw(|f: &mut ratatui::Frame| {
                    f.render_widget(config.theme.apply_to(test), f.area());
                })?;
            }
            State::Results(results) => {
                terminal.draw(|f: &mut ratatui::Frame| {
                    f.render_widget(config.theme.apply_to(results), f.area());
                })?;
            }
        }
        Ok(())
    }
}

fn main() -> io::Result<()> {
    let mut opt = Opt::parse();
    if opt.debug {
        dbg!(&opt);
    }

    let config = opt.config();
    if opt.debug {
        dbg!(&config);
    }

    if let Some(Command::Completions { shell }) = opt.command {
        generate(shell, &mut Opt::command(), "ttypo", &mut io::stdout());
        return Ok(());
    }

    if opt.list_languages {
        opt.languages()
            .unwrap()
            .for_each(|name| println!("{}", name.to_str().expect("Ill-formatted language name.")));

        return Ok(());
    }

    // Validate language up front (language mode only). Fail early before any
    // terminal takeover so error + help render cleanly.
    if opt.contents.is_none()
        && let Err(lang) = opt.validate_language(&config)
    {
        eprintln!("error: language \"{}\" not found.\n", lang);
        let _ = Opt::command().print_help();
        std::process::exit(1);
    }

    let backend = CrosstermBackend::new(io::stdout());
    let mut terminal = Terminal::new(backend)?;

    // File/stdin mode: read contents BEFORE entering the alt screen so stdin
    // still points at the real TTY/pipe.
    let mut file_contents: Option<(Vec<String>, Vec<DisplayLine>)> = if opt.contents.is_some() {
        let r = opt.gen_contents().expect(
            "Couldn't get test contents. Make sure the specified language actually exists.",
        );
        if r.0.is_empty() {
            eprintln!("Error: the provided file or language contains no words to type.");
            eprintln!("If you specified a file, make sure it isn't empty.");
            std::process::exit(1);
        }
        Some(r)
    } else {
        None
    };

    terminal::enable_raw_mode()?;
    execute!(
        io::stdout(),
        cursor::Hide,
        cursor::SavePosition,
        terminal::EnterAlternateScreen,
    )?;
    terminal.clear()?;

    // Outer "session" loop: re-entered when the user hits 'm' on the results
    // screen to return to the main menu. File mode never re-enters since 'm'
    // is disabled there.
    'outer: loop {
        if opt.contents.is_none() {
            let t = title::Title::new(
                opt.language
                    .clone()
                    .unwrap_or_else(|| config.default_language.clone()),
                opt.words,
                opt.sudden_death,
                opt.no_backtrack,
                opt.no_backspace,
                opt.ascii,
                opt.languages_sorted(),
            );
            match title::run(&mut terminal, &config, t)? {
                title::Outcome::Quit => break 'outer,
                title::Outcome::Start(t) => {
                    opt.language = Some(t.language);
                    opt.words = t.words;
                    opt.sudden_death = t.sudden_death;
                    opt.no_backtrack = t.no_backtrack;
                    opt.no_backspace = t.no_backspace;
                    opt.ascii = t.ascii;
                }
            }
        }

        let (contents, lines) = match file_contents.take() {
            Some(fc) => fc,
            None => opt.gen_contents().unwrap_or_else(|| {
                let _ = teardown();
                eprintln!("Couldn't get test contents.");
                std::process::exit(1);
            }),
        };
        if contents.is_empty() {
            let _ = teardown();
            eprintln!("Error: the provided file or language contains no words to type.");
            std::process::exit(1);
        }

        let source = match &opt.contents {
            Some(path) if path.as_os_str() == "-" => "stdin".to_string(),
            Some(path) => path
                .file_name()
                .map(|n| n.to_string_lossy().into_owned())
                .unwrap_or_else(|| path.display().to_string()),
            None => opt
                .language
                .clone()
                .unwrap_or_else(|| config.default_language.clone()),
        };

        let saved_contents: Option<(Vec<String>, Vec<DisplayLine>)> = opt
            .contents
            .is_some()
            .then(|| (contents.clone(), lines.clone()));
        let is_file_mode = saved_contents.is_some();

        let make_test = |contents: Vec<String>, lines: Vec<DisplayLine>, source: String| {
            Test::new(
                contents,
                !opt.no_backtrack,
                opt.sudden_death,
                !opt.no_backspace,
                lines,
                opt.ascii,
                source,
            )
        };

        let restart_contents = || -> (Vec<String>, Vec<DisplayLine>) {
            saved_contents
                .as_ref()
                .map(|(c, l)| (c.clone(), l.clone()))
                .unwrap_or_else(|| {
                    opt.gen_contents().expect(
                        "Couldn't get test contents. Make sure the specified language actually exists.",
                    )
                })
        };

        let mut paused_test: Option<Test> = None;
        let mut state = State::Test(make_test(contents, lines, source.clone()));

        state.render_into(&mut terminal, &config)?;
        'session: loop {
            // Poll with timeout so the status bar (timer/WPM) updates live
            if !event::poll(Duration::from_millis(200))? {
                // Redraw for timer updates
                state.render_into(&mut terminal, &config)?;
                continue;
            }
            let event = event::read()?;

            // handle exit controls
            match event {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    kind: KeyEventKind::Press,
                    modifiers: KeyModifiers::CONTROL,
                    ..
                }) => break 'outer,
                Event::Key(KeyEvent {
                    code: KeyCode::Esc,
                    kind: KeyEventKind::Press,
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    state = match state {
                        State::Test(test) => {
                            let mut results = Results::from(&test);
                            results.is_repeat = is_file_mode;
                            paused_test = Some(test);
                            State::Results(results)
                        }
                        State::Results(_) => break 'outer,
                    };
                }
                _ => {}
            }

            match state {
                State::Test(ref mut test) => {
                    if let Event::Key(key) = event {
                        test.handle_key(key);
                        if test.complete {
                            let mut results = Results::from(&*test);
                            results.is_repeat = is_file_mode;
                            paused_test = None;
                            state = State::Results(results);
                        }
                    }
                }
                State::Results(ref result) => {
                    if let Event::Key(KeyEvent {
                        code: KeyCode::Char(c),
                        kind: KeyEventKind::Press,
                        ..
                    }) = event
                    {
                        match c.to_ascii_lowercase() {
                            'r' => {
                                let (new_contents, new_lines) = restart_contents();
                                if new_contents.is_empty() {
                                    continue;
                                }
                                state =
                                    State::Test(make_test(new_contents, new_lines, source.clone()));
                            }
                            'p' => {
                                if result.missed_words.is_empty() {
                                    continue;
                                }
                                let mut practice_words: Vec<String> = result
                                    .missed_words
                                    .iter()
                                    .flat_map(|(w, _)| std::iter::repeat_n(w.clone(), 5))
                                    .collect();
                                practice_words.shuffle(&mut rand::rng());
                                state = State::Test(make_test(
                                    practice_words,
                                    Vec::new(),
                                    "practice".to_string(),
                                ));
                            }
                            'c' => {
                                if let Some(test) = paused_test.take() {
                                    state = State::Test(test);
                                }
                            }
                            'q' => break 'outer,
                            'm' if !is_file_mode => break 'session,
                            _ => {}
                        }
                    }
                }
            }

            state.render_into(&mut terminal, &config)?;
        }
    }

    teardown()?;

    Ok(())
}

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

    fn make_opt(path: PathBuf, ascii: bool) -> Opt {
        Opt {
            contents: Some(path),
            debug: false,
            words: num::NonZeroUsize::new(50).unwrap(),
            config: None,
            language_file: None,
            language: None,
            list_languages: false,
            no_backtrack: false,
            sudden_death: false,
            no_backspace: false,
            ascii,
            command: None,
        }
    }

    #[test]
    fn gen_contents_empty_file_returns_empty_vec() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("empty.txt");
        fs::File::create(&path).unwrap();

        let (contents, lines) = make_opt(path, false).gen_contents().unwrap();
        assert!(contents.is_empty(), "empty file should produce empty vec");
        assert!(lines.is_empty());
    }

    #[test]
    fn gen_contents_splits_words() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("words.txt");
        let mut f = fs::File::create(&path).unwrap();
        writeln!(f, "hello world rust").unwrap();

        let (contents, lines) = make_opt(path, false).gen_contents().unwrap();
        assert_eq!(contents, vec!["hello", "world", "rust"]);
        assert_eq!(lines.len(), 1);
        assert_eq!(lines[0].word_start, 0);
        assert_eq!(lines[0].word_count, 3);
    }

    #[test]
    fn gen_contents_preserves_unicode() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("unicode.txt");
        let mut f = fs::File::create(&path).unwrap();
        writeln!(f, "hello\u{2014}world \u{201c}quoted\u{201d}").unwrap();

        let (contents, _) = make_opt(path, false).gen_contents().unwrap();
        assert_eq!(
            contents,
            vec!["hello\u{2014}world", "\u{201c}quoted\u{201d}"]
        );
    }

    #[test]
    fn gen_contents_multiline_tracks_lines() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("multi.txt");
        let mut f = fs::File::create(&path).unwrap();
        write!(f, "first line\nsecond line\n\nfourth line").unwrap();

        let (contents, lines) = make_opt(path, false).gen_contents().unwrap();
        assert_eq!(
            contents,
            vec!["first", "line", "second", "line", "fourth", "line"]
        );
        // 4 lines: line 1, line 2, empty line, line 4
        assert_eq!(lines.len(), 4);
        assert_eq!((lines[0].word_start, lines[0].word_count), (0, 2));
        assert_eq!((lines[1].word_start, lines[1].word_count), (2, 2));
        assert_eq!(lines[2].word_count, 0); // empty line preserved
        assert_eq!((lines[3].word_start, lines[3].word_count), (4, 2));
    }

    #[test]
    fn gen_contents_preserves_whitespace_only_lines() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("spaces.txt");
        let mut f = fs::File::create(&path).unwrap();
        write!(f, "hello\n   \n  \t  \nworld").unwrap();

        let (contents, lines) = make_opt(path, false).gen_contents().unwrap();
        assert_eq!(contents, vec!["hello", "world"]);
        // 4 lines total: "hello", whitespace-only, whitespace-only, "world"
        assert_eq!(lines.len(), 4);
        assert_eq!(lines[0].word_count, 1);
        assert_eq!(lines[1].word_count, 0); // whitespace-only preserved
        assert_eq!(lines[2].word_count, 0);
        assert_eq!(lines[3].word_count, 1);
    }

    #[test]
    fn gen_contents_keeps_all_unicode_tokens() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("alluni.txt");
        let mut f = fs::File::create(&path).unwrap();
        write!(f, "hello \u{2014}\u{2014}\u{2014} world").unwrap();

        let (contents, _) = make_opt(path, false).gen_contents().unwrap();
        assert_eq!(contents, vec!["hello", "\u{2014}\u{2014}\u{2014}", "world"]);
    }

    #[test]
    fn gen_contents_preserves_punctuation() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("punct.txt");
        let mut f = fs::File::create(&path).unwrap();
        write!(f, "it's a \"test\" (100%); done!").unwrap();

        let (contents, _) = make_opt(path, false).gen_contents().unwrap();
        assert_eq!(contents, vec!["it's", "a", "\"test\"", "(100%);", "done!"]);
    }

    #[test]
    fn gen_contents_expands_tabs_in_indent() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("tabs.txt");
        let mut f = fs::File::create(&path).unwrap();
        write!(f, "hello\n\tindented\n\t\tdeep").unwrap();

        let (contents, lines) = make_opt(path, false).gen_contents().unwrap();
        assert_eq!(contents, vec!["hello", "indented", "deep"]);
        assert_eq!(lines[0].indent, "");
        assert_eq!(lines[1].indent, "    "); // 1 tab = 4 spaces
        assert_eq!(lines[2].indent, "        "); // 2 tabs = 8 spaces
    }

    #[test]
    fn gen_contents_strips_control_chars_from_words() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("ctrl.txt");
        let mut f = fs::File::create(&path).unwrap();
        write!(f, "hel\x07lo wor\x00ld").unwrap();

        let (contents, _) = make_opt(path, false).gen_contents().unwrap();
        assert_eq!(contents, vec!["hello", "world"]);
    }
}