/*! An easy-to-use prometheus-compatible metrics library

# Writing metrics

A "metric" is a named value of type `f64`.  There is a single global set of
metrics; you can update it from any thread.

```
use epimetheus::metric;

metric!(foobar).set(12.3);
metric!(foobar).add(0.7);
```

If you increment a metric which has never been set, it is considered to
start from zero.

```
# use epimetheus::metric;
metric!(barqux).add(6.5);
// now barqux = 6.5
```

## Labels

The base part of the name is fixed statically at at compile-time.  However,
a metric's name may also include "labels", which are dynamic.

```
# use epimetheus::metric;
let user_id = 7;
metric!(login_attempts{user=user_id}).add(1.0);
```

The label values can be anything which implements `Display`.

```
# use epimetheus::metric;
# let user_id = 0;
# let passwd = 0;
# let try_log_in = |_, _| 0;
// enum LoginResult { Success, BadUsername, BadPassword }
// impl Display for LoginResult { ... }

let result = try_log_in(user_id, passwd);
metric!(login_attempts{user=user_id, result=result}).add(1.0);
```

Labels can be useful, but they come at a performance cost (see README).

# Seeing your metrics

You need to call `spawn_http_server()` somewhere near the start of your
program.  (Personally I do it straight after I initialize the log backend.)
This will spawn a new thread which will serve metrics over HTTP.

```
epimetheus::spawn_http_server();
```

Connect to the server to see the current values of all metrics.  Metrics appear
in the output after being updated for the first time.

```text
$ curl localhost:9898
barqux 6.5
foobar 20
login_attempts{result="Success",user="7"} 1
login_attempts{user="7"} 1
```

By default the HTTP server runs on port 9898, but you can change this by
setting the `RUST_METRICS_PORT` environment variable.  Tip: If you want to
specify the metrics port in your application itself, you can do so like this:

```
std::env::set_var("RUST_METRICS_PORT", "1234");
epimetheus::spawn_http_server();
```

By the way, if you try to update a metric before calling `spawn_http_server()`,
the update is thrown away silently, and _quickly_.  This means that if you want
to disable your metrics, you can just omit the call to `spawn_http_server()`
and all your metric updates become fairly cheap no-ops.

```
# struct Opts { enable_metrics: bool }
# let opts = Opts { enable_metrics: false };
if opts.enable_metrics {
    epimetheus::spawn_http_server();
}
```
*/

use crossbeam_channel::Sender;
use log::*;
use once_cell::sync::OnceCell;
use std::collections::BTreeMap;
use std::fmt;
use std::io::prelude::*;
use std::net::{Ipv4Addr, SocketAddr};
use std::time::Duration;

struct State {
    chan: Sender<(Metric, Action)>,
    port: u16,
}

/// The global update channel.
static STATE: OnceCell<State> = OnceCell::new();

