nyanit/src/main.rs

// <Nyanit - init program fot initrd.>
// Copyright (C) <2026>  <Kira Foundation>

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use ratatui::widgets::Borders;
use ratatui::{Frame, Terminal, backend::TermionBackend, border, widgets::Paragraph};
use ratatui::{
    buffer::Buffer,
    layout::{Constraint, Layout, Rect},
    style::{Style, Stylize},
    symbols::border,
    text::{Line, Text},
    widgets::{Block, List, ListState, Widget},
};
use std::io::{self, Read, stdout};
use std::process::ExitCode;

use termion::event::{Event, Key};
use termion::input::TermRead; // Required for .keys()
use termion::raw::IntoRawMode;

mod nyan_system;

use nyan_system::{ExecCommandError, exec_command};

const BEE_CAT: &str = include_str!("bee_cat.txt");

const ROOT_PARTLABEL: &str = "bee-root";
const ROOT_DEV: &str = "/dev/disk/by-partlabel/bee-root";
const TMP_ROOT: &str = "/mnt/bee_root";
const TMP_ROOT_SNAPS: &str = "/mnt/bee-root/bzz_snaps";
const SYS_INIT: &str = "/sbin/init";

// fn bee_root_is_present() -> bool {
//    nyan_system::list_gpt_labeled_devs().iter().filter(|s| s.ends_with(ROOT_PARTLABEL)).count() > 0
// }

// fn main() -> ExitCode {
//     println!("Looking for {}", ROOT_PARTLABEL);
//     if bee_root_is_present() {
//         println!("{} found!", ROOT_PARTLABEL);
//     }
//     else {
//         print!("Cannot find {}!", ROOT_PARTLABEL);
//         print!("Exiting...");
//         return ExitCode::FAILURE;
//     }
//     // 1 get labels and see if there is bzzpss-root in them

//     ExitCode::SUCCESS
// }

#[derive(Debug, PartialEq)]
enum NyanitState {
    Init,
    WaitForF2(Vec<String>),
    ChooseSnapshot(Vec<String>),
    EnterPassword,
    Error(String),
    HandOff(String),
}

pub fn init_stage() -> Result<NyanitState, ExecCommandError> {
    // mount sysfss
    nyan_system::prepare_env()?;

    // check if device is encrypted
    if nyan_system::bzzpss_dev_encrypted(ROOT_DEV)? {
        // if encrypted return EnterPassword so used get promted for pass
        return Ok(NyanitState::EnterPassword);
    }

    // if unencrypted
    // mount bee-root
    nyan_system::mount_bee_root(TMP_ROOT, ROOT_DEV)?;

    // list snaps if any
    let snaps_list = nyan_system::list_bzzpss_snaps(TMP_ROOT_SNAPS)?;
    // if any snaps detected whait for user to press F2
    if snaps_list.len() > 0 {
        return Ok(NyanitState::WaitForF2(snaps_list));
    }

    // tell to finish boot
    Ok(NyanitState::HandOff(TMP_ROOT.into()))
}

pub struct NyanitTUI {
    counter: u32,
    exit: bool,
    state: NyanitState,
    list_state: ListState,
}

impl NyanitTUI {
    pub fn new() -> io::Result<Self> {
        Ok(Self {
            counter: 0,
            exit: false,
            state: NyanitState::Init,
            list_state: ListState::default().with_selected(Some(0)),
        })
    }

    // pub fn bzz_unwrap<T,E>(&mut self, val: Result<T,E>) -> T {
    //     match val {
    //         Ok(v) => v,
    //     }
    // }
    /// runs the application's main loop until the user quits
    pub fn run(&mut self) -> io::Result<()> {
        self.state = match init_stage() {
            Ok(st) => st,
            Err(ex) => NyanitState::Error(format!("Error while trying to bzzz-pss things: {}", ex)),
        };

        // if drive unencrypted and there is no root snapshots finish boot
        if let NyanitState::HandOff(root_mnt) = &self.state {
            if let Err(ex) = nyan_system::hand_off_control(root_mnt, SYS_INIT) {
                self.state = NyanitState::Error(ex.to_string())
            }
        }

        // 1. Enable raw mode
        let stdout = stdout().into_raw_mode()?;
        // 2. Setup Terminal
        let backend = TermionBackend::new(stdout);
        let mut terminal = Terminal::new(backend)?;

        // init tty input and start actually drawing things now!
        let mut tty_input: io::Bytes<std::fs::File> = termion::get_tty()?.bytes();
        while !self.exit {
            terminal.draw(|frame| self.draw(frame))?;
            self.handle_events(&mut tty_input)?;
        }
        Ok(())
    }

    fn draw(&mut self, frame: &mut Frame) {
        let area = frame.area();
        // 1. Create the outer block with borders
        let outer_block = Block::bordered()
            .title(Line::from(" BZZ-PSS ").centered())
            .light_yellow();
        // 2. Calculate the inner area excluding the borders
        let inner_area = outer_block.inner(area);
        // 3. Render the border block itself
        frame.render_widget(outer_block, area);

        // 1. Create vertical layout: Top (flexible), Center (fixed height), Bottom (flexible)
        let vertical_chunks = Layout::vertical([
            Constraint::Length(7),      // logo top part
            Constraint::Percentage(80), // Height of the paragraph block
            Constraint::Length(3),      // footer
        ])
        .split(inner_area);

        // 2. Create horizontal layout inside the center chunk: Left (flexible), Center (fixed width), Right (flexible)
        let horizontal_chunks = Layout::horizontal([
            Constraint::Fill(2),
            Constraint::Length(11),
            Constraint::Fill(2),
        ])
        .split(vertical_chunks[0]);

        // 3. Render the logo in the upper chunk
        let paragraph = Paragraph::new(BEE_CAT).left_aligned().light_yellow();
        frame.render_widget(paragraph, horizontal_chunks[1]);

        // now check app state
        match &self.state {
            NyanitState::Error(ex) => {
                // render big red wolf in the center with err messaghe
                let err_paragraph = Paragraph::new(ex.as_str()).centered().light_red().bold();
                frame.render_widget(err_paragraph, vertical_chunks[1]);
            }
            NyanitState::WaitForF2(_) => {
                // show press F2 for options text
                let f2_paragraph = Paragraph::new("Press F2 for options")
                    .centered()
                    .magenta()
                    .bold();
                frame.render_widget(f2_paragraph, vertical_chunks[1]);
            }
            NyanitState::ChooseSnapshot(snaps_list) => {
                // render menu for snap choosing
                frame.render_stateful_widget(
                    List::new(snaps_list.iter().map(String::as_str))
                        .block(Block::default().borders(Borders::ALL))
                        .highlight_symbol(">>")
                        .highlight_style(Style::new().bold()),
                    vertical_chunks[1],
                    &mut self.list_state,
                );
            }
            _ => (),
        };
    }

    fn handle_events(&mut self, tty_input: &mut io::Bytes<std::fs::File>) -> io::Result<()> {
        if let Some(item) = tty_input.next() {
            let read_byte = item?;
            let e = termion::event::parse_event(read_byte, tty_input)?;
            match e {
                Event::Key(Key::Up) => {
                    self.counter += 1;
                }
                Event::Key(Key::Down) => {
                    self.counter = self.counter.saturating_sub(1u32);
                }
                _ => (),
            }
        }
        Ok(())
    }
}

fn main() -> ExitCode {
    let mut app = NyanitTUI::new().unwrap();
    app.run().unwrap();
    // Terminal restores automatically when `stdout` (RawTerminal) is dropped

    ExitCode::SUCCESS
}