// Panopticon - A libre program analysis library for machine code
// Copyright (C) 2014-2018  The Panopticon Developers
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

//! Universally unique identifiers.
//!
//! Used to give long-term identifier to functions and other data structures.
//! Folows the RON spec:
//! https://github.com/gritzko/swarm-ron-docs/blob/master/uid.md

use std::time::SystemTime;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::cmp;
use std::fmt;
use std::result;
use std::str::FromStr;

use chrono::{Utc, DateTime, Datelike, Timelike};
use quickcheck::{Arbitrary, Gen};

use {Str, Result};

static UUID_NODE_ID: AtomicUsize = ATOMIC_USIZE_INIT;
static UUID_SEQUENCE: AtomicUsize = ATOMIC_USIZE_INIT;

/// Universally unique identifier.
///
/// UUIDs consist of an origin and an optional timestamp, both 64 bit integers.
#[derive(Debug, PartialEq, Eq, Clone, Hash)]
pub struct Uuid {
    origin: u64,
    timestamp: u64,
}

impl cmp::PartialOrd for Uuid {
    fn partial_cmp(&self, other: &Uuid) -> Option<cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl cmp::Ord for Uuid {
    fn cmp(&self, other: &Uuid) -> cmp::Ordering {
        let a = self.origin.cmp(&other.origin);
        if a == cmp::Ordering::Equal {
            self.timestamp.cmp(&other.timestamp)
        } else {
            a
        }
    }
}

impl fmt::Display for Uuid {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Uuid{ origin: 0, timestamp: 0 } =>
                f.write_str("0"),
            Uuid{ origin, timestamp: 0 } =>
                f.write_str(&Self::encode_b64(*origin)),
            Uuid{ origin, timestamp } => {
                f.write_str(&Self::encode_b64(*origin))?;
                f.write_str("-")?;
                f.write_str(&Self::encode_b64(*timestamp))
            }
        }
    }
}

impl Default for Uuid {
    fn default() -> Self {
        Uuid{ origin: 0, timestamp: 0 }
    }
}

impl Arbitrary for Uuid {
    fn arbitrary<G: Gen>(g: &mut G) -> Self {
        if g.gen_weighted_bool(50) {
            if g.gen::<bool>() {
                Uuid{
                    origin: 0,
                    timestamp: g.gen::<u64>() & 0xfffffffffffffff
                }
            } else {
                Uuid{
                    origin: g.gen::<u64>() & 0xfffffffffffffff,
                    timestamp: 0,
                }
            }
        } else {
            Uuid{
                origin: g.gen::<u64>() & 0xfffffffffffffff,
                timestamp: g.gen::<u64>() & 0xfffffffffffffff
            }
        }
    }
}

impl FromStr for Uuid {
    type Err = ();
    fn from_str(s: &str) -> result::Result<Uuid, ()> {
        let mut s = s.as_bytes().iter().cloned();
        let origin = Self::decode_b64(&mut s).map_err(|_| ())?;

        if s.next() == Some(b'-') {
            let timestamp = Self::decode_b64(&mut s).map_err(|_| ())?;

            Ok(Uuid{ origin: origin, timestamp: timestamp })
        } else {
            Ok(Uuid{ timestamp: 0, origin: origin })
        }
    }
}

impl Uuid {
    /// Sets the default UUID origin. `s` must not be longer than 10 characters and only consist of
    /// [0-9a-zA-Z~_].
    pub fn set_node_id(s: &str) -> Result<()> {
        let mut i = s.as_bytes().iter().cloned();
        let val = Self::decode_b64(&mut i)?;

        if i.next().is_some() { return Err("Not a valid UUID".into()); }
        UUID_NODE_ID.store(val as usize, Ordering::Relaxed);
        Ok(())
    }

    /// The current default UUID origin. Initially "0".
    pub fn node_id() -> Str {
        Self::encode_b64(UUID_NODE_ID.load(Ordering::Relaxed) as u64)
    }

