use inquire::{Select, Text};
use std::fs::{self, File};
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use crate::elgamal::{decrypt_text, encrypt_text, OpenKey, SecretKey};
use crate::mpn::Number;
mod algo;
mod elgamal;
mod mpn;
mod program;

fn files_recursive(dir: &Path) -> io::Result<Vec<PathBuf>> {
    let mut dir_front = Vec::new();
    let mut files = Vec::new();
    dir_front.push(dir.to_owned());
    while dir_front.len() != 0 {
        let dir = dir_front.pop().unwrap();
        for entry in fs::read_dir(dir)? {
            let path = entry?.path();
            if path.is_dir() {
                dir_front.push(path.clone());
            } else {
                files.push(path.clone());
            }
        }
    }
    return Ok(files);
}

fn dirs_recursive(dir: &Path) -> io::Result<Vec<PathBuf>> {
    let mut dir_front = Vec::new();
    let mut dirs = Vec::new();
    dir_front.push(dir.to_owned());
    dirs.push(dir.to_owned());
    while dir_front.len() != 0 {
        let dir = dir_front.pop().unwrap();
        for entry in fs::read_dir(dir)? {
            let path = entry?.path();
            if path.is_dir() {
                dir_front.push(path.clone());
                dirs.push(path.clone());
            }
        }
    }
    return Ok(dirs);
}

fn write_keys(
    open: OpenKey,
    open_path: String,
    secret: SecretKey,
    secret_path: String,
) {
    let mut open_file = File::create(open_path).unwrap();
    let key = format!("{} {} {}", open.y, open.g, open.p);
    open_file.write_all(key.as_bytes()).unwrap();

    let mut secret_file = File::create(secret_path).unwrap();
    let key = format!("{}", secret.x);
    secret_file.write_all(key.as_bytes()).unwrap();
}

fn read_open_key(path: String) -> OpenKey {
    let open_text = fs::read_to_string(path).unwrap();
    let numbers = open_text.split(" ").collect::<Vec<&str>>();
    let y = Number::parse(numbers[0], 10).unwrap();
    let g = Number::parse(numbers[1], 10).unwrap();
    let p = Number::parse(numbers[2], 10).unwrap();
    return OpenKey { y, g, p };
}

fn read_secret_key(path: String) -> SecretKey {
    let secret_text = fs::read_to_string(path).unwrap();
    let x = Number::parse(&secret_text, 10).unwrap();
    return SecretKey { x };
}

fn main() {
    let operation = Select::new(
        "Выберите тип операции:",
        vec!["Генерация ключей", "Зашифрование", "Расшифрование"],
    )
    .prompt()
    .unwrap();

    if operation == "Генерация ключей" {
        let dir_entries = dirs_recursive(Path::new("."))
            .unwrap()
            .into_iter()
            .map(|e| e.to_str().unwrap().to_owned())
            .collect::<Vec<String>>();
        let result_path =
            Select::new("Выберите куда сохранить результат:", dir_entries)
                .prompt()
                .unwrap();
        let open_name = Text::new("Имя для файла с открытым ключом:")
            .prompt()
            .unwrap();
        let secret_name = Text::new("Имя для файла с закрытым ключом:")
            .prompt()
            .unwrap();
        let p_text = Text::new("Введите простое p:").prompt().unwrap();
        let p = Number::parse(&p_text, 10).unwrap();
        if !algo::rabin_miller_test(&p, 10) {
            println!("Число p должно быть простым");
            return;
        }
        let (op, cl) = elgamal::generate_keypair(&p);
        let open_path = Path::new(&result_path)
            .join(&open_name)
            .to_string_lossy()
            .to_string();
        let secret_path = Path::new(&result_path)
            .join(&secret_name)
            .to_string_lossy()
            .to_string();
        write_keys(op, open_path, cl, secret_path);
    }

    if operation == "Зашифрование" {
        let entries = files_recursive(Path::new("."))
            .unwrap()
            .into_iter()
            .map(|e| e.to_str().unwrap().to_owned())
            .collect::<Vec<String>>();
        let open_name =
            Select::new("Выберите файл с открытым ключом:", entries.clone())
                .prompt()
                .unwrap();
        let plain_file =
            Select::new("Выберите файл для зашифрования:", entries)
                .prompt()
                .unwrap();

        let dir_entries = dirs_recursive(Path::new("."))
            .unwrap()
            .into_iter()
            .map(|e| e.to_str().unwrap().to_owned())
            .collect::<Vec<String>>();
        let result_path =
            Select::new("Выберите куда сохранить результат:", dir_entries)
                .prompt()
                .unwrap();

        let cipher_name =
            Text::new("Введите название для файла с шифротекстом:")
                .prompt()
                .unwrap();

        let open_key = read_open_key(open_name);
        let plaintext = fs::read_to_string(plain_file).unwrap();
        println!("{plaintext}");
        println!("Отфильтрованный текст:");
        let plaintext = elgamal::filter_text(&plaintext);
        println!("{plaintext}");
        let cipher = encrypt_text(&plaintext, open_key);

        let mut elements = Vec::new();
        for (a, b) in &cipher {
            let elem = format!("{},{}", a, b);
            elements.push(elem);
        }
        let ciphertext = elements.join(":");
        println!("Полученный шифротекст: {ciphertext}");
        let cipher_path = Path::new(&result_path)
            .join(&cipher_name)
            .to_string_lossy()
            .to_string();
        fs::write(cipher_path, ciphertext.as_bytes())
            .expect("Не получилось записать шифротекст");
    }

    if operation == "Расшифрование" {
        let entries = files_recursive(Path::new("."))
            .unwrap()
            .into_iter()
            .map(|e| e.to_str().unwrap().to_owned())
            .collect::<Vec<String>>();
        let secret_name =
            Select::new("Выберите файл с закрытым ключом:", entries.clone())
                .prompt()
                .unwrap();
        let open_name =
            Select::new("Выберите файл с открытым ключом:", entries.clone())
                .prompt()
                .unwrap();
        let cipher_file =
            Select::new("Выберите файл для расшифрования:", entries)
                .prompt()
                .unwrap();

        let dir_entries = dirs_recursive(Path::new("."))
            .unwrap()
            .into_iter()
            .map(|e| e.to_str().unwrap().to_owned())
            .collect::<Vec<String>>();
        let result_path =
            Select::new("Выберите куда сохранить результат:", dir_entries)
                .prompt()
                .unwrap();
        let plain_name =
            Text::new("Введите название для файла с расшифровкой:")
                .prompt()
                .unwrap();

        let open = read_open_key(open_name);
        let secret = read_secret_key(secret_name);
        let ciphertext = fs::read_to_string(cipher_file).unwrap();
        let plaintext = decrypt_text(&ciphertext, secret, open);
        println!("Расшифрованный текст:");
        println!("{plaintext}");

        let plain_path = Path::new(&result_path)
            .join(&plain_name)
            .to_string_lossy()
            .to_string();
        fs::write(plain_path, plaintext.as_bytes())
            .expect("Не получилось записать расшифровку");
    }
}