use core::fmt::Display;
use rand::rngs::OsRng;
use rand::{CryptoRng, Rng, RngCore};
use std::collections::HashMap;
use std::ops::{BitAnd, BitOr, BitXor};
use tiny_keccak::{Hasher, Shake};

fn encrypt(key_a: [u8; 16], key_b: [u8; 16], out_key: [u8; 16]) -> [u8; 32] {
    let mut shake = Shake::v256();
    shake.update(&key_a);
    shake.update(&key_b);
    let mut enc_out = [0; 32];
    shake.finalize(&mut enc_out);
    for i in 0..16 {
        enc_out[i + 16] ^= out_key[i];
    }
    enc_out
}

fn decrypt(key_a: [u8; 16], key_b: [u8; 16], encrypted: [u8; 32]) -> Option<[u8; 16]> {
    let mut shake = Shake::v256();
    shake.update(&key_a);
    shake.update(&key_b);
    let mut dec_out = [0; 32];
    shake.finalize(&mut dec_out);
    if dec_out[..16] != encrypted[..16] {
        return None;
    }
    for i in 0..32 {
        dec_out[i] ^= encrypted[i];
    }
    let mut key = [0; 16];
    key[..16].copy_from_slice(&dec_out[16..(16 + 16)]);
    Some(key)
}

#[derive(Clone, Debug)]
struct Wire {
    sym: Option<String>,
    keys: [[u8; 16]; 2],
    evaluator: bool,
}

#[derive(Clone, Debug)]
struct Gate {
    left: Box<Circuit>,
    right: Box<Circuit>,
    out: Wire,
    table: [[[u8; 32]; 2]; 2],
    sym: String,
}

#[derive(Clone, Debug)]
enum Circuit {
    Gate(Gate),
    Wire(Wire),
}

#[derive(Clone, Debug)]
struct HiddenWire {
    sym: Option<String>,
}

#[derive(Clone, Debug)]
struct HiddenGate {
    left: Box<HiddenCircuit>,
    right: Box<HiddenCircuit>,
    table: [[[u8; 32]; 2]; 2],
    sym: String,
}

#[derive(Clone, Debug)]
enum HiddenCircuit {
    Gate(HiddenGate),
    Wire(HiddenWire),
}

#[derive(Clone, Debug)]
struct FullyHiddenCircuit {
    table: [[u8; 32]; 2],
    circuit: HiddenCircuit,
}

impl Display for Wire {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.sym {
            Some(sym) => write!(f, "{}", sym),
            None => write!(f, ""),
        }
    }
}

impl Display for Circuit {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Circuit::Wire(w) => write!(f, "{}", w),
            Circuit::Gate(c) => write!(f, "{}", c),
        }
    }
}

impl Display for Gate {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} {} {}", self.left, self.sym, self.right)
    }
}

impl Wire {
    fn new_rng(sym: String, evaluator: bool, rng: &mut (impl RngCore + CryptoRng)) -> Self {
        let mut keys = [[0; 16]; 2];
        rng.fill_bytes(&mut keys[0]);
        rng.fill_bytes(&mut keys[1]);
        Self {
            sym: Some(sym),
            keys,
            evaluator,
        }
    }

    fn new(sym: String, evaluator: bool) -> Self {
        Self::new_rng(sym, evaluator, &mut OsRng)
    }

    fn to_circ(&self) -> Circuit {
        Circuit::Wire(self.clone())
    }

    fn empty(rng: &mut (impl RngCore + CryptoRng)) -> Self {
        let mut keys = [[0; 16]; 2];
        rng.fill_bytes(&mut keys[0]);
        rng.fill_bytes(&mut keys[1]);
        Self {
            sym: None,
            keys,
            evaluator: false,
        }
    }
}

impl Gate {
    fn new(
        left: Circuit,
        right: Circuit,
        value: [[bool; 2]; 2],
        sym: String,
        rng: &mut (impl RngCore + CryptoRng),
    ) -> Self {
        let linp_keys = match left {
            Circuit::Wire(ref w) => w.keys,
            Circuit::Gate(ref c) => c.out.keys,
        };

        let rinp_keys = match right {
            Circuit::Wire(ref w) => w.keys,
            Circuit::Gate(ref c) => c.out.keys,
        };

        let out = Wire::empty(rng);

        let mut table = [[[0; 32]; 2]; 2];

        let i = rng.gen::<bool>() as usize;
        let j = rng.gen::<bool>() as usize;
        for left_bit in [i, 1 - i] {
            for right_bit in [j, 1 - j] {
                let enc = encrypt(
                    linp_keys[left_bit],
                    rinp_keys[right_bit],
                    out.keys[value[left_bit][right_bit] as usize],
                );

                table[left_bit][right_bit] = enc;
            }
        }

        Self {
            left: Box::new(left),
            right: Box::new(right),
            table,
            sym,
            out,
        }
    }

