use num::bigint::BigInt;
use num::cast::ToPrimitive;
use std::convert::TryInto;
use super::constants::{Op, AtomMinor, CompoundMinor};

pub type Error = std::io::Error;
pub type Result = std::result::Result<(), Error>;

pub trait Writer {
    fn write_annotation_prefix(&mut self) -> Result;
    fn write_placeholder_ref(&mut self, v: usize) -> Result;
    fn write_noop(&mut self) -> Result;

    fn write_bool(&mut self, v: bool) -> Result;

    fn write_f32(&mut self, v: f32) -> Result;
    fn write_f64(&mut self, v: f64) -> Result;

    fn write_i8(&mut self, v: i8) -> Result;
    fn write_u8(&mut self, v: u8) -> Result;
    fn write_i16(&mut self, v: i16) -> Result;
    fn write_u16(&mut self, v: u16) -> Result;
    fn write_i32(&mut self, v: i32) -> Result;
    fn write_u32(&mut self, v: u32) -> Result;
    fn write_i64(&mut self, v: i64) -> Result;
    fn write_u64(&mut self, v: u64) -> Result;
    fn write_int(&mut self, v: &BigInt) -> Result;

    fn write_string(&mut self, v: &str) -> Result;
    fn write_bytes(&mut self, v: &[u8]) -> Result;
    fn write_symbol(&mut self, v: &str) -> Result;

    fn open_record(&mut self, field_count: usize) -> Result;
    fn open_sequence(&mut self, item_count: usize) -> Result;
    fn open_set(&mut self, item_count: usize) -> Result;
    fn open_dictionary(&mut self, entry_count: usize) -> Result;

    fn close_record(&mut self) -> Result;
    fn close_sequence(&mut self) -> Result;
    fn close_set(&mut self) -> Result;
    fn close_dictionary(&mut self) -> Result;

    fn stream_string(&mut self) -> Result;
    fn stream_bytes(&mut self) -> Result;
    fn stream_symbol(&mut self) -> Result;

    fn stream_record(&mut self) -> Result;
    fn stream_sequence(&mut self) -> Result;
    fn stream_set(&mut self) -> Result;
    fn stream_dictionary(&mut self) -> Result;

    fn close_stream(&mut self) -> Result;
}

pub fn varint<W: std::io::Write>(w: &mut W, mut v: usize) -> Result {
    loop {
        if v < 128 {
            return w.write_all(&[v as u8])
        } else {
            w.write_all(&[((v & 0x7f) + 128) as u8])?;
            v = v >> 7;
        }
    }
}

pub fn write_op<W: std::io::Write>(w: &mut W, op: Op, arg: u8) -> Result {
    w.write_all(&[(u8::from(op) << 4) | (arg & 15)])
}

pub fn write_header<W: std::io::Write>(w: &mut W, op: Op, wirelength: usize) -> Result {
    if wirelength < 15 {
        write_op(w, op, wirelength as u8)
    } else {
        write_op(w, op, 15)?;
        varint(w, wirelength)
    }
}

pub fn write_atom<W: std::io::Write>(w: &mut W, minor: AtomMinor, bs: &[u8]) -> Result {
    write_header(w, Op::Atom(minor), bs.len())?;
    w.write_all(bs)
}

impl<W: std::io::Write> Writer for W {
    fn write_annotation_prefix(&mut self) -> Result {
        write_header(self, Op::Misc(0), 5)
    }

    fn write_placeholder_ref(&mut self, v: usize) -> Result {
        write_header(self, Op::Misc(1), v)
    }

    fn write_noop(&mut self) -> Result {
        write_op(self, Op::Reserved(3), 15)
    }

    fn write_bool(&mut self, v: bool) -> Result {
        write_op(self, Op::Misc(0), if v { 1 } else { 0 })
    }

    fn write_f32(&mut self, v: f32) -> Result {
        write_op(self, Op::Misc(0), 2)?;
        self.write_all(&u32::to_be_bytes(f32::to_bits(v)))
    }

    fn write_f64(&mut self, v: f64) -> Result {
        write_op(self, Op::Misc(0), 3)?;
        self.write_all(&u64::to_be_bytes(f64::to_bits(v)))
    }

    fn write_i8(&mut self, v: i8) -> Result {
        if v >= 0 && v <= 12 { return write_op(self, Op::Misc(3), v as u8) }
        if v >= -3 && v < 0 { return write_op(self, Op::Misc(3), (v + 16) as u8) }
        write_atom(self, AtomMinor::SignedInteger, &[v as u8])
    }

    fn write_u8(&mut self, v: u8) -> Result {
        if let Ok(w) = v.try_into() { return self.write_i8(w) }
        write_atom(self, AtomMinor::SignedInteger, &[0, v])
    }

    fn write_i16(&mut self, v: i16) -> Result {
        if let Ok(w) = v.try_into() { return self.write_i8(w) }
        write_atom(self, AtomMinor::SignedInteger, &[(v >> 8) as u8, (v & 255) as u8])
    }

    fn write_u16(&mut self, v: u16) -> Result {
        if let Ok(w) = v.try_into() { return self.write_i16(w) }
        write_atom(self, AtomMinor::SignedInteger, &[0, (v >> 8) as u8, (v & 255) as u8])
    }

