use crate::algo;
use crate::mpn::Number;
use std::collections::HashMap;

#[derive(Debug)]
pub struct OpenKey {
    pub y: Number,
    pub g: Number,
    pub p: Number,
}

#[derive(Debug)]
pub struct SecretKey {
    pub x: Number,
}

fn factorization(x: &Number) -> Vec<Number> {
    let mut factors = Vec::new();
    let mut n = x.clone();
    let mut i: Number = 2.into();
    let mut flag = true;
    while i.clone() * i.clone() <= n.clone() && flag {
        if n.clone() % i.clone() == 0.into() {
            factors.push(i.clone());
            while n.clone() % i.clone() == 0.into() {
                n = n.clone() / i.clone();
                if algo::rabin_miller_test(&n, 10) {
                    flag = false;
                    break;
                }
            }
        }
        i = i.clone() + 1.into();
    }
    if n.clone() > 1.into() {
        factors.push(n.clone());
    }
    return factors;
}

pub fn is_primitive_root(x: &Number, p: &Number, powers: &Vec<Number>) -> bool {
    if !algo::rabin_miller_test(p, 10) {
        todo!("примитивный корень для составных чисел");
    }
    for l in powers {
        if x.pow_mod(&l, p).unwrap() == 1.into() {
            return false;
        }
    }
    return true;
}

pub fn generate_keypair(p: &Number) -> (OpenKey, SecretKey) {
    let p1 = p.clone() - 1.into();
    let powers = factorization(&p1);
    let mut g;
    loop {
        g = p.random_below();
        if gcd(&g, &p1) == 1.into() {
            continue;
        }
        if g.pow_mod(&(p1.clone() / 2.into()), p).unwrap() != 1.into() {
            continue;
        }
        if is_primitive_root(&g, &p, &powers) {
            break;
        }
    }

    let x = g.random_below();
    let y = g.pow_mod(&x, &p).unwrap();

    println!("Найден примитивный корень: {g}");
    println!("x = {x}");
    println!("y = {y}");

    return (OpenKey { y, g, p: p.clone() }, SecretKey { x });
}

pub const ALPHABET: &str =
    "абвгдеёжзиклмнопрстуфхцчшщъыьэюя1234567890!\"'$%#,.;:{}[]<>*\n\t";

fn get_encoding() -> HashMap<char, String> {
    let mut map = HashMap::new();
    let mut enc = 111;
    for ch in ALPHABET.chars() {
        let mut txt = enc.to_string();
        while txt.contains('0') {
            enc += 1;
            txt = enc.to_string();
        }
        map.insert(ch, txt);
        enc += 1;
    }
    return map;
}

fn get_inv_encoding() -> HashMap<String, char> {
    let encoding = get_encoding();
    let mut inv = HashMap::new();
    for (k, v) in &encoding {
        inv.insert(v.clone(), k.clone());
    }
    return inv;
}

fn gcd(a: &Number, b: &Number) -> Number {
    if *b == 0.into() {
        return a.clone();
    } else {
        return gcd(&b.clone(), &(a.clone() % b.clone()));
    }
}

fn session_key(p: &Number) -> Number {
    let p1 = p.clone() - 1.into();
    loop {
        let k = p1.random_below();
        if gcd(&k, &p1) == 1.into() {
            return k;
        }
    }
}

pub fn filter_text(text: &String) -> String {
    let mut chars = String::new();
    for ch in text.to_lowercase().chars() {
        if ALPHABET.contains(ch) {
            chars.push(ch);
        }
    }
    return chars;
}

pub fn encrypt_text(text: &String, open: OpenKey) -> Vec<(Number, Number)> {
    let encoding = get_encoding();
    println!("Кодировка: {encoding:?}");
    let mut ciphertext = String::new();
    for ch in text.chars() {
        ciphertext.extend(encoding[&ch].chars());
    }

    println!("Перекодированный текст: {ciphertext}");

    let mut result = Vec::new();
    let ciphertext = ciphertext.chars().collect::<Vec<char>>();
    let mut i = 0;
    loop {
        let mut num = String::new();
        while i < ciphertext.len()
            && (num.len() == 0 || Number::parse(&num, 10).unwrap() < open.p)
        {
            num.push(ciphertext[i]);
            i += 1;
        }
        let num2;
        if i == ciphertext.len() {
            let nn = Number::parse(&num, 10).unwrap();
            if nn >= open.p {
                i -= 1;
                num2 = Number::parse(&num[..num.len() - 1], 10).unwrap();
            } else {
                num2 = nn.clone();
            }
        } else {
            i -= 1;
            num2 = Number::parse(&num[..num.len() - 1], 10).unwrap();
        }
        let num = num2;
        println!("Блок: {num}");
        let k = session_key(&open.p);
        let a = open.g.pow_mod(&k, &open.p).unwrap();
        let b = (open.y.pow_mod(&k, &open.p).unwrap() * num.clone())
            % open.p.clone();
        result.push((a, b));
        if i == ciphertext.len() {
            break;
        }
    }
    return result;
}

pub fn decrypt_text(cipher: &String, key: SecretKey, open: OpenKey) -> String {
    let mut encoded = String::new();
    for elem in cipher.split(":") {
        let ab = elem.split(",").collect::<Vec<&str>>();
        let a = Number::parse(ab[0], 10).unwrap();
        let b = Number::parse(ab[1], 10).unwrap();

        println!("Считанные a и b: {a} {b}");
        let pow = open.p.clone() - 1.into() - key.x.clone();
        let m =
            (b.clone() * a.pow_mod(&pow, &open.p).unwrap()) % open.p.clone();
        println!("Расшифровка: {m}");

        encoded.extend(format!("{m}").chars());
    }

    println!("Закодированный текст: {encoded}");
    let inv_enc = get_inv_encoding();
    println!("Обратная кодировка: {inv_enc:?}");
    let encoded = encoded.chars().into_iter().collect::<Vec<char>>();
    let mut result = String::new();
    let mut i = 0;
    loop {
        let mut num = String::new();
        while i < encoded.len() && !inv_enc.contains_key(&num) {
            num.push(encoded[i]);
            i += 1;
        }
        result.push(inv_enc[&num]);
        if i == encoded.len() {
            break;
        }
    }

    return result;
}