pick_fast 0.1.18

use std::{sync::Arc, time::Duration};
mod dns_server;
use aok::{OK, Void};
use dns_server::{DNS_SERVER_LI, DnsServer};
use futures::StreamExt;
use hickory_resolver::{
  Resolver,
  config::{NameServerConfig, ResolverConfig},
  proto::xfer::Protocol,
};
use pick_fast::PickFast;
use race::Race;

#[static_init::constructor(0)]
extern "C" fn _log_init() {
  log_init::init();
}

// Create resolver with specific DNS server / 使用指定 DNS 服务器创建解析器
fn create_resolver(
  server: &DnsServer,
) -> Resolver<hickory_resolver::name_server::TokioConnectionProvider> {
  let ns = NameServerConfig::new(std::net::SocketAddr::new(server.ip, 53), Protocol::Udp);
  let mut config = ResolverConfig::new();
  config.add_name_server(ns);

  let provider = hickory_resolver::name_server::TokioConnectionProvider::default();
  Resolver::builder_with_config(config, provider).build()
}

#[tokio::test]
async fn test_iter_with_race_dns() -> Void {
  let lb = Arc::new(PickFast::<DnsServer, pick_fast::Inverse>::new(
    DNS_SERVER_LI,
  ));
  const HOST: &str = "qq.com";

  println!("Testing iter() with race crate for real DNS resolution...");

  let lb_clone = lb.clone();
  // DNS resolution function / DNS 解析函数
  let resolve_dns = move |(index, node): &(usize, &pick_fast::Node<DnsServer>)| {
    let start = std::time::Instant::now();
    let lb = lb_clone.clone();
    let index = *index;
    let resolver = create_resolver(&node.data);
    let server_ip = node.data.ip;

    async move {
      match resolver.lookup_ip(HOST).await {
        Ok(response) => {
          if let Some(addr) = response.iter().next() {
            let duration = start.elapsed();
            let latency_us = duration.as_micros() as u32;

            // Successful: update latency weight / 成功：更新延时权重
            lb.set(index, latency_us);

            println!("  ✅ {HOST} via {server_ip} -> {addr} in {duration:?} ({latency_us}μs)");
            Ok(addr)
          } else {
            let duration = start.elapsed();
            lb.failed(index);

            let error = std::io::Error::new(
              std::io::ErrorKind::NotFound,
              format!("No address found for {HOST}"),
            );
            println!("  ❌ {HOST} via {server_ip} failed after {duration:?}: {error}");
            Err(error)
          }
        }
        Err(e) => {
          let duration = start.elapsed();
          lb.failed(index);

          let error = std::io::Error::other(e.to_string());
          println!("  ❌ {HOST} via {server_ip} failed after {duration:?}: {e}");
          Err(error)
        }
      }
    }
  };

  // Use race for staggered DNS resolution / 使用 race 进行阶梯式 DNS 解析
  let mut race = Race::new(Duration::from_millis(500), resolve_dns, lb.iter());

  println!("Starting staggered DNS resolution with 500ms intervals...");

  let mut resolved_ip = None;
  while let Some((_args, result)) = race.next().await {
    match result {
      Ok(addr) => {
        println!("resolve_dns {addr}");
        resolved_ip = Some(addr);
        break;
      }
      Err(e) => {
        println!("⚠️  Resolution attempt failed: {e}");
      }
    }
  }

  // Mark remaining tasks as failed / 将剩余任务标记为失败
  for ((index, node), _) in race.ing {
    lb.failed(index);
    log::debug!("cancel {:?} {HOST}", node.data.ip);
  }

  if resolved_ip.is_some() {
    println!("✅ Race-based DNS resolution completed successfully");
    println!("This demonstrates how PickFast + Race provides fast DNS failover");
  } else {
    println!("❌ All DNS resolutions failed (network issue?)");
  }

  println!("Real DNS resolution test completed");

  OK
}