//! Minimal RouterOS API client for graceful MikroTik power-off.
//!
//! Protocol per the official MikroTik "Python3 Example" template
//! (help.mikrotik.com, updated 2026-02): plain TCP (default port 8728),
//! variable-length word framing (no flags, no CRC encoding), sentences
//! terminated by an empty word. Login sends `/login =name= =password=`;
//! modern RouterOS answers `!done` outright, pre-6.43 answers
//! `!done =ret=<hex>` and expects an MD5 challenge response.
//!
//! We only need one command: `/system/shutdown`.

use anyhow::{Context, Result};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::TcpStream;
use tracing::info;

use crate::config::MikrotikDevice;

/// Liveness probe for `health`: connect + login, nothing else.
/// `Ok` = the router is up AND the configured credentials work;
/// any failure (unreachable, timeout, rejected login) = the entity is off,
/// exactly like an HA switch that does not answer.
pub async fn ping_device(dev: &MikrotikDevice) -> Result<()> {
    let target = format!("{}:{}", dev.address, dev.api_port);
    let mut s = tokio::time::timeout(Duration::from_secs(5), TcpStream::connect(&target))
        .await
        .with_context(|| format!("connecting to RouterOS at {}", target))?
        .with_context(|| format!("connect to RouterOS at {}", target))?;
    s.set_nodelay(true)?;
    tokio::time::timeout(Duration::from_secs(10), login_stream(&mut s, dev))
        .await
        .map_err(|_| anyhow::anyhow!("RouterOS login timed out ({})", dev.name))?
}

/// Connect + login, then shut the router down via `/system/shutdown`.
/// One connection per call — the router is about to die anyway.
pub async fn shutdown_device(dev: &MikrotikDevice) -> Result<()> {
    let target = format!("{}:{}", dev.address, dev.api_port);
    let mut s = tokio::time::timeout(Duration::from_secs(10), TcpStream::connect(&target))
        .await
        .with_context(|| format!("connecting to RouterOS at {}", target))?
        .with_context(|| format!("connect to RouterOS at {}", target))?;
    s.set_nodelay(true)?;
    login_stream(&mut s, dev).await?;
    shutdown_stream(&mut s, dev).await
}

/// The login exchange on an authenticated-or-not stream: `/login` with
/// username + password, then `!done` (modern) or the legacy MD5
/// challenge round-trip (pre-6.43).
async fn login_stream<S>(s: &mut S, dev: &MikrotikDevice) -> Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    send_sentence(
        s,
        &[
            "/login",
            &format!("=name={}", dev.username),
            &format!("=password={}", dev.password),
        ],
    )
    .await
    .context("sending /login")?;

    loop {
        let mut words = read_sentence(s).await.context("reading /login reply")?;
        if words.first().is_some_and(|w| w == "#routeros") {
            words.remove(0); // tolerate the legacy banner word
        }
        let rest: Vec<&str> = words.iter().skip(1).map(|w| w.as_str()).collect();
        match words.first().map(|s| s.as_str()) {
            Some("!trap") => {
                anyhow::bail!("RouterOS login rejected ({}): {}", dev.name, rest.join(" "))
            }
            Some("!done") => {
                let Some(hex) = rest.iter().find_map(|w| w.strip_prefix("=ret=")) else {
                    return Ok(()); // no challenge — logged in
                };
                let challenge = hex_decode(hex)
                    .map_err(|e| anyhow::anyhow!("login challenge is not hex: {hex} ({e})"))?;
                let mut pre = vec![0u8];
                pre.extend_from_slice(dev.password.as_bytes());
                pre.extend_from_slice(&challenge);
                let md = md5::compute(&pre);
                send_sentence(
                    s,
                    &[
                        "/login",
                        &format!("=name={}", dev.username),
                        &format!("=response=00{}", hex_encode(md.as_slice())),
                    ],
                )
                .await
                .context("sending login challenge response")?;
            }
            _ => anyhow::bail!(
                "unexpected RouterOS reply during login ({}): {}",
                dev.name,
                words.join(" ")
            ),
        }
    }
}

/// `/system/shutdown` on a logged-in stream; read replies until
/// `!done` (`!trap` = the router refused, e.g. policy).
async fn shutdown_stream<S>(s: &mut S, dev: &MikrotikDevice) -> Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    info!("RouterOS graceful shutdown: {}", dev.name);
    let r = tokio::time::timeout(Duration::from_secs(10), async {
        send_sentence(s, &["/system/shutdown"])
            .await
            .context("sending /system/shutdown")?;
        loop {
            let mut words = read_sentence(s).await.context("reading shutdown reply")?;
            if words.first().is_some_and(|w| w == "#routeros") {
                words.remove(0);
            }
            match words.first().map(|s| s.as_str()) {
                Some("!trap") => {
                    anyhow::bail!(
                        "RouterOS {} refused /system/shutdown: {}",
                        dev.name,
                        words
                            .iter()
                            .skip(1)
                            .map(|w| w.as_str())
                            .collect::<Vec<_>>()
                            .join(" ")
                    );
                }
                Some("!done") => break,
                _ => {} // interim `.*` progress lines
            }
        }
        Ok::<(), anyhow::Error>(())
    })
    .await;
    r.context("RouterOS shutdown timed out (10 s)")?
}