/// A named metric;  it has a associated global mutable `f64` value.
///
/// You can create these by hand, but you might find it more convenient to
/// use the `metric!()` macro.
pub struct Metric {
    pub name: &'static str,
    pub labels: Labels,
}
type Labels = Vec<(&'static str, Box<dyn fmt::Display + Send>)>;
enum Action {
    Inc(f64),
    Set(f64),
}

impl Metric {
    /// Set the metric to the specified value.
    #[inline]
    pub fn set(self, x: f64) {
        send_chan((self, Action::Set(x)));
    }

    /// Increment the metric by the specified amount.
    #[inline]
    pub fn add(self, x: f64) {
        send_chan((self, Action::Inc(x)));
    }
}

#[inline]
fn send_chan(x: (Metric, Action)) {
    // If there's no HTTP thread, drop the update
    if let Some(state) = STATE.get() {
        // If the HTTP thread dies then the Receiver will be dropped and
        // send() will return an error.  This is fine, so we ignore it.
        let _ = state.chan.send(x);
    }
}

/// Refer to a metric.
#[macro_export]
macro_rules! metric {
    ($name:ident) => {
        epimetheus::Metric {
            name: stringify!($name),
            labels: Vec::new(),
        }
    };
    ($name:ident{$($key:ident = $val:expr),*}) => {
        epimetheus::Metric {
            name: stringify!($name),
            labels: vec![$((stringify!($key), Box::new($val))),*],
        }
    };
}

#[derive(PartialEq, Eq, PartialOrd, Ord)]
struct LabelsDisplay(BTreeMap<&'static str, String>);
impl LabelsDisplay {
    fn new(labels: Labels) -> LabelsDisplay {
        let labels = labels
            .into_iter()
            .map(|(k, v)| (k, v.to_string()))
            .collect::<BTreeMap<_, _>>();
        LabelsDisplay(labels)
    }
}
impl fmt::Display for LabelsDisplay {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut labels = self.0.iter();
        if let Some(head) = labels.next() {
            f.write_str("{")?;
            f.write_str(head.0)?;
            f.write_str("=")?;
            fmt::Debug::fmt(head.1, f)?;
            for (k, v) in labels {
                f.write_str(",")?;
                f.write_str(k)?;
                f.write_str("=")?;
                fmt::Debug::fmt(v, f)?;
            }
            f.write_str("}")?;
        }
        Ok(())
    }
}

/// Get the port of a running server if it exists.
pub fn existing_server_port() -> Option<u16> {
    STATE.get().map(|s| s.port)
}

pub fn spawn_http_server() {
    if let Err(e) = try_spawn_http_server() {
        warn!("HTTP thread failed to start: {}", e);
    }
}

fn try_spawn_http_server() -> Result<(), impl std::error::Error> {
    let (tx, rx) = crossbeam_channel::unbounded();

    // This thread stores the current value of all the metrics in a map.
    // (We use a BTreeMap so the metrics are nicely sorted when we print them.)
    let mut metrics = BTreeMap::<(&'static str, LabelsDisplay), f64>::default();

    // We're going to bind a socket and listen for incoming connections from
    // HTTP clients.  When we get a connection, we'll:
    //
    // 1. drain any updates from the channel and apply them to the local
    //    metrics map; and then
    // 2. render the map to prometheus exposition format and send it to
    //    the client.
    //
    // However, we want to ensure that the channel is regularly drained, even
    // when there are no new connections coming in.  (Otherwise, we'd have
    // a memory leak and - worse - the metric update latency would suffer.)
    // Therefore we set a timeout while we wait, and drain the channel when
    // someone connects or a timeout occurs.
    let port = std::env::var("RUST_METRICS_PORT")
        .ok()
        .and_then(|x| x.parse::<u16>().ok())
        .unwrap_or(9898);
    STATE
        .set(State { chan: tx, port })
        .map_err(|_| ())
        .expect("run_http_server() was called twice!");

    use mio::{net::TcpListener, Events, Interest, Poll, Token};
    let mut sock = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, port)))?;
    let mut poll = Poll::new()?;
    poll.registry()
        .register(&mut sock, Token(0), Interest::READABLE)?;
    info!("Listening on port {}", port);

    let mut events = Events::with_capacity(64);
    std::thread::spawn(move || loop {
        poll.poll(&mut events, Some(Duration::from_secs(20)))
            .unwrap();
        // No point reading `events` - there's only one thing it could be.
        events.clear();
        for (metric, action) in rx.try_iter() {
            let labels = LabelsDisplay::new(metric.labels);
            match action {
                Action::Inc(x) => *metrics.entry((metric.name, labels)).or_insert(0.0) += x,
                Action::Set(x) => *metrics.entry((metric.name, labels)).or_insert(0.0) = x,
            }
        }
        while let Ok((conn, _)) = sock.accept() {
            if let Err(e) = handle_http_client(conn, &metrics) {
                warn!("{}", e);
            }
        }
    });
    Ok::<(), std::io::Error>(())
}

/// This is the world's simplest HTTP implementation.  It completely
/// ignores the request, unconditionally sending the same response.
/// This response comes with no headers or anything - just a body.
fn handle_http_client(
    mut conn: mio::net::TcpStream,
    metrics: &BTreeMap<(&'static str, LabelsDisplay), f64>,
) -> Result<(), std::io::Error> {
    // We don't care about the request, but some HTTP clients get
    // upset if you don't at least read it.  Unfortunate.
    let mut buf = [0; 128];
    loop {
        // conn is non-blocking, so we need to retry if we get EAGAIN
        match conn.read(&mut buf) {
            Ok(_) => break,
            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => (),
            Err(e) => return Err(e),
        }
    }
    writeln!(conn, "HTTP/1.1 200 OK\r\n")?;
    for ((name, labels), val) in metrics {
        writeln!(conn, "{}{} {}", name, labels, val)?;
    }
    Ok(())
}