use base64::{engine::general_purpose, Engine as _};
use ircparser;
use pgp::crypto::sym::SymmetricKeyAlgorithm;
use pgp::ser::Serialize;
use pgp::{Deserializable, Message, SignedPublicKey, SignedSecretKey};
use rand::prelude::*;
use std::collections::HashMap;
use std::net::{Shutdown, TcpStream};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;
use std::{env, fs};

mod listener_server;
mod writer_client;

macro_rules! try_recv {
    ($rx:ident) => {
        match $rx.try_recv() {
            Ok(recieved) => recieved,
            Err(_e) => match _e {
                mpsc::TryRecvError::Empty => String::new(),
                mpsc::TryRecvError::Disconnected => {
                    panic!("Error!")
                }
            },
        }
    };
}

fn bytes_to_privmsg_base64(message: Vec<u8>, reciever: &str) -> String {
    let mut command = String::new();
    for line in general_purpose::STANDARD
        .encode(message)
        .chars()
        .collect::<Vec<char>>()
        .chunks(50)
        .map(|c| c.iter().collect::<String>())
    {
        command.push_str(&format!("PRIVMSG {reciever} :{line}\r\n"));
    }

    println!("{}", command);
    command
}

fn encrypt(key: &SignedPublicKey, message: &str) -> Result<Vec<u8>, pgp::errors::Error> {
    let message = Message::new_literal("none", message);
    let mut rng = StdRng::from_entropy();

    let message = message.encrypt_to_keys(&mut rng, SymmetricKeyAlgorithm::AES128, &[key])?;

    Ok(message.to_bytes()?)
}

fn decrypt<'a>(
    key: &'a SignedSecretKey,
    message: Vec<u8>,
    password: &'a str,
) -> Result<String, pgp::errors::Error> {
    let message = Message::from_bytes(message.as_slice())?; // Convert message bytes to message object
    let message = message.decrypt(|| password.to_string(), &[key])?.0; // Decrypt
    let message = message.map(|x| x.unwrap()).collect::<Vec<Message>>(); // Get all messages
    let message = &message[0]; // Get first message
    let message = message.get_content()?.unwrap_or(Vec::new()); // Get message content as Vec<u8>
    let message = String::from_utf8(message).unwrap(); // Convert to String

    Ok(message)
}