// --- Framing -------------------------------------------------------------
//
// Word = variable-length length prefix + raw UTF-8 bytes (no flags).
//   <0x80                1-byte length
//   0x80..0xC0           (len | 0x8000),      2 bytes total
//   0xC0..0xE0           (len | 0xC00000),    3 bytes total
//   0xE0..0xF0           (len | 0xE0000000),  4 bytes total
//   0xF0..               5 bytes total, first is a flag only

async fn write_word<W: AsyncWrite + Unpin>(w: &mut W, s: &str) -> std::io::Result<()> {
    let data = s.as_bytes();
    let len = data.len();
    let mut hdr = [0u8; 5];
    let n = if len < 0x80 {
        hdr[0] = len as u8;
        1
    } else if len < 0x4000 {
        let v = (len | 0x8000) as u16;
        hdr[0..2].copy_from_slice(&v.to_be_bytes());
        2
    } else if len < 0x200000 {
        let v = (len | 0xC00000) as u32;
        hdr[0..3].copy_from_slice(&v.to_be_bytes()[1..]);
        3
    } else if len < 0x10000000 {
        let v = (len | 0xE0000000) as u32;
        hdr[0..4].copy_from_slice(&v.to_be_bytes());
        4
    } else {
        hdr[0] = 0xF0;
        hdr[1..5].copy_from_slice(&len.to_be_bytes());
        5
    };
    w.write_all(&hdr[..n]).await?;
    w.write_all(data).await?;
    Ok(())
}

async fn read_u8<R: AsyncRead + Unpin>(r: &mut R) -> std::io::Result<u32> {
    let mut x = [0u8; 1];
    r.read_exact(&mut x).await?;
    Ok(u32::from(x[0]))
}

async fn read_word_len<R: AsyncRead + Unpin>(r: &mut R) -> std::io::Result<u32> {
    let mut c = [0u8; 1];
    r.read_exact(&mut c).await?;
    Ok(if c[0] < 0x80 {
        u32::from(c[0])
    } else if c[0] < 0xC0 {
        ((c[0] & 0x3F) as u32) << 8 | read_u8(r).await?
    } else if c[0] < 0xE0 {
        ((c[0] & 0x1F) as u32) << 16 | read_u8(r).await? << 8 | read_u8(r).await?
    } else if c[0] < 0xF0 {
        ((c[0] & 0x0F) as u32) << 24
            | read_u8(r).await? << 16
            | read_u8(r).await? << 8
            | read_u8(r).await?
    } else {
        read_u8(r).await? << 24
            | read_u8(r).await? << 16
            | read_u8(r).await? << 8
            | read_u8(r).await?
    })
}

async fn read_word<R: AsyncRead + Unpin>(r: &mut R) -> std::io::Result<Vec<u8>> {
    let len = read_word_len(r).await?;
    let mut buf = vec![0u8; len as usize];
    r.read_exact(&mut buf).await?;
    Ok(buf)
}

/// One sentence = words + a terminating empty word.
async fn send_sentence<W: AsyncWrite + Unpin>(w: &mut W, words: &[&str]) -> std::io::Result<()> {
    for word in words {
        write_word(w, word).await?;
    }
    write_word(w, "").await
}

async fn read_sentence<R: AsyncRead + Unpin>(r: &mut R) -> std::io::Result<Vec<String>> {
    let mut words = Vec::new();
    loop {
        let raw = read_word(r).await?;
        if raw.is_empty() {
            return Ok(words);
        }
        words.push(String::from_utf8_lossy(&raw).into_owned());
    }
}

fn hex_decode(s: &str) -> std::result::Result<Vec<u8>, String> {
    let bytes = s.as_bytes();
    if !bytes.len().is_multiple_of(2) {
        return Err("odd number of hex digits".to_string());
    }
    let val = |c: u8| -> std::result::Result<u8, String> {
        match c {
            b'0'..=b'9' => Ok(c - b'0'),
            b'a'..=b'f' => Ok(c - b'a' + 10),
            b'A'..=b'F' => Ok(c - b'A' + 10),
            _ => Err("not a hex digit".to_string()),
        }
    };
    (0..bytes.len())
        .step_by(2)
        .map(|i| Ok(val(bytes[i])? * 16 + val(bytes[i + 1])?))
        .collect()
}

