// Adapted from : https://docs.rs/prometheus/0.13.0/src/prometheus/encoder/text.rs.html#21
    // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.


use std::io::{self, Write};

use prometheus::Result;
use prometheus::proto::{self, MetricFamily, MetricType};

use prometheus::{Encoder};
use serde_json::{Map, Value, json};

/// The text format of metric family.
pub const TEXT_FORMAT: &str = "text/plain; version=0.0.4";

const POSITIVE_INF: &str = "+Inf";

/// An implementation of an [`Encoder`] that converts a [`MetricFamily`] proto message
/// into json format.

#[derive(Debug, Default)]
pub struct JsonEncoder;

impl JsonEncoder {
    /// Create a new json encoder.
    pub fn new() -> JsonEncoder {
        JsonEncoder
    }
    /// Appends metrics to a given `String` buffer.
    ///
    /// This is a convenience wrapper around `<JsonEncoder as Encoder>::encode`.
    pub fn encode_utf8(&self, metric_families: &[MetricFamily], buf: &mut String) -> Result<()> {
        // Note: it's important to *not* re-validate UTF8-validity for the
        // entirety of `buf`. Otherwise, repeatedly appending metrics to the
        // same `buf` will lead to quadratic behavior. That's why we use
        // `WriteUtf8` abstraction to skip the validation.
        self.encode_impl(metric_families, &mut StringBuf(buf))?;
        Ok(())
    }
    /// Converts metrics to `String`.
    ///
    /// This is a convenience wrapper around `<JsonEncoder as Encoder>::encode`.
    pub fn encode_to_string(&self, metric_families: &[MetricFamily]) -> Result<String> {
        let mut buf = String::new();
        self.encode_utf8(metric_families, &mut buf)?;
        Ok(buf)
    }

    fn encode_impl(
        &self,
        metric_families: &[MetricFamily],
        writer: &mut dyn WriteUtf8,
    ) -> Result<()> {

        let mut map = Map::new();

        for mf in metric_families {
            
            // Add entry for the metric
            let name : &str = mf.get_name();
            
            let mut mf_map = Map::new();

            // Add Help
            let help : &str = mf.get_help();
            mf_map.insert("help".to_string(), json!{help}); 

            // Write `# TYPE` header.
            let metric_type : MetricType = mf.get_field_type();
            let lowercase_type = json!(format!("{:?}", metric_type).to_lowercase());
            mf_map.insert("type".to_string(), lowercase_type); 

            let mut debug_counter = 0;

            for m in mf.get_metric() {
                println!("{}", debug_counter);
                debug_counter += 1;         // TODO : Remove
                // Metric
                match metric_type {
                    MetricType::COUNTER => {
                        mf_map.insert("counter".to_string(), json!(m.get_counter().get_value()));
                        extra_info(&mut mf_map, m);
                        // f64
                    }
                    MetricType::GAUGE => {
                        mf_map.insert("gauge".to_string(), json!(m.get_gauge().get_value()));
                        extra_info(&mut mf_map, m);
                        // f64
                    }
                    MetricType::HISTOGRAM => {
                        let h = m.get_histogram();
                        let mut upper_bounds : Vec<Value> = vec![];
                        let mut cumulative_counts : Vec<Value> = vec![];
                        let mut inf_seen = false;
                        
                        for b in h.get_bucket() {
                            let upper_bound = b.get_upper_bound();           // f64
                            let cumulative_count = b.get_cumulative_count(); // f64

                            upper_bounds.push(json!(upper_bound));
                            cumulative_counts.push(json!(cumulative_count));

                            if upper_bound.is_sign_positive() && upper_bound.is_infinite() {
                                inf_seen = true;
                            }
                        }
                        if !inf_seen {
                            upper_bounds.push(json!(POSITIVE_INF));
                            cumulative_counts.push(json!(h.get_sample_count()));
                        }
                        let names = [
                            "cumulative_counts".to_string(),
                            "upper_bounds".to_string(),
                            "sum".to_string(),
                            "counts".to_string()
                        ];

                        let values = [
                            json!(cumulative_counts),
                            json!(upper_bounds),
                            json!(h.get_sample_sum()),
                            json!(h.get_sample_count())
                        ];
                        for (key, value) in names.into_iter().zip(values.into_iter()) {
                            mf_map.insert(key, value);
                        }
                        extra_info(&mut mf_map, m);
                    }

                    MetricType::SUMMARY => {
                        let s = m.get_summary();
                        let mut quantiles = vec![];
                        let mut values = vec![];

                        for q in s.get_quantile() {
                            quantiles.push(json!(q.get_quantile()));
                            values.push(q.get_value());
                        }

                        let names = [
                            "sum".to_string(),
                            "quantiles".to_string()
                        ];

                        let values = [
                            json!(s.get_sample_sum()),
                            json!(s.get_sample_count())
                        ];
                        for (key, value) in names.into_iter().zip(values.into_iter()) {
                            mf_map.insert(key, value);
                        }
                        extra_info(&mut mf_map, m);
                    }
                    MetricType::UNTYPED => {
                        unimplemented!();
                    }
                }
            }
            map.insert(name.to_string(), json!(mf_map));
        }

        let x = serde_json::to_vec(&map).unwrap();
        //println!{"{}", &x};
        writer.write_all(&x)?; // String 
        Ok(())
    }
    
}

