//! A minimal CoAP server over UDP.
//!
//! CoAP (RFC 7252) is "HTTP semantics on a UDP datagram": same methods, paths
//! and query strings, but each request/response is a single datagram carrying a
//! compact binary header. `coap-lite` handles that binary codec; the UDP socket
//! and the routing are right here so nothing is hidden.

use std::net::SocketAddr;
use std::time::Instant;

use coap_lite::{
    CoapOption, CoapRequest, CoapResponse, ContentFormat, Packet, RequestType,
    ResponseType,
};
use tokio::net::UdpSocket;

/// Max UDP payload we accept. RFC 7252 recommends staying under the ~1152 B
/// that fits an unfragmented IPv6 datagram; anything larger needs block-wise
/// transfer (RFC 7959).
const MAX_DATAGRAM: usize = 1152;

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let host = std::env::var("COAP_HOST").unwrap_or_else(|_| "0.0.0.0".into());
    let port = std::env::var("COAP_PORT").unwrap_or_else(|_| "5683".into());
    let addr = format!("{host}:{port}");

    let socket = UdpSocket::bind(&addr).await?;
    let started = Instant::now();

    println!("CoAP server (UDP) listening on coap://{addr}");
    println!("  GET  /hello?name=Imran");
    println!("  POST /hello   (payload: Imran)");
    println!("  GET  /health");

    let mut buf = [0u8; MAX_DATAGRAM];

    loop {
        tokio::select! {
            _ = tokio::signal::ctrl_c() => {
                println!("\nshutting down");
                return Ok(());
            }
            result = socket.recv_from(&mut buf) => {
                let (len, peer) = match result {
                    Ok(v) => v,
                    Err(e) => { eprintln!("recv error: {e}"); continue; }
                };

                // Handling is pure CPU work with no awaits, so doing it inline
                // keeps the hot path allocation-light and ordered. Spawn a task
                // here instead once a handler needs to do real I/O.
                if let Some(reply) = handle(&buf[..len], peer, started)
                    && let Err(e) = socket.send_to(&reply, peer).await
                {
                    eprintln!("send error to {peer}: {e}");
                }
            }
        }
    }
}

/// Decode one datagram, route it, and encode the reply.
/// Returns `None` when no reply should be sent (unparsable or non-request).
fn handle(datagram: &[u8], peer: SocketAddr, started: Instant) -> Option<Vec<u8>> {
    let packet = match Packet::from_bytes(datagram) {
        Ok(p) => p,
        Err(e) => {
            eprintln!("malformed packet from {peer}: {e}");
            return None;
        }
    };

    let mut request = CoapRequest::from_packet(packet, peer);

    // `None` means the datagram was an ACK/RST, not something we answer.
    let mut response = CoapResponse::new(&request.message)?;

    let path = format!("/{}", request.get_path());
    let method = *request.get_method();
    println!("{method:?} {}{} from {peer}", path, query_string(&request.message));

    route(method, &path, &mut request, &mut response, started);

    match response.message.to_bytes() {
        Ok(bytes) => Some(bytes),
        Err(e) => {
            eprintln!("encode error: {e}");
            None
        }
    }
}

fn route(
    method: RequestType,
    path: &str,
    request: &mut CoapRequest<SocketAddr>,
    response: &mut CoapResponse,
    started: Instant,
) {
    match (method, path) {
        (RequestType::Get, "/hello") => {
            let name = query_param(&request.message, "name");
            reply(response, ResponseType::Content, ContentFormat::TextPlain, greet(name.as_deref()));
        }

        (RequestType::Post, "/hello") => {
            let name = name_from_payload(&request.message.payload)
                .or_else(|| query_param(&request.message, "name"));

            match name {
                Some(name) => reply(
                    response,
                    ResponseType::Created,
                    ContentFormat::TextPlain,
                    greet(Some(&name)),
                ),
                None => reply(
                    response,
                    ResponseType::BadRequest,
                    ContentFormat::TextPlain,
                    "Bad Request: send a name as the payload".into(),
                ),
            }
        }

        (RequestType::Get, "/health") => {
            let body = format!(
                r#"{{"status":"ok","uptime":{}}}"#,
                started.elapsed().as_secs()
            );
            reply(response, ResponseType::Content, ContentFormat::ApplicationJSON, body);
        }

        // Path exists but the method doesn't: 4.05, per RFC 7252.
        (_, "/hello") | (_, "/health") => reply(
            response,
            ResponseType::MethodNotAllowed,
            ContentFormat::TextPlain,
            format!("Method Not Allowed: {method:?} {path}"),
        ),

        _ => reply(
            response,
            ResponseType::NotFound,
            ContentFormat::TextPlain,
            format!("Not Found: {path}"),
        ),
    }
}

fn greet(name: Option<&str>) -> String {
    let name = name.map(str::trim).filter(|n| !n.is_empty()).unwrap_or("World");
    format!("Hello, {name}!")
}

fn reply(
    response: &mut CoapResponse,
    status: ResponseType,
    format: ContentFormat,
    body: String,
) {
    response.set_status(status);
    response.message.set_content_format(format);
    response.message.payload = body.into_bytes();
}