fn hex_encode(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        s.push_str(&format!("{:02x}", b));
    }
    s
}

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

    fn device() -> MikrotikDevice {
        MikrotikDevice {
            name: "t".into(),
            address: "127.0.0.1".into(),
            username: "admin".into(),
            password: "secret".into(),
            api_port: 8728,
            on_entity: None,
        }
    }

    #[tokio::test]
    async fn framing_roundtrip_all_bands() {
        for n in [0usize, 1, 127, 128, 16383, 16384, 2_097_151] {
            // duplex is a bounded pipe — write and read must run concurrently
            let (mut a, mut b) = duplex(64 * 1024);
            let s = "x".repeat(n);
            let s2 = s.clone();
            let writer = tokio::spawn(async move { write_word(&mut a, &s2).await });
            let reader = tokio::spawn(async move { read_word(&mut b).await });
            let got = reader.await.unwrap().unwrap();
            writer.await.unwrap().unwrap();
            assert_eq!(got.len(), n, "size {n}");
            assert_eq!(String::from_utf8(got).unwrap(), s);
        }
    }

    #[tokio::test]
    async fn framing_band5_roundtrip() {
        // 2 MiB+ payload exercises the 5-byte header band.
        let (mut a, mut b) = duplex(64 * 1024);
        let s = "y".repeat(2_100_000);
        let s2 = s.clone();
        let writer = tokio::spawn(async move { write_word(&mut a, &s2).await });
        let reader = tokio::spawn(async move { read_word(&mut b).await });
        let got = reader.await.unwrap().unwrap();
        writer.await.unwrap().unwrap();
        assert_eq!(got.len(), 2_100_000);
    }

    #[tokio::test]
    async fn sentence_terminates_on_empty_word() {
        let (mut a, mut b) = duplex(64 * 1024);
        tokio::spawn(async move {
            send_sentence(&mut a, &["/login", "=name=admin"])
                .await
                .unwrap();
        });
        let words = read_sentence(&mut b).await.unwrap();
        assert_eq!(words, vec!["/login", "=name=admin"]);
    }

    /// Client login exchange against a mock RouterOS that answers modern
    /// style: `!done` with no challenge.
    #[tokio::test]
    async fn login_modern_ok() {
        let (mut c, mut s) = duplex(64 * 1024);
        tokio::spawn(async move {
            let sent = read_sentence(&mut s).await.unwrap();
            assert_eq!(sent, vec!["/login", "=name=admin", "=password=secret"]);
            send_sentence(&mut s, &["!done"]).await.unwrap();
        });
        login_stream(&mut c, &device()).await.unwrap();
    }

    /// Legacy challenge: `!done =ret=<hex>` → client must answer with the
    /// MD5 response, then `!done`.
    #[tokio::test]
    async fn login_legacy_challenge() {
        let (mut c, mut s) = duplex(64 * 1024);
        tokio::spawn(async move {
            let sent = read_sentence(&mut s).await.unwrap();
            assert_eq!(sent[0], "/login");
            send_sentence(&mut s, &["!done", "=ret=deadbeef"])
                .await
                .unwrap();
            let resp = read_sentence(&mut s).await.unwrap();
            // expected: md5(0x00 || "secret" || 0xde 0xad 0xbe 0xef)
            let mut pre = vec![0u8];
            pre.extend_from_slice(b"secret");
            pre.extend_from_slice(&[0xde, 0xad, 0xbe, 0xef]);
            let md = md5::compute(&pre);
            let want = format!("=response=00{}", hex_encode(md.as_slice()));
            assert_eq!(
                resp,
                vec!["/login".to_string(), "=name=admin".to_string(), want],
                "challenge response"
            );
            send_sentence(&mut s, &["!done"]).await.unwrap();
        });
        login_stream(&mut c, &device()).await.unwrap();
    }

    #[tokio::test]
    async fn login_trap_fails() {
        let (mut c, mut s) = duplex(64 * 1024);
        tokio::spawn(async move {
            let _ = read_sentence(&mut s).await;
            send_sentence(&mut s, &["!trap", "=message=wrong pass"])
                .await
                .unwrap();
        });
        let e = login_stream(&mut c, &device()).await.unwrap_err();
        assert!(e.to_string().contains("wrong pass"), "{}", e);
    }

    #[tokio::test]
    async fn shutdown_done() {
        let (mut c, mut s) = duplex(64 * 1024);
        tokio::spawn(async move {
            let sent = read_sentence(&mut s).await.unwrap();
            assert_eq!(sent, vec!["/system/shutdown"]);
            send_sentence(&mut s, &["!done"]).await.unwrap();
        });
        shutdown_stream(&mut c, &device()).await.unwrap();
    }

    #[tokio::test]
    async fn shutdown_trap_fails() {
        let (mut c, mut s) = duplex(64 * 1024);
        tokio::spawn(async move {
            let _ = read_sentence(&mut s).await;
            send_sentence(&mut s, &["!trap", "=message=denied"])
                .await
                .unwrap();
        });
        let e = shutdown_stream(&mut c, &device()).await.unwrap_err();
        assert!(e.to_string().contains("denied"), "{}", e);
    }

    #[test]
    fn hex_roundtrip() {
        let bytes: Vec<u8> = vec![0xde, 0xad, 0x00, 0xff, 0x01];
        assert_eq!(hex_decode(&hex_encode(&bytes)).unwrap(), bytes);
        assert!(hex_decode("abc").is_err());
        assert!(hex_decode("zz").is_err());
    }
}