impl Encoder for JsonEncoder {
    fn encode<W: Write>(&self, metric_families: &[MetricFamily], writer: &mut W) -> Result<()> {
        self.encode_impl(metric_families, &mut *writer)
    }

    fn format_type(&self) -> &str {
        TEXT_FORMAT
    }
}

// Adds into a map m.timestamp and m.LabelPair.
// names and values must be of the same length
fn extra_info(
    map : &mut Map<String, Value>,
    mc: &proto::Metric
) -> () {

    let timestamp = mc.get_timestamp_ms();
    if timestamp != 0 {
        map.insert("timestamp".to_string(), json!(timestamp));
    }

    for lp in mc.get_label() {
        map.insert(lp.get_name().to_string(), json!(lp.get_value()));
    }
}

trait WriteUtf8 {
    fn write_all(&mut self, text: &[u8]) -> io::Result<()>;
}

impl<W: Write> WriteUtf8 for W {
    fn write_all(&mut self, text: &[u8]) -> io::Result<()> {
        Write::write_all(self, text)
    }
}

/// Coherence forbids to impl `WriteUtf8` directly on `String`, need this
/// wrapper as a work-around.
struct StringBuf<'a>(&'a mut String);

impl WriteUtf8 for StringBuf<'_> {
    fn write_all(&mut self, text: &[u8]) -> io::Result<()> {
        self.0.push_str(std::str::from_utf8(text).unwrap());
        Ok(())
    }
}


#[cfg(test)]
mod tests {

    use super::*;
    use prometheus::Counter;
    use prometheus::Gauge;
    use prometheus::{Histogram, HistogramOpts};
    use prometheus::Opts;
    use prometheus::core::Collector;

