use axum::{Json, Router, extract::State, routing::get};
use native_tls::TlsConnector;
use printstats::*;
use reqwest::Client;
use rumqttc::{AsyncClient, Event, MqttOptions, Packet, QoS, Transport};
use std::collections::HashMap;
use std::fs;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;

#[tokio::main]
async fn main() {
    let content = fs::read_to_string("config.toml").expect("Couldn't read config.toml");
    let config = Config::load(&content);
    let state: Arc<Mutex<HashMap<String, PrinterState>>> = Arc::new(Mutex::new(HashMap::new()));

    for printer in config.printers {
        match printer {
            Printer::Prusa { name, host, api_key, .. } => {
                let state_clone = state.clone();
                tokio::spawn(poll_prusa(name, host, api_key, state_clone));
            }
            Printer::Bambu { name, host, serial_number, access_code } => {
                println!("Found Bambu: {} at {}", name, host);
                let state_clone = Arc::clone(&state);
                tokio::spawn(poll_bambu(name, host, serial_number, access_code, state_clone));
            }
        }
    }

    tracing_subscriber::fmt::init();

    let app = Router::new()
        .route("/", get(root))
        .with_state(Arc::clone(&state));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    let _ = axum::serve(listener, app).await;
}

async fn poll_prusa(
    name: String,
    host: String,
    api_key: String,
    state: Arc<Mutex<HashMap<String, PrinterState>>>,
) {
    let client = Client::new();
    loop {
        let result: Result<(), Box<dyn std::error::Error + Send + Sync>> = async {
            let response = client
                .get(format!("http://{}/api/v1/status", host))
                .header("X-Api-Key", &api_key)
                .send()
                .await?
                .json::<PrusaStatus>()
                .await?;
            let mut lock = state.lock().await;
            let entry = lock.entry(name.clone()).or_default();
            extract_status_from_prusa(&response, entry);
            Ok(())
        }
        .await;
        if let Err(e) = result {
            eprintln!("Error polling Prusa printer {}: {}", name, e);
        }
        tokio::time::sleep(Duration::from_secs(5)).await;
    }
}

async fn poll_bambu(
    name: String,
    host: String,
    serial_number: String,
    access_code: String,
    state: Arc<Mutex<HashMap<String, PrinterState>>>,
) {
    let mut mqttoptions = MqttOptions::new(&name, &host, 8883);
    mqttoptions.set_keep_alive(Duration::from_secs(5));
    mqttoptions.set_credentials("bblp", &access_code);

    // Bambu printers use TLS with a self-signed certificate
    let connector = TlsConnector::builder()
        .danger_accept_invalid_certs(true)
        .build()
        .unwrap();
    mqttoptions.set_transport(Transport::tls_with_config(connector.into()));

    let (client, mut eventloop) = AsyncClient::new(mqttoptions, 10);
    client
        .subscribe(format!("device/{}/report", serial_number), QoS::AtMostOnce)
        .await
        .unwrap();

    loop {
        // eventloop.poll() yields back to Tokio when there's no data
        match eventloop.poll().await {
            Ok(Event::Incoming(Packet::Publish(p))) => {
                println!("{:?}", &p.payload);
                match serde_json::from_slice::<BambuStatus>(&p.payload) {
                    Ok(msg) => {
                        let mut lock = state.lock().await;
                        let entry = lock.entry(name.clone()).or_default();
                        extract_status_from_bambu(&msg, entry);
                        println!("Updated state for {}: {:?}", &name, p.payload);
                    }
                    Err(e) => eprintln!("Failed to deserialize BambuStatus: {}", e),
                }
            }
            Ok(_) => {}
            Err(e) => {
                eprintln!("MQTT error: {:?}", e);
                tokio::time::sleep(Duration::from_secs(5)).await; // Simple retry
            }
        }
    }
}

async fn root(
    State(state): State<Arc<Mutex<HashMap<String, PrinterState>>>>,
) -> Json<HashMap<String, PrinterState>> {
    let lock = state.lock().await;
    Json(lock.clone())
}