/// Re-join the URI-Query options for logging, e.g. `?name=Imran`.
fn query_string(packet: &Packet) -> String {
    match packet.get_option(CoapOption::UriQuery) {
        Some(options) if !options.is_empty() => {
            let pairs: Vec<String> = options
                .iter()
                .map(|o| String::from_utf8_lossy(o).into_owned())
                .collect();
            format!("?{}", pairs.join("&"))
        }
        _ => String::new(),
    }
}

/// Read one `?key=value` pair out of the URI-Query options.
fn query_param(packet: &Packet, key: &str) -> Option<String> {
    let options = packet.get_option(CoapOption::UriQuery)?;
    for option in options.iter() {
        let pair = String::from_utf8_lossy(option);
        if let Some(value) = pair.strip_prefix(&format!("{key}=")) {
            return Some(percent_decode(value));
        }
    }
    None
}

/// Accept a raw-text payload, or `{"name":"Imran"}`.
fn name_from_payload(payload: &[u8]) -> Option<String> {
    let body = String::from_utf8_lossy(payload).trim().to_string();
    if body.is_empty() {
        return None;
    }
    if body.starts_with('{') {
        return json_string_field(&body, "name");
    }
    Some(body)
}

/// Pull one top-level string field out of flat JSON. Enough for an MVP body
/// like `{"name":"Imran"}`; swap in `serde_json` when the shape grows.
fn json_string_field(body: &str, field: &str) -> Option<String> {
    let after_key = body.split_once(&format!("\"{field}\""))?.1;
    let after_colon = after_key.split_once(':')?.1;
    let start = after_colon.find('"')? + 1;
    let value = &after_colon[start..];
    let end = value.find('"')?;
    Some(value[..end].to_string())
}

/// A CoAP client is supposed to percent-decode before splitting a URI into
/// options, but not all of them do — decode leniently so `?name=Md%20Imran`
/// works either way. Invalid escapes are left as-is.
fn percent_decode(input: &str) -> String {
    let bytes = input.as_bytes();
    let mut out = Vec::with_capacity(bytes.len());
    let mut i = 0;

    while i < bytes.len() {
        match bytes[i] {
            b'%' if i + 2 < bytes.len() => {
                let hex = std::str::from_utf8(&bytes[i + 1..i + 3]).ok();
                match hex.and_then(|h| u8::from_str_radix(h, 16).ok()) {
                    Some(byte) => {
                        out.push(byte);
                        i += 3;
                    }
                    None => {
                        out.push(bytes[i]);
                        i += 1;
                    }
                }
            }
            b'+' => {
                out.push(b' ');
                i += 1;
            }
            byte => {
                out.push(byte);
                i += 1;
            }
        }
    }

    String::from_utf8_lossy(&out).into_owned()
}

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

    #[test]
    fn greets_world_when_no_name() {
        assert_eq!(greet(None), "Hello, World!");
        assert_eq!(greet(Some("")), "Hello, World!");
        assert_eq!(greet(Some("   ")), "Hello, World!");
    }

    #[test]
    fn greets_the_given_name() {
        assert_eq!(greet(Some("Imran")), "Hello, Imran!");
        assert_eq!(greet(Some("  Imran  ")), "Hello, Imran!");
    }

    #[test]
    
    fn reads_name_from_text_and_json_payloads() {
        assert_eq!(name_from_payload(b"Imran").as_deref(), Some("Imran"));
        assert_eq!(
            name_from_payload(br#"{"name":"Imran"}"#).as_deref(),
            Some("Imran")
        );
        assert_eq!(
            name_from_payload(br#"{ "name" : "Md Imran" }"#).as_deref(),
            Some("Md Imran")
        );
        assert_eq!(name_from_payload(b""), None);
        assert_eq!(name_from_payload(b"   "), None);
        assert_eq!(name_from_payload(br#"{"other":"x"}"#), None);
    }

    #[test]
    fn decodes_percent_escapes() {
        assert_eq!(percent_decode("Md%20Imran"), "Md Imran");
        assert_eq!(percent_decode("Md+Imran"), "Md Imran");
        assert_eq!(percent_decode("Imran"), "Imran");
        // Invalid escapes are passed through rather than dropped.
        assert_eq!(percent_decode("100%"), "100%");
        assert_eq!(percent_decode("%zz"), "%zz");
    }

    #[test]
    fn reads_query_param_from_uri_query_options() {
        let mut packet = Packet::new();
        packet.add_option(CoapOption::UriQuery, b"name=Md%20Imran".to_vec());
        packet.add_option(CoapOption::UriQuery, b"lang=en".to_vec());

        assert_eq!(query_param(&packet, "name").as_deref(), Some("Md Imran"));
        assert_eq!(query_param(&packet, "lang").as_deref(), Some("en"));
        assert_eq!(query_param(&packet, "missing"), None);
        assert_eq!(query_string(&packet), "?name=Md%20Imran&lang=en");
    }
}