fn get_nick(userstring: &str) -> String {
    let userstring = userstring.chars().collect::<Vec<char>>();
    let start_pos = userstring.iter().position(|&x| x == ':').unwrap() + 1;
    let end_pos = userstring.iter().position(|&x| x == '!').unwrap();
    userstring[start_pos..end_pos].iter().collect::<String>()
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut args = env::args().collect::<Vec<String>>();

    let server = args[1].clone();
    let port = args[2].clone();

    let default_password = String::new();
    let passwd = args.pop().unwrap_or(default_password);

    let stream = TcpStream::connect(format!("{server}:6697"))?;

    let public_key = fs::read("public.gpg")?;
    let secret_key = SignedSecretKey::from_bytes(fs::read("secret.gpg")?.as_slice())?;

    let reader_stream = match stream.try_clone() {
        Ok(stream) => stream,
        Err(_error) => {
            let _ = stream.shutdown(Shutdown::Both);
            panic!("Failed to create the reader stream")
        }
    };

    let writer_stream = match stream.try_clone() {
        Ok(stream) => stream,
        Err(_error) => {
            let _ = stream.shutdown(Shutdown::Both);
            let _ = reader_stream.shutdown(Shutdown::Both);
            panic!("Failed to create the writer stream")
        }
    };

    let (listener_channel_send_tx, listener_channel_rx) = mpsc::channel();
    let (listener_channel_tx, listener_channel_recv_rx) = mpsc::channel();

    let (writer_channel_tx, writer_channel_send_rx) = mpsc::channel();
    let (writer_channel_recv_tx, writer_channel_rx) = mpsc::channel();

    thread::spawn(move || {
        listener_server::listen_to_client(listener_channel_send_tx, listener_channel_recv_rx, &port)
    });
    let tmp_server = server.clone();
    thread::spawn(|| {
        writer_client::write_to_server(
            writer_stream,
            tmp_server,
            writer_channel_send_rx,
            writer_channel_recv_tx,
        )
    });

    let mut keys: HashMap<String, SignedPublicKey> = HashMap::new();
    let mut userstring = String::new();

    loop {
        let recieved = try_recv!(listener_channel_rx);

        if !recieved.is_empty() {
            let command = &ircparser::parse(&recieved).expect("Got an invalid IRC instruction")[0];

            if command.command == "PRIVMSG" && !command.params[0].starts_with("#") {
                let reciever = &command.params[0];
                if !keys.contains_key(reciever) {
                    writer_channel_tx.send(format!("PRIVMSG {reciever} :START_KEY\r\n"))?;

                    writer_channel_tx
                        .send(bytes_to_privmsg_base64(public_key.clone(), reciever))?;
                    writer_channel_tx.send(format!("PRIVMSG {reciever} :END_KEY\r\n"))?;

                    let mut foreign_key: Vec<String> = Vec::new();
                    while !foreign_key.contains(&"END_KEY".to_string()) {
                        let recieved_raw = writer_channel_rx.recv()?;

                        let recieved = &ircparser::parse(&recieved_raw).unwrap()[0];
                        let begin_source_reciever = format!(":{reciever}!");
                        if recieved.command != "PRIVMSG"
                            || !recieved
                                .source
                                .clone()
                                .unwrap_or("".to_string())
                                .starts_with(&begin_source_reciever)
                        {
                            listener_channel_tx.send(recieved_raw)?;
                            continue;
                        }

                        foreign_key.push(recieved.params[1].clone());
                    }
                    foreign_key.pop();

                    println!("Got a foreign key from {reciever}");
                    let foreign_key = foreign_key.concat();
                    dbg!(&foreign_key);

                    let foreign_key = general_purpose::STANDARD.decode(foreign_key)?;

                    dbg!("It's decoded");

                    let key = SignedPublicKey::from_bytes(foreign_key.as_slice())?;

                    dbg!("And I now got it's SignedPublicKey format");

                    keys.insert(reciever.to_string(), key);
                }

                let foreign_key = keys.get(reciever).unwrap();

                writer_channel_tx.send(format!("PRIVMSG {reciever} START_MESSAGE\r\n"))?;
                writer_channel_tx.send(bytes_to_privmsg_base64(
                    encrypt(&foreign_key, &command.params[1])?,
                    reciever,
                ))?;
                writer_channel_tx.send(format!("PRIVMSG {reciever} END_MESSAGE\r\n"))?;
            } else {
                writer_channel_tx.send(recieved.replace("127.0.0.1", &server))?;
            }
        }

        let recieved = try_recv!(writer_channel_rx);
        if !recieved.is_empty() {
            let recieved_parsed = &ircparser::parse(&recieved).unwrap()[0];

            if recieved_parsed.command != "PRIVMSG"
                || recieved_parsed
                    .params
                    .get(0)
                    .unwrap_or(&String::new())
                    .starts_with("#")
            {
                listener_channel_tx.send(recieved.replace(&server, "127.0.0.1"))?;
                continue;
            }
            
            dbg!(&recieved_parsed);

            if recieved_parsed.params[1] == "START_MESSAGE" {
                let reciever = get_nick(&recieved_parsed.source.clone().unwrap());
                let sample_privmsg = recieved.clone();

                let mut encrypted_message: Vec<String> = Vec::new();
                while !encrypted_message.contains(&"END_MESSAGE".to_string()) {
                    let recieved_raw = writer_channel_rx.recv()?;
                    let recieved = &ircparser::parse(&recieved_raw).unwrap()[0];

                    let begin_source_reciever = format!(":{reciever}!");
                    if recieved.command != "PRIVMSG"
                        || !recieved
                            .source
                            .clone()
                            .unwrap_or("".to_string())
                            .starts_with(&begin_source_reciever)
                    {
                        listener_channel_tx.send(recieved_raw.replace(&server, "127.0.0.1"))?;
                        continue;
                    }

                    encrypted_message.push(recieved.params[1].clone());
                }
                encrypted_message.pop();

                let encrypted_message = encrypted_message.join("");
                let encrypted_message = general_purpose::STANDARD.decode(encrypted_message)?;

                let decrypted_message = decrypt(&secret_key, encrypted_message, &passwd)?;
                listener_channel_tx
                    .send(sample_privmsg.replace("START_MESSAGE", &decrypted_message))?;
            } else if recieved_parsed.params[1] == "START_KEY" {
                let reciever = get_nick(&recieved_parsed.source.clone().unwrap());
                let to_send = bytes_to_privmsg_base64(public_key.clone(), &reciever);
                writer_channel_tx.send(to_send)?;
                writer_channel_tx.send(format!("PRIVMSG {reciever} END_KEY\r\n"))?;

                dbg!("Started Getting a key");
                let reciever = get_nick(&recieved_parsed.source.clone().unwrap());
                let mut foreign_key: Vec<String> = Vec::new();
                while !foreign_key.contains(&"END_KEY".to_string()) {
                    let recieved_raw = writer_channel_rx.recv()?;

                    let recieved = &ircparser::parse(&recieved_raw).unwrap()[0];
                    let begin_source_reciever = format!(":{reciever}!");
                    if recieved.command != "PRIVMSG"
                        || !recieved
                            .source
                            .clone()
                            .unwrap_or("".to_string())
                            .starts_with(&begin_source_reciever)
                    {
                        listener_channel_tx.send(recieved_raw.replace(&server, "127.0.0.1"))?;
                        continue;
                    }

                    foreign_key.push(recieved.params[1].clone());
                }
                foreign_key.pop();
                dbg!("Got a foreign key from");
                dbg!(&reciever);

                let foreign_key = foreign_key.concat();

                let foreign_key = general_purpose::STANDARD.decode(foreign_key)?;
                dbg!("Decoded the key");

                let key = SignedPublicKey::from_bytes(foreign_key.as_slice())?;
                dbg!("Deserialized the key");

                keys.insert(reciever.to_string(), key);
            }
        }
    }
}