use std::borrow::Cow;

use bencode::{Dict, Object};

/// Allows the encode function to take something that can be converted into an object,
/// or a reference to an object.
pub trait Encodable<'a> {
    fn get_object(self) -> Cow<'a, Object>;
}

impl<'a, T> Encodable<'a> for T where T: Into<Object> {
    fn get_object(self) -> Cow<'a, Object> {
        Cow::Owned(self.into())
    }
}

impl<'a> Encodable<'a> for &'a Object {
    fn get_object(self) -> Cow<'a, Object> {
        Cow::Borrowed(self)
    }
}

pub fn encode<'a, O>(obj: O) -> Vec<u8> where O: Encodable<'a> {
    let mut buff = Vec::new();
    encode_object(&mut buff, &obj.get_object());
    buff
}

fn encode_object(buff: &mut Vec<u8>, obj: &Object) {
    match *obj {
        Object::Int(num) => encode_int(buff, num),
        Object::Bytes(ref bytes) => encode_bytes(buff, bytes),
        Object::List(ref list) => encode_list(buff, list),
        Object::Dict(ref dict) => encode_dict(buff, dict),
    }
}

fn encode_int(buff: &mut Vec<u8>, num: i64) {
    buff.extend_from_slice(format!("i{}e", num).as_bytes());
}

fn encode_bytes(buff: &mut Vec<u8>, bytes: &[u8]) {
    buff.extend_from_slice(format!("{}:", bytes.len()).as_bytes());
    buff.extend_from_slice(bytes);
}

fn encode_list(buff: &mut Vec<u8>, list: &[Object]) {
    buff.push(b'l');
    for obj in list {
        encode_object(buff, obj);
    }
    buff.push(b'e');
}

fn encode_dict(buff: &mut Vec<u8>, dict: &Dict) {
    buff.push(b'd');
    for (key, val) in dict.iter() {
        encode_bytes(buff, key);
        encode_object(buff, val);
    }
    buff.push(b'e');
}

#[test]
fn test_int_pos() {
    assert_eq!(encode(1337), b"i1337e");
}

#[test]
fn test_int_neg() {
    assert_eq!(encode(-1337), b"i-1337e");
}

#[test]
fn test_bytes() {
    assert_eq!(encode("hello"), b"5:hello");
}

#[test]
fn test_list() {
    let list = vec![Object::Int(1), Object::Int(2), Object::Int(3)];
    assert_eq!(encode(list), b"li1ei2ei3ee");
}

#[test]
fn test_dict() {
    use bencode::Dict;

    let mut dict = Dict::new();
    dict.insert("hello", 1337);
    assert_eq!(encode(dict), b"d5:helloi1337ee")
}