    /// New "transcendent" UUID with origin `s` and no timestamp. Used to identify "global"
    /// entities like types. `s` must a array of [0-9a-zA-Z~_] no longer than 10 characters.
    pub fn known(s: &[u8]) -> Result<Self> {
        if s.len() <= 10 {
            Ok(Uuid{ origin: Self::decode_b64(&mut s.iter().cloned())?, timestamp: 0 })
        } else {
            Err("string too long".into())
        }
    }

    /// New UUID with the default origin (see `node_id` and `set_node_id`) and the current time.
    /// Ignoring leap seconds and other events that mess with the system time all calls to this
    /// functions return unique UUID (duh).
    pub fn now() -> Self {
        Uuid{
            origin: UUID_NODE_ID.load(Ordering::Relaxed) as u64,
            timestamp:  Self::timestamp()
        }
    }

    /// New timestamped UUID with origin `s`. See `now`. `s` must a array of
    /// [0-9a-zA-Z~_] no longer than 10 characters.
    pub fn now_foreign(s: &[u8]) -> Result<Self> {
        if s.len() <= 10 {
                Ok(Uuid{
                    origin: Self::decode_b64(&mut s.iter().cloned())?,
                    timestamp:  Self::timestamp()
                })
        } else {
            Err("location too long".into())
        }
    }

    /// Parses a 16 byte UUID from `b`. Must be a 16 byte slice.
    pub fn from_bytes(b: &[u8]) -> Result<Uuid> {
        if b.len() != 16 {
            Err("wrong byte slice size".into())
        } else {
            let origin = (b[0] as u64) << 56
                       | (b[1] as u64) << 48
                       | (b[2] as u64) << 40
                       | (b[3] as u64) << 32
                       | (b[4] as u64) << 24
                       | (b[5] as u64) << 16
                       | (b[6] as u64) << 8
                       | (b[7] as u64) << 0;
            let timestamp = (b[8] as u64) << 56
                          | (b[9] as u64) << 48
                          | (b[10] as u64) << 40
                          | (b[11] as u64) << 32
                          | (b[12] as u64) << 24
                          | (b[13] as u64) << 16
                          | (b[14] as u64) << 8
                          | (b[15] as u64) << 0;

            Ok(Uuid{
                origin: origin, timestamp: timestamp
            })
        }
    }

    /// Returns this UUID as a byte array.
    pub fn as_bytes(&self) -> [u8; 16] {
        [
            (self.origin >> 56 & 0xff) as u8,
            (self.origin >> 48 & 0xff) as u8,
            (self.origin >> 40 & 0xff) as u8,
            (self.origin >> 32 & 0xff) as u8,
            (self.origin >> 24 & 0xff) as u8,
            (self.origin >> 16 & 0xff) as u8,
            (self.origin >> 8 & 0xff) as u8,
            (self.origin >> 0 & 0xff) as u8,
            (self.timestamp >> 56 & 0xff) as u8,
            (self.timestamp >> 48 & 0xff) as u8,
            (self.timestamp >> 40 & 0xff) as u8,
            (self.timestamp >> 32 & 0xff) as u8,
            (self.timestamp >> 24 & 0xff) as u8,
            (self.timestamp >> 16 & 0xff) as u8,
            (self.timestamp >> 8 & 0xff) as u8,
            (self.timestamp >> 0 & 0xff) as u8,
        ]
    }

