//! Reconnect coverage for socket types beyond PUSH/PULL.

use std::net::Ipv4Addr;
use std::time::Duration;

use omq_compio::endpoint::Host;
use omq_compio::options::ReconnectPolicy;
use omq_compio::{Endpoint, Message, Options, Socket, SocketType};

fn tcp_ep(port: u16) -> Endpoint {
    Endpoint::Tcp {
        host: Host::Ip(Ipv4Addr::LOCALHOST.into()),
        port,
    }
}

fn fast_reconnect() -> Options {
    Options {
        reconnect: ReconnectPolicy::Fixed(Duration::from_millis(30)),
        ..Default::default()
    }
}

async fn rebind<F: Fn() -> Socket>(ep: &Endpoint, make: F) -> Socket {
    let s = make();
    for _ in 0..40 {
        if s.bind(ep.clone()).await.is_ok() {
            return s;
        }
        compio::time::sleep(Duration::from_millis(25)).await;
    }
    panic!("could not rebind {ep:?} after 40 attempts");
}

// ── REQ / REP ────────────────────────────────────────────────────────────────

#[compio::test]
async fn req_rep_reconnect_after_server_restart() {
    let rep1 = Socket::new(SocketType::Rep, Options::default());
    let ep = rep1.bind(tcp_ep(0)).await.unwrap();

    let req = Socket::new(SocketType::Req, fast_reconnect());
    req.connect(ep.clone()).await.unwrap();

    req.send(Message::single("ping1")).await.unwrap();
    let got = compio::time::timeout(Duration::from_secs(2), rep1.recv())
        .await
        .expect("initial recv timed out")
        .unwrap();
    assert_eq!(got.part_bytes(0).unwrap().as_ref(), b"ping1");
    rep1.send(Message::single("pong1")).await.unwrap();
    compio::time::timeout(Duration::from_secs(2), req.recv())
        .await
        .expect("initial reply timed out")
        .unwrap();

    rep1.close().await.unwrap();
    let rep2 = rebind(&ep, || Socket::new(SocketType::Rep, Options::default())).await;

    req.send(Message::single("ping2")).await.unwrap();
    let got2 = compio::time::timeout(Duration::from_secs(3), rep2.recv())
        .await
        .expect("post-restart recv timed out")
        .unwrap();
    assert_eq!(got2.part_bytes(0).unwrap().as_ref(), b"ping2");
    rep2.send(Message::single("pong2")).await.unwrap();
    let reply2 = compio::time::timeout(Duration::from_secs(2), req.recv())
        .await
        .expect("post-restart reply timed out")
        .unwrap();
    assert_eq!(reply2.part_bytes(0).unwrap().as_ref(), b"pong2");
}

#[compio::test]
async fn req_state_machine_survives_drop_mid_cycle() {
    let rep1 = Socket::new(SocketType::Rep, Options::default());
    let ep = rep1.bind(tcp_ep(0)).await.unwrap();

    let req = Socket::new(SocketType::Req, fast_reconnect());
    req.connect(ep.clone()).await.unwrap();

    req.send(Message::single("a")).await.unwrap();
    compio::time::timeout(Duration::from_secs(2), rep1.recv())
        .await
        .expect("warm-up recv timed out")
        .unwrap();
    rep1.send(Message::single("a-reply")).await.unwrap();
    compio::time::timeout(Duration::from_secs(2), req.recv())
        .await
        .expect("warm-up reply timed out")
        .unwrap();

    req.send(Message::single("b")).await.unwrap();
    compio::time::sleep(Duration::from_millis(50)).await;
    rep1.close().await.unwrap();

    let rep2 = rebind(&ep, || Socket::new(SocketType::Rep, Options::default())).await;

    req.send(Message::single("c")).await.unwrap();
    let got = compio::time::timeout(Duration::from_secs(3), rep2.recv())
        .await
        .expect("post-drop recv timed out")
        .unwrap();
    assert_eq!(got.part_bytes(0).unwrap().as_ref(), b"c");
    rep2.send(Message::single("c-reply")).await.unwrap();
    let reply = compio::time::timeout(Duration::from_secs(2), req.recv())
        .await
        .expect("post-drop reply timed out")
        .unwrap();
    assert_eq!(reply.part_bytes(0).unwrap().as_ref(), b"c-reply");
}

// ── PUB / SUB ────────────────────────────────────────────────────────────────

#[compio::test]
async fn pub_sub_reconnect_replays_subscriptions() {
    let pub1 = Socket::new(SocketType::Pub, Options::default());
    let ep = pub1.bind(tcp_ep(0)).await.unwrap();

    let sub = Socket::new(SocketType::Sub, fast_reconnect());
    sub.connect(ep.clone()).await.unwrap();
    sub.subscribe("x.").await.unwrap();
    compio::time::sleep(Duration::from_millis(100)).await;

    pub1.send(Message::single("x.hello")).await.unwrap();
    let m1 = compio::time::timeout(Duration::from_millis(500), sub.recv())
        .await
        .expect("initial recv timed out")
        .unwrap();
    assert_eq!(m1.part_bytes(0).unwrap().as_ref(), b"x.hello");

    pub1.close().await.unwrap();
    let pub2 = rebind(&ep, || Socket::new(SocketType::Pub, Options::default())).await;

    compio::time::sleep(Duration::from_millis(150)).await;

    pub2.send(Message::single("x.world")).await.unwrap();
    pub2.send(Message::single("y.ignored")).await.unwrap();
    pub2.send(Message::single("x.again")).await.unwrap();

    let m2 = compio::time::timeout(Duration::from_secs(2), sub.recv())
        .await
        .expect("post-restart recv timed out")
        .unwrap();
    assert_eq!(m2.part_bytes(0).unwrap().as_ref(), b"x.world");
    let m3 = compio::time::timeout(Duration::from_millis(500), sub.recv())
        .await
        .expect("second post-restart recv timed out")
        .unwrap();
    assert_eq!(m3.part_bytes(0).unwrap().as_ref(), b"x.again");
}

// ── DEALER / ROUTER ──────────────────────────────────────────────────────────

#[compio::test]
async fn dealer_router_reconnect_after_router_restart() {
    let router1 = Socket::new(SocketType::Router, Options::default());
    let ep = router1.bind(tcp_ep(0)).await.unwrap();

    let dealer = Socket::new(
        SocketType::Dealer,
        Options {
            reconnect: ReconnectPolicy::Fixed(Duration::from_millis(30)),
            ..Options::default().identity(bytes::Bytes::from_static(b"d1"))
        },
    );
    dealer.connect(ep.clone()).await.unwrap();
    compio::time::sleep(Duration::from_millis(100)).await;

    dealer.send(Message::single("hello")).await.unwrap();
    let got1 = compio::time::timeout(Duration::from_secs(2), router1.recv())
        .await
        .expect("initial recv timed out")
        .unwrap();
    assert_eq!(got1.part_bytes(0).unwrap().as_ref(), b"d1");
    assert_eq!(got1.part_bytes(1).unwrap().as_ref(), b"hello");

    router1.close().await.unwrap();
    let router2 = rebind(&ep, || Socket::new(SocketType::Router, Options::default())).await;

    dealer.send(Message::single("after")).await.unwrap();
    let got2 = compio::time::timeout(Duration::from_secs(3), router2.recv())
        .await
        .expect("post-restart recv timed out")
        .unwrap();
    assert_eq!(got2.part_bytes(0).unwrap().as_ref(), b"d1");
    assert_eq!(got2.part_bytes(1).unwrap().as_ref(), b"after");
}