    #[test]
    fn test_json_encoder() {
        let counter_opts = Opts::new("test_counter", "test help")
            .const_label("a", "1")
            .const_label("b", "2");
        let counter = Counter::with_opts(counter_opts).unwrap();
        counter.inc();

        let mf = counter.collect();
        let mut writer = Vec::<u8>::new();
        let encoder = JsonEncoder::new();
        let txt = encoder.encode(&mf, &mut writer);
        assert!(txt.is_ok());

        // Object({"test_counter": Object({"a": String("1"), "b": String("2"), "counter": Number(1.0), "help": String("test help"), "type": String("counter")})})
        let v: Value = serde_json::from_slice(&writer).unwrap();
        

        assert_eq!(v["test_counter"]["help"], "test help");
        assert_eq!(v["test_counter"]["a"], "1");
        assert_eq!(v["test_counter"]["b"], "2");
        assert_eq!(v["test_counter"]["type"], "counter");
        assert_eq!(v["test_counter"]["counter"], 1.0);
        


        let gauge_opts = Opts::new("test_gauge", "test help")
            .const_label("a", "1")
            .const_label("b", "2");
        let gauge = Gauge::with_opts(gauge_opts).unwrap();
        gauge.inc();
        gauge.set(42.0);

        let mf = gauge.collect();
        writer.clear();
        let txt = encoder.encode(&mf, &mut writer);
        assert!(txt.is_ok());
        let v: Value = serde_json::from_slice(&writer).unwrap();
        println!("{:?}", v);
        assert_eq!(v["test_gauge"]["help"], "test help");
        assert_eq!(v["test_gauge"]["a"], "1");
        assert_eq!(v["test_gauge"]["b"], "2");
        assert_eq!(v["test_gauge"]["type"], "gauge");
        assert_eq!(v["test_gauge"]["gauge"], json!(42.0));

    }

    
    #[test]
    fn test_text_encoder_histogram() {
        let opts = HistogramOpts::new("test_histogram", "test help").const_label("a", "1");
        let histogram = Histogram::with_opts(opts).unwrap();
        histogram.observe(0.25);

        let mf = histogram.collect();
        let mut writer = Vec::<u8>::new();
        let encoder = JsonEncoder::new();
        let res = encoder.encode(&mf, &mut writer);
        assert!(res.is_ok());

        let v: Value = serde_json::from_slice(&writer).unwrap();

        println!("{:?}", v);
        assert_eq!(v["test_histogram"]["help"], "test help");
        assert_eq!(v["test_histogram"]["type"], "histogram");
        assert_eq!(v["test_histogram"]["a"], "1");
        assert_eq!(v["test_histogram"]["counts"], json!(1));
        assert_eq!(v["test_histogram"]["sum"], json!(0.25));
        assert_eq!(v["test_histogram"]["cumulative_counts"], json!([0,0,0,0,0,1,1,1,1,1,1,1]));
        let array_test = [json!(0.005), json!(0.01), json!(0.025), json!(0.05), json!(0.1), json!(0.25), json!(0.5), json!(1.0), json!(2.5), json!(5.0), json!(10.0), json!(POSITIVE_INF)];
        assert_eq!(v["test_histogram"]["upper_bounds"], json!(array_test));


    }

    #[test]
    fn test_text_encoder_summary() {
        use prometheus::proto::{Metric, Quantile, Summary};
        use protobuf::RepeatedField;

        let mut metric_family = MetricFamily::default();
        metric_family.set_name("test_summary".to_string());
        metric_family.set_help("This is a test summary statistic".to_string());
        metric_family.set_field_type(MetricType::SUMMARY);

        let mut summary = Summary::default();
        summary.set_sample_count(5.0 as u64);
        summary.set_sample_sum(15.0);

        let mut quantile1 = Quantile::default();
        quantile1.set_quantile(50.0);
        quantile1.set_value(3.0);

        let mut quantile2 = Quantile::default();
        quantile2.set_quantile(100.0);
        quantile2.set_value(5.0);

        summary.set_quantile(RepeatedField::from_vec(vec!(quantile1, quantile2)));

        let mut metric = Metric::default();
        metric.set_summary(summary);
        metric_family.set_metric(RepeatedField::from_vec(vec!(metric)));

        let mut writer = Vec::<u8>::new();
        let encoder = JsonEncoder::new();
        let res = encoder.encode(&vec![metric_family], &mut writer);
        assert!(res.is_ok());
        
        let v: Value = serde_json::from_slice(&writer).unwrap();
        println!("{:?}", v);
        assert_eq!(v["test_summary"]["help"], "This is a test summary statistic");
        assert_eq!(v["test_summary"]["type"], "summary");
        assert_eq!(v["test_summary"]["quantiles"], json!(5));
        assert_eq!(v["test_summary"]["sum"], json!(15));

    }
 
    #[test]
    fn test_text_encoder_to_string() {
        let counter_opts = Opts::new("test_counter", "test help")
            .const_label("a", "1")
            .const_label("b", "2");
        let counter = Counter::with_opts(counter_opts).unwrap();
        counter.inc();

        let mf = counter.collect();

        let encoder = JsonEncoder::new();
        let txt = encoder.encode_to_string(&mf);
        let txt = txt.unwrap();
        let v : Value = serde_json::from_str(txt.as_str()).unwrap();


        assert_eq!(v["test_counter"]["help"], "test help");
        assert_eq!(v["test_counter"]["a"], "1");
        assert_eq!(v["test_counter"]["b"], "2");
        assert_eq!(v["test_counter"]["counter"], json!(1.0));
    }
    
}