use opentelemetry::{global, metrics::Counter, KeyValue};
use prometheus::Registry;
use rocket::{fairing::{Fairing, Info, Kind}, get, routes, Request, Response, State};
use tracing::error;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
use std::process;
use url::Url;

struct MetricsState {
    counter: Counter<u64>,
}

#[rocket::main]
async fn main() -> Result<(), tracing_loki::Error> {
    let logs_provider = init_logs_provider("http://127.0.0.1:3100");
    let (meter_provider, registry) = init_meter_provider();
    
    // Create a meter from the above MeterProvider.
    let meter = global::meter("mylibraryname");

    // Create a Counter Instrument.
    let counter = meter.u64_counter("my_counter").init();

    let _ = rocket::build()
        .mount("/", routes![metrics, hello_world])
        .manage(registry)
        .attach(MetricsState { counter })
        .launch().await;

    meter_provider.shutdown().unwrap();
    logs_provider.abort();
    Ok(())
}

#[get("/")]
fn hello_world() -> &'static str {
    "Hello world." 
}

// "/metrics" Is where prometheus expects to gather metrics at.
// If you change this make sure your prometheus config reflects the change.
#[get("/metrics")]
fn metrics(state: &State<Registry>) -> String {
    let mut buffer = String::new();
    let encoder = prometheus::TextEncoder::new();
    let metric_families = state.gather();
    encoder.encode_utf8(&metric_families, &mut buffer).unwrap();
    buffer
}

fn init_logs_provider(url: &str) -> tokio::task::JoinHandle<()> {
    // Setup logs
    let (layer, task) = tracing_loki::builder()
        /*
         * Labels are for static things
         * Fields are for dynamic things
         */
        .label("program", "webserver").expect("Invalid log label name")
        .extra_field("version", "v1.0").expect("Invalid log field name")
        .extra_field("pid", format!("{}", process::id())).expect("Invalid log field name")
        .build_url(Url::parse(url).unwrap()).expect("Invalid log url endpoint");
    tracing_subscriber::registry()
        .with(layer)
        // Stdout if disired:
        // .with(tracing_subscriber::fmt::Layer::new())
        .init();
    tokio::spawn(task)
}

fn init_meter_provider() -> (opentelemetry_sdk::metrics::SdkMeterProvider, Registry) {
    use opentelemetry_sdk::metrics::SdkMeterProvider;

    let registry = Registry::new();
    let exporter = opentelemetry_prometheus::exporter()
        .with_registry(registry.clone())
        .build().unwrap();
    let provider = SdkMeterProvider::builder()
        .with_reader(exporter)
//        .with_resource(Resource::new([KeyValue::new(
//            "service.name",
//            "metrics-basic-example",
//        )]))
        .build();
    global::set_meter_provider(provider.clone());
    (provider, registry)
}


#[rocket::async_trait]
impl Fairing for MetricsState {
    fn info(&self) -> Info {
        Info {
            name: "Route metris provider",
            kind: Kind::Response
        }
    }

    async fn on_response<'r>(&self, req: &'r Request<'_>, res: &mut Response<'r>) {
        let path = req.uri().path().as_str().to_owned();
        let method = req.method().as_str();
        let status = res.status().to_string();
        let query = req.query_fields().map(|f| format!("{}={},", f.name, f.value)).collect::<String>();

        self.counter.add(
            1, &[
                KeyValue::new("Path", path),
                KeyValue::new("Method", method),
                KeyValue::new("Status", status),
                KeyValue::new("Query", query),
            ], );
        }
}