    fn timestamp() -> u64 {
        use std::time::Duration;
        use std::thread::sleep;

        let now = SystemTime::now();
        let seq = (UUID_SEQUENCE.fetch_add(1, Ordering::SeqCst) & 0b111111_111111) as u64;

        // we wait for the clock to advance after sequence number overflows to avoid creating
        // duplicated uuids. This can happen if we try to generate more than 4k instances in
        // under one ms.
        if seq == 0b111111_111111 {
            while now == SystemTime::now() {
                sleep(Duration::from_millis(1));
            }
        }

        let dt: DateTime<Utc> = now.into();
        let months = (2010 + dt.year() as u64 * 12 + dt.month0() as u64) & 0b111111_111111;;
        let ts = (months << (8 * 6))
               | ((dt.day0() as u64   & 0b111111) << (7 * 6))
               | ((dt.hour() as u64   & 0b111111) << (6 * 6))
               | ((dt.minute() as u64 & 0b111111) << (5 * 6))
               | ((dt.second() as u64 & 0b111111) << (4 * 6))
               | (((dt.nanosecond() as u64 / 1_000_000) & 0b111111_111111) << (2 * 6))
               | seq;

        ts
    }

    fn decode_b64<I: Iterator<Item=u8> + Clone>(i: &mut I) -> Result<u64> {
        let mut chars = i.clone();
        let mut ret = 0;
        let mut num = 0;

        for j in (0..10).rev() {
            let c = match chars.next() {
                Some(a) if a >= b'0' && a <= b'9' => a - b'0',
                Some(a) if a >= b'A' && a <= b'Z' => a - b'A' + 10,
                Some(a) if a == b'_' => 36,
                Some(a) if a >= b'a' && a <= b'z' => a - b'a' + 37,
                Some(a) if a == b'~' => 63,
                _ => {for _ in 0..num { i.next(); }
                    return Ok(ret);
                }
            };

            num += 1;
            ret |= (c as u64) << (j * 6);
        }

        for _ in 0..num { i.next(); }
        Ok(ret)
    }

    fn encode_b64(i: u64) -> Str {
        let mut ret = String::new();

        for idx in 0..10 {
            let idx = (9 - idx) * 6;
            let val = (i >> idx) & 0b111111;

            let c = match val as u8 {
                x@0...9   => b'0' + x,
                x@10...35 => b'A' + x - 10,
                36        => b'_',
                x@37...62 => b'a' + x - 37,
                63        => b'~',
                _ => unreachable!(),
            };

            ret.push(char::from(c));
        }

        ret = ret.trim_end_matches('0').to_string();

        if ret.is_empty() {
            "0".into()
        } else {
            ret.into()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn ts_unique() {
        let mut ts = (0..(5*0x1_0000usize))
            .map(|_| Uuid::now_foreign(b"test").unwrap()).collect::<Vec<_>>();

        ts.sort();
        assert_eq!(ts.len(), 5*0x1_0000);
        ts.dedup();
        assert_eq!(ts.len(), 5*0x1_0000);
    }

    quickcheck!{
        fn base64_roundtrip(v: u64) -> bool {
            let w = Uuid::decode_b64(&mut Uuid::encode_b64(v).to_string().into_bytes().into_iter()).unwrap();
            w == v
        }
    }

    quickcheck!{
        fn as_bytes_roundtrip(v: Uuid) -> bool {
            let w = Uuid::from_bytes(&v.as_bytes()).unwrap();
            w == v
        }
    }

    quickcheck!{
        fn from_str_roundtrip(v: Uuid) -> bool {
            let w = Uuid::from_str(&format!("{}",v));
            w.unwrap() == v
        }
    }

    #[test]
    fn decode_b64_one_dir_known() {
        assert_eq!(Uuid::decode_b64(&mut b"inc0000000".iter().cloned()).ok(), Some(824893205576155136));
        assert_eq!(Uuid::decode_b64(&mut b"inc".iter().cloned()).ok(), Some(824893205576155136));
    }

    #[test]
    fn node_id() {
        assert_eq!(Uuid::node_id(), "0");
        assert!(Uuid::set_node_id("test").is_ok());
        assert_eq!(Uuid::node_id(), "test");
        assert!(Uuid::set_node_id("test{").is_err());
        assert_eq!(Uuid::node_id(), "test");
    }
}