    fn write_i32(&mut self, v: i32) -> Result {
        if let Ok(w) = v.try_into() { return self.write_i16(w) }
        if v >= -(2 << 23) && v < (2 << 23) {
            return write_atom(self, AtomMinor::SignedInteger, &[(v >> 16) as u8,
                                                                (v >> 8) as u8,
                                                                (v & 255) as u8]);
        }
        return write_atom(self, AtomMinor::SignedInteger, &[(v >> 24) as u8,
                                                            (v >> 16) as u8,
                                                            (v >> 8) as u8,
                                                            (v & 255) as u8]);
    }

    fn write_u32(&mut self, v: u32) -> Result {
        if let Ok(w) = v.try_into() { return self.write_i32(w) }
        return write_atom(self, AtomMinor::SignedInteger, &[0,
                                                            (v >> 24) as u8,
                                                            (v >> 16) as u8,
                                                            (v >> 8) as u8,
                                                            (v & 255) as u8]);
    }

    fn write_i64(&mut self, v: i64) -> Result {
        if let Ok(w) = v.try_into() { return self.write_i32(w) }
        if v >= -(2 << 39) && v < (2 << 39) {
            return write_atom(self, AtomMinor::SignedInteger, &[(v >> 32) as u8,
                                                                (v >> 24) as u8,
                                                                (v >> 16) as u8,
                                                                (v >> 8) as u8,
                                                                (v & 255) as u8]);
        }
        if v >= -(2 << 47) && v < (2 << 47) {
            return write_atom(self, AtomMinor::SignedInteger, &[(v >> 40) as u8,
                                                                (v >> 32) as u8,
                                                                (v >> 24) as u8,
                                                                (v >> 16) as u8,
                                                                (v >> 8) as u8,
                                                                (v & 255) as u8]);
        }
        if v >= -(2 << 55) && v < (2 << 55) {
            return write_atom(self, AtomMinor::SignedInteger, &[(v >> 48) as u8,
                                                                (v >> 40) as u8,
                                                                (v >> 32) as u8,
                                                                (v >> 24) as u8,
                                                                (v >> 16) as u8,
                                                                (v >> 8) as u8,
                                                                (v & 255) as u8]);
        }
        return write_atom(self, AtomMinor::SignedInteger, &[(v >> 56) as u8,
                                                            (v >> 48) as u8,
                                                            (v >> 40) as u8,
                                                            (v >> 32) as u8,
                                                            (v >> 24) as u8,
                                                            (v >> 16) as u8,
                                                            (v >> 8) as u8,
                                                            (v & 255) as u8]);
    }

    fn write_u64(&mut self, v: u64) -> Result {
        if let Ok(w) = v.try_into() { return self.write_i64(w) }
        return write_atom(self, AtomMinor::SignedInteger, &[0,
                                                            (v >> 56) as u8,
                                                            (v >> 48) as u8,
                                                            (v >> 40) as u8,
                                                            (v >> 32) as u8,
                                                            (v >> 24) as u8,
                                                            (v >> 16) as u8,
                                                            (v >> 8) as u8,
                                                            (v & 255) as u8]);
    }

    fn write_int(&mut self, v: &BigInt) -> Result {
        match v.to_i8() {
            Some(n) => self.write_i8(n),
            None => write_atom(self, AtomMinor::SignedInteger, &v.to_signed_bytes_be()),
        }
    }

    fn write_string(&mut self, v: &str) -> Result {
        write_atom(self, AtomMinor::String, v.as_bytes())
    }

    fn write_bytes(&mut self, v: &[u8]) -> Result {
        write_atom(self, AtomMinor::ByteString, v)
    }

    fn write_symbol(&mut self, v: &str) -> Result {
        write_atom(self, AtomMinor::Symbol, v.as_bytes())
    }

    fn open_record(&mut self, field_count: usize) -> Result {
        write_header(self, Op::Compound(CompoundMinor::Record), field_count + 1)
    }

    fn open_sequence(&mut self, item_count: usize) -> Result {
        write_header(self, Op::Compound(CompoundMinor::Sequence), item_count)
    }

    fn open_set(&mut self, item_count: usize) -> Result {
        write_header(self, Op::Compound(CompoundMinor::Set), item_count)
    }

    fn open_dictionary(&mut self, entry_count: usize) -> Result {
        write_header(self, Op::Compound(CompoundMinor::Dictionary), entry_count << 1)
    }

    fn close_record(&mut self) -> Result { Ok(()) }
    fn close_sequence(&mut self) -> Result { Ok(()) }
    fn close_set(&mut self) -> Result { Ok(()) }
    fn close_dictionary(&mut self) -> Result { Ok(()) }

    fn stream_string(&mut self) -> Result {
        write_op(self, Op::Misc(2), Op::Atom(AtomMinor::String).into())
    }

    fn stream_bytes(&mut self) -> Result {
        write_op(self, Op::Misc(2), Op::Atom(AtomMinor::ByteString).into())
    }

    fn stream_symbol(&mut self) -> Result {
        write_op(self, Op::Misc(2), Op::Atom(AtomMinor::Symbol).into())
    }

    fn stream_record(&mut self) -> Result {
        write_op(self, Op::Misc(2), Op::Compound(CompoundMinor::Record).into())
    }

    fn stream_sequence(&mut self) -> Result {
        write_op(self, Op::Misc(2), Op::Compound(CompoundMinor::Sequence).into())
    }

    fn stream_set(&mut self) -> Result {
        write_op(self, Op::Misc(2), Op::Compound(CompoundMinor::Set).into())
    }

    fn stream_dictionary(&mut self) -> Result {
        write_op(self, Op::Misc(2), Op::Compound(CompoundMinor::Dictionary).into())
    }

    fn close_stream(&mut self) -> Result {
        write_op(self, Op::Misc(0), 4)
    }
}