    fn to_circ(&self) -> Circuit {
        Circuit::Gate(self.clone())
    }
}

impl Circuit {
    fn evaluate_re(&self, vals: HashMap<String, bool>) -> [u8; 16] {
        match self {
            Circuit::Gate(c) => {
                let lv = c.left.evaluate_re(vals.clone());
                let rv = c.right.evaluate_re(vals);
                let mut out = [0; 16];
                for i in [0, 1] {
                    for j in [0, 1] {
                        if let Some(x) = decrypt(lv, rv, c.table[i][j]) {
                            out = x;
                        }
                    }
                }
                out
            }
            Circuit::Wire(w) => match &w.sym {
                // No OT
                Some(c) => w.keys[(vals[c]) as usize],
                None => panic!("This should not happen."),
            },
        }
    }

    pub fn evaluate(&self, vals: HashMap<String, bool>) -> bool {
        let out_label = self.evaluate_re(vals);

        match self {
            Circuit::Gate(c) => {
                if out_label == c.out.keys[0] {
                    false
                } else if out_label == c.out.keys[1] {
                    true
                } else {
                    panic!("Error!");
                }
            }
            Circuit::Wire(w) => {
                if out_label == w.keys[0] {
                    false
                } else if out_label == w.keys[1] {
                    true
                } else {
                    panic!("Error!");
                }
            }
        }
    }

    pub fn to_hidden(
        &self,
    ) -> (
        FullyHiddenCircuit,
        HashMap<String, [[u8; 16]; 2]>,
        HashMap<String, [[u8; 16]; 2]>,
    ) {
        let (circ, map1, map2) = self.to_hidden_re();
        match self {
            Circuit::Gate(c) => {
                let out_keys = c.out.keys;
                let mut table = [[0; 32]; 2];

                let i = OsRng.gen::<bool>() as usize;

                for bit in [i, 1 - i] {
                    let enc = encrypt(out_keys[bit], [0; 16], [bit as u8; 16]);
                    table[bit] = enc;
                }
                (
                    FullyHiddenCircuit {
                        table,
                        circuit: circ,
                    },
                    map1,
                    map2,
                )
            }
            Circuit::Wire(w) => {
                let out_keys = w.keys;
                let mut table = [[0; 32]; 2];
                let i = OsRng.gen::<bool>() as usize;
                for bit in [i, 1 - i] {
                    let enc = encrypt(out_keys[bit], [0; 16], [bit as u8; 16]);
                    table[bit] = enc;
                }
                (
                    FullyHiddenCircuit {
                        table,
                        circuit: circ,
                    },
                    map1,
                    map2,
                )
            }
        }
    }

    fn to_hidden_re(
        &self,
    ) -> (
        HiddenCircuit,
        HashMap<String, [[u8; 16]; 2]>,
        HashMap<String, [[u8; 16]; 2]>,
    ) {
        match self {
            Circuit::Gate(c) => {
                let (c1, sender1, eval1) = c.left.to_hidden_re();
                let (c2, sender2, eval2) = c.right.to_hidden_re();
                (
                    HiddenCircuit::Gate(HiddenGate {
                        left: Box::new(c1),
                        right: Box::new(c2),
                        table: c.table,
                        sym: c.sym.clone(),
                    }),
                    sender1.into_iter().chain(sender2).collect(),
                    eval1.into_iter().chain(eval2).collect(),
                )
            }
            Circuit::Wire(w) => {
                let out = match w.sym.clone() {
                    Some(s) => s,
                    None => panic!("Need wire description."),
                };
                if w.evaluator {
                    let sender = HashMap::new();
                    let mut eval = HashMap::new();
                    eval.insert(out, w.keys);
                    (
                        HiddenCircuit::Wire(HiddenWire { sym: w.sym.clone() }),
                        sender,
                        eval,
                    )
                } else {
                    let mut sender = HashMap::new();
                    let eval = HashMap::new();
                    sender.insert(out, w.keys);
                    (
                        HiddenCircuit::Wire(HiddenWire { sym: w.sym.clone() }),
                        sender,
                        eval,
                    )
                }
            }
        }
    }
}

impl BitAnd for Circuit {
    type Output = Circuit;

