ai-gui/src/app.rs

use std::{sync::mpsc::{self, Receiver, Sender}, time::Duration};

use chrono::Local;
use eframe::{
    egui::{
        self, CentralPanel, ComboBox, DragValue, Grid, Layout, ScrollArea, Sense, Slider, TextEdit,
        Window,
    },
    emath::Align,
    epaint::Vec2,
};
use egui_notify::Toasts;

use crate::{
    ai::{generate_with_mpsc, Params, Response, ResponseType},
    config::{AI, CONFIG},
    history::{read_all_history, write_history, HistoryEntry},
};

pub struct App {
    toasts: Toasts,

    output: String,
    selected_ai: AI,
    params: Params,

    history_open: bool,

    history: Vec<HistoryEntry>,
    open_history_entry: Option<HistoryEntry>,

    tx: Sender<Response>,
    rx: Receiver<Response>,
}

impl Default for App {
    fn default() -> Self {
        let (tx, rx) = mpsc::channel();

        Self {
            toasts: Default::default(),
            output: "Awaiting generation...".to_string(),
            selected_ai: CONFIG.ai[0].clone(),
            params: Default::default(),
            history_open: false,
            history: vec![],
            open_history_entry: None,
            tx,
            rx,
        }
    }
}

impl eframe::App for App {
    fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
        if let Ok(recieved) = self.rx.try_recv() {
            if recieved.response_type == ResponseType::Success {
                self.output = recieved.output.trim().to_string();
            } else {
                self.output = "Failed to generate!".to_string();
                self.toasts
                    .error("Failed to generate! Error printed to console.");

                eprintln!("{}", recieved.output);
            }

            match write_history(HistoryEntry {
                date: Local::now(),
                selected_ai: self.selected_ai.name.clone(),
                params: self.params.clone(),
                result: self.output.clone(),
            }) {
                Ok(_) => {}
                Err(error) => {
                    self.toasts
                        .error(format!("Failed to write history! {}", error));
                }
            }
        }
        CentralPanel::default().show(ctx, |ui| {
            Window::new("History")
                .open(&mut self.history_open)
                .collapsible(false)
                .hscroll(true)
                .show(ctx, |ui| {
                    if let Some(entry) = self.open_history_entry.clone() {
                        ui.set_max_width(580.0);

                        ui.horizontal(|ui| {
                            // maybe make these param indicators look more like the main ui somehow
                            ui.label(format!("{} - AI", entry.selected_ai));

                            ui.separator();

                            ui.label(format!(
                                "{} - Temperature",
                                entry.params.temperature.to_string()
                            ));

                            ui.separator();

                            ui.label(format!(
                                "{} - Max Tokens",
                                entry.params.max_tokens.to_string()
                            ));

                            ui.with_layout(Layout::right_to_left(Align::Max), |ui| {
                                if ui.button("Back").clicked() {
                                    self.open_history_entry = None;
                                }
                            });
                        });

                        ui.with_layout(Layout::left_to_right(Align::Min), |ui| {
                            ScrollArea::vertical()
                                .id_source("history_prompt_scroll")
                                .show(ui, |ui| {
                                    ui.add(TextEdit::multiline(&mut entry.params.prompt.as_str()));
                                });

                            ScrollArea::vertical()
                                .id_source("history_output_scroll")
                                .show(ui, |ui| {
                                    ui.add(TextEdit::multiline(&mut entry.result.as_str()));
                                });
                        })
                        .response
                        .rect
                        .width();
                    } else {
                        ScrollArea::vertical()
                            .id_source("history_scroll")
                            .show(ui, |ui| {
                                Grid::new("history_grid").num_columns(1).show(
                                    ui,
                                    |ui| {
                                        for entry in &self.history {
                                            let response = ui
                                                .horizontal(|ui| {
                                                    ui.label(&entry.selected_ai);
                                                    ui.label(
                                                        entry
                                                            .date
                                                            .format("%Y-%m-%d %H:%M:%S")
                                                            .to_string(),
                                                    )
                                                })
                                                .response;

                                            let clickable =
                                                ui.allocate_rect(response.rect, Sense::click());

                                            if clickable.clicked() {
                                                self.open_history_entry = Some(entry.clone());
                                            }

                                            ui.end_row()
                                        }
                                    },
                                );
                            });

                    }
                });

            ui.vertical(|ui| {
                if ui.button("Generate!").clicked() {
                    self.output = "Generating...".to_string();

                    generate_with_mpsc(
                        self.selected_ai.clone(),
                        self.params.clone(),
                        self.tx.clone(),
                    )
                }

                ui.horizontal(|ui| {
                    ComboBox::from_label("AI")
                        .selected_text(&self.selected_ai.name)
                        .show_ui(ui, |ui| {
                            for ai in &CONFIG.ai {
                                ui.selectable_value(&mut self.selected_ai, ai.clone(), &ai.name);
                            }
                        });

                    ui.separator();

                    ui.add(
                        Slider::new(&mut self.params.temperature, 0.0..=2.0)
                            .text("Temperature")
                            .show_value(true),
                    );

                    ui.separator();

                    ui.add(DragValue::new(&mut self.params.max_tokens));
                    ui.label("Max Tokens");

                    ui.with_layout(Layout::right_to_left(Align::Max), |ui| {
                        ui.add_visible_ui(CONFIG.app.history_dir.is_some(), |ui| {
                            if ui.button("History").clicked() {
                                // add proper error handling
                                self.history = read_all_history().unwrap();

                                self.history_open = true;
                            }
                        });
                    });
                });

                ui.with_layout(Layout::left_to_right(Align::Min), |ui| {
                    let size = ui.available_size();

                    ScrollArea::vertical()
                        .id_source("prompt_scroll")
                        .show(ui, |ui| {
                            ui.add_sized(
                                Vec2::new(size.x / 2.0, size.y),
                                TextEdit::multiline(&mut self.params.prompt),
                            );
                        });

                    ScrollArea::vertical()
                        .id_source("output_scroll")
                        .show(ui, |ui| {
                            ui.add_sized(
                                Vec2::new(size.x / 2.0, size.y),
                                TextEdit::multiline(&mut self.output.as_str()),
                            );
                        });
                });
            });

            self.toasts.show(ctx);
            ctx.request_repaint_after(Duration::from_secs(1))
        });
    }
}