    fn bitand(self, rhs: Self) -> Self::Output {
        Circuit::Gate(Gate::new(
            self,
            rhs,
            [[false, false], [false, true]],
            "&".to_string(),
            &mut OsRng,
        ))
    }
}

impl BitOr for Circuit {
    type Output = Circuit;

    fn bitor(self, rhs: Self) -> Self::Output {
        Circuit::Gate(Gate::new(
            self,
            rhs,
            [[false, true], [true, true]],
            "|".to_string(),
            &mut OsRng,
        ))
    }
}

impl BitXor for Circuit {
    type Output = Circuit;

    fn bitxor(self, rhs: Self) -> Self::Output {
        Circuit::Gate(Gate::new(
            self,
            rhs,
            [[false, true], [true, false]],
            "^".to_string(),
            &mut OsRng,
        ))
    }
}

impl From<Wire> for HiddenWire {
    fn from(value: Wire) -> Self {
        Self { sym: value.sym }
    }
}

impl From<Gate> for HiddenGate {
    fn from(value: Gate) -> Self {
        Self {
            left: Box::new((*value.left).into()),
            right: Box::new((*value.right).into()),
            table: value.table,
            sym: value.sym,
        }
    }
}

impl From<Circuit> for HiddenCircuit {
    fn from(value: Circuit) -> Self {
        match value {
            Circuit::Gate(c) => Self::Gate(c.into()),
            Circuit::Wire(w) => Self::Wire(w.into()),
        }
    }
}

impl HiddenCircuit {
    fn evaluate_re(&self, vals: HashMap<String, [u8; 16]>) -> [u8; 16] {
        match self {
            HiddenCircuit::Gate(c) => {
                let lv = c.left.evaluate_re(vals.clone());
                let rv = c.right.evaluate_re(vals);
                let mut out = [0; 16];
                for i in [0, 1] {
                    for j in [0, 1] {
                        if let Some(x) = decrypt(lv, rv, c.table[i][j]) {
                            out = x;
                        }
                    }
                }
                out
            }
            HiddenCircuit::Wire(w) => match &w.sym {
                Some(c) => vals[c],
                None => panic!("This should not happen."),
            },
        }
    }
}

impl FullyHiddenCircuit {
    pub fn evaluate(
        &self,
        sender: HashMap<String, [u8; 16]>,
        eval: HashMap<String, [u8; 16]>,
    ) -> bool {
        let out_label = self
            .circuit
            .evaluate_re(sender.into_iter().chain(eval).collect());

        let mut out = false;
        for i in [0, 1] {
            if let Some(x) = decrypt(out_label, [0; 16], self.table[i]) {
                out = x[0] != 0;
            }
        }
        out
    }
}

pub fn map_values_to_keys(
    keys: HashMap<String, [[u8; 16]; 2]>,
    values: Vec<(String, bool)>,
) -> HashMap<String, [u8; 16]> {
    let mut ret = HashMap::new();
    for i in values {
        match keys.get(&i.0) {
            Some(v) => ret.insert(i.0, v[i.1 as usize]),
            None => continue,
        };
    }
    ret
}

#[cfg(test)]
mod test {

    use crate::garble_classic::Wire;
    use crate::garble_classic::{decrypt, encrypt, map_values_to_keys};
    use std::collections::HashMap;

    #[test]
    fn enc() {
        let a = decrypt([2; 16], [4; 16], encrypt([2; 16], [4; 16], [9; 16]));
        assert!(a == Some([9; 16]));
    }

    #[test]
    fn public_circuit() {
        let a = Wire::new("a".to_string(), true).to_circ();
        let b = Wire::new("b".to_string(), false).to_circ();
        let c = Wire::new("c".to_string(), false).to_circ();
        let circ = (a | b) & c;

        let mut h = HashMap::new();
        h.insert("a".to_string(), false);
        h.insert("b".to_string(), false);
        h.insert("c".to_string(), true);

        let out = circ.evaluate(h);
        assert!(out == false);
    }

    #[test]
    fn garbled_circuit() {
        let a = Wire::new("a".to_string(), true).to_circ();
        let b = Wire::new("b".to_string(), false).to_circ();
        let c = Wire::new("c".to_string(), false).to_circ();
        let circ = (a | b) & c;

        let (hidden_circ, sender, eval) = circ.to_hidden();

        let values_sender = map_values_to_keys(
            sender,
            vec![("b".to_string(), false), ("c".to_string(), true)],
        );

        let values_eval = map_values_to_keys(eval, vec![("a".to_string(), true)]);

        let out = hidden_circ.evaluate(values_sender, values_eval);

        assert!(out == true);
    }
}