pick_fast 0.1.1

use std::{
  net::IpAddr,
  sync::{
    Arc,
    atomic::{AtomicU32, Ordering},
  },
  thread,
};

use aok::{OK, Void};
use pick_fast::PickFast;

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

#[derive(Debug, Clone, Copy)]
struct DnsServer {
  ip: IpAddr,
}

#[test]
fn test() -> Void {
  const SERVERS: [DnsServer; 8] = [
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(8, 8, 8, 8)),
    }, // 100ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(1, 1, 1, 1)),
    }, // 80ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(223, 5, 5, 5)),
    }, // 5ms (最快)
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(208, 67, 222, 222)),
    }, // 60ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(9, 9, 9, 9)),
    }, // 40ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(1, 0, 0, 1)),
    }, // 20ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(114, 114, 114, 114)),
    }, // 70ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(180, 76, 76, 76)),
    }, // 90ms
  ];

  let lb = Arc::new(PickFast::<DnsServer, 8, pick_fast::Inverse>::new(SERVERS));

  println!("Load Balancer initialized with {} nodes.", SERVERS.len());

  let handles: Vec<_> = (0..8)
    .map(|_| {
      let c_lb = lb.clone();
      thread::spawn(move || {
        for _ in 0..1000 {
          let node = c_lb.pick();
          let latency = match node.ip {
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(8, 8, 8, 8)) => 100_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(1, 1, 1, 1)) => 80_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(223, 5, 5, 5)) => 5_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(208, 67, 222, 222)) => 60_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(9, 9, 9, 9)) => 40_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(1, 0, 0, 1)) => 20_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(114, 114, 114, 114)) => 70_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(180, 76, 76, 76)) => 90_000,
            _ => 100_000,
          };

          c_lb.set(node.index, latency);
        }
      })
    })
    .collect();

  for h in handles {
    h.join().unwrap();
  }

  let weight_li: &[AtomicU32; 8] = unsafe { std::mem::transmute(&lb.as_ref().weight_li) };

  let w_slow = weight_li[0].load(Ordering::Relaxed);
  let w_fast = weight_li[2].load(Ordering::Relaxed);

  println!("Slow Node Weight Google (8.8.8.8, 100ms): {}", w_slow);
  println!("Fast Node Weight AliDNS (223.5.5.5, 5ms): {}", w_fast);
  println!(
    "Ratio: {:.2} (Expected ~20.0)",
    w_fast as f64 / w_slow as f64
  );
  OK
}

#[test]
fn test_pick_count_with_chart() -> Void {
  const SERVERS: [DnsServer; 8] = [
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(8, 8, 8, 8)),
    }, // 100ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(1, 1, 1, 1)),
    }, // 80ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(223, 5, 5, 5)),
    }, // 5ms (最快)
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(208, 67, 222, 222)),
    }, // 60ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(9, 9, 9, 9)),
    }, // 40ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(1, 0, 0, 1)),
    }, // 20ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(114, 114, 114, 114)),
    }, // 70ms
    DnsServer {
      ip: IpAddr::V4(std::net::Ipv4Addr::new(180, 76, 76, 76)),
    }, // 90ms
  ];

  const LATENCIES: [u32; 8] = [100, 80, 5, 60, 40, 20, 70, 90];

  let lb = Arc::new(PickFast::<DnsServer, 8, pick_fast::Inverse>::new(SERVERS));

  let pick_counts: Arc<[AtomicU32; 8]> = Arc::new([const { AtomicU32::new(0) }; 8]);

  println!("Running 10000 picks to verify fast node is selected more than slow node...");

  let handles: Vec<_> = (0..8)
    .map(|_| {
      let c_lb = lb.clone();
      let c_counts = pick_counts.clone();
      thread::spawn(move || {
        for _ in 0..1250 {
          let node = c_lb.pick();

          c_counts[node.index].fetch_add(1, Ordering::Relaxed);

          let latency = match node.ip {
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(8, 8, 8, 8)) => 100_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(1, 1, 1, 1)) => 80_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(223, 5, 5, 5)) => 5_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(208, 67, 222, 222)) => 60_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(9, 9, 9, 9)) => 40_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(1, 0, 0, 1)) => 20_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(114, 114, 114, 114)) => 70_000,
            ip if ip == IpAddr::V4(std::net::Ipv4Addr::new(180, 76, 76, 76)) => 90_000,
            _ => 100_000,
          };

          c_lb.set(node.index, latency);
        }
      })
    })
    .collect();

  for h in handles {
    h.join().unwrap();
  }

  let mut counts = [0; 8];
  for (i, counter) in pick_counts.iter().enumerate() {
    counts[i] = counter.load(Ordering::Relaxed);
  }

  println!("\n=== 节点选择统计 ===");
  for (i, &count) in counts.iter().enumerate() {
    println!(
      "Node {}: {}ms -> {} times",
      SERVERS[i].ip, LATENCIES[i], count
    );
  }

  let slow_count = counts[0];
  let fast_count = counts[2];
  println!("\n慢节点 (8.8.8.8, 100ms) 被选中: {} 次", slow_count);
  println!("快节点 (223.5.5.5, 5ms) 被选中: {} 次", fast_count);
  println!("比例: {:.2}", fast_count as f64 / slow_count as f64);

  assert!(
    fast_count > slow_count,
    "Fast node should be picked more than slow node"
  );

  // 按选中次数排序数据 / Sort data by selection count
  let mut indexed_data: Vec<_> = SERVERS
    .iter()
    .zip(counts.iter())
    .zip(LATENCIES.iter())
    .enumerate()
    .collect();
  indexed_data.sort_by_key(|&(_, ((_, &count), _))| std::cmp::Reverse(count));

  let sorted_counts: Vec<u32> = indexed_data.iter().map(|&(_, ((_, &count), _))| count).collect();
  let sorted_latencies: Vec<u32> = indexed_data.iter().map(|&(_, ((..), &latency))| latency).collect();
  let sorted_servers: Vec<&DnsServer> = indexed_data.iter().map(|&(_, ((server, _), _))| server).collect();

  // 生成 SVG 图表 / Generate SVG charts
  draw_svg_histogram(&sorted_servers, &sorted_latencies, &sorted_counts)?;

  OK
}



fn draw_svg_histogram(servers: &[&DnsServer], latencies: &[u32], counts: &[u32]) -> Void {
  use std::fs;

  // 确保readme目录存在 / Ensure readme directory exists
  fs::create_dir_all("readme")?;

  // 生成中文版图表 / Generate Chinese version chart
  draw_3d_chart(servers, latencies, counts, "readme/rank-zh.svg", true)?;

  // 生成英文版图表 / Generate English version chart
  draw_3d_chart(servers, latencies, counts, "readme/rank-en.svg", false)?;

  println!("SVG图表已保存到 readme/rank-zh.svg 和 readme/rank-en.svg");
  println!("SVG charts saved to readme/rank-zh.svg and readme/rank-en.svg");

  OK
}

fn draw_3d_chart(
  servers: &[&DnsServer],
  latencies: &[u32],
  counts: &[u32],
  filename: &str,
  is_chinese: bool,
) -> Void {
  use std::fs;

  use svg::{
    Document,
    node::element::{Group, Polygon, Rectangle, Text},
  };

  let width = 1000;
  let height = 480; // 减少整体高度 / Reduce overall height
  let margin = 50; // 减少边距 / Reduce margin
  let title_margin = 60; // 减少标题区域高度 / Reduce title area height
  let chart_width = width - 2 * margin;
  let chart_height = height - 2 * margin - title_margin - 60; // 减少底部空间 / Reduce bottom space

  let max_count = *counts.iter().max().unwrap_or(&1);

  // 3D参数 / 3D parameters
  let depth = 40.0;
  let angle_rad: f64 = 0.5; // 约30度 / About 30 degrees
  let dx = depth * angle_rad.cos();
  let dy = depth * angle_rad.sin();

  let mut document = Document::new()
    .set("viewBox", (0, 0, width, height))
    .set("width", width)
    .set("height", height);

  // 定义统一的浅蓝色 / Define uniform light blue colors
  let front_color = "rgb(147, 197, 253)"; // 浅蓝色正面 / Light blue front
  let top_color = "rgb(191, 219, 254)"; // 更浅蓝色顶面 / Lighter blue top
  let right_color = "rgb(96, 165, 250)"; // 稍深蓝色侧面 / Slightly darker blue side

  // 透明背景，不添加背景矩形 / Transparent background, no background rectangle

  // 标题 / Title
  let title = if is_chinese {
    "PickFast 使用演示：DNS 响应延时 与 选中次数"
  } else {
    "PickFast Demo: DNS Response Latency vs Selection Count"
  };

  let title_text = Text::new(title)
    .set("x", width / 2)
    .set("y", 35) // 标题位置上移 / Move title up
    .set("text-anchor", "middle")
    .set("font-family", "Arial, sans-serif")
    .set("font-size", 22) // 稍微减小字体 / Slightly reduce font size
    .set("font-weight", "bold")
    .set("fill", "rgb(30, 41, 59)");
  document = document.add(title_text);

  // 绘制3D柱状图 / Draw 3D bars
  let bar_width = chart_width as f64 / 8.0 * 0.7;
  let bar_spacing = chart_width as f64 / 8.0;

  for (i, &count) in counts.iter().enumerate() {
    if count == 0 {
      continue;
    }

    let x = margin as f64 + i as f64 * bar_spacing + bar_spacing * 0.15;
    let bar_height = (count as f64 / max_count as f64) * chart_height as f64;
    let y = margin as f64 + title_margin as f64 + chart_height as f64 - bar_height; // 图表起始位置下移 / Move chart start position down

    let mut group = Group::new();

    // 正面 / Front face
    let front_face = Rectangle::new()
      .set("x", x)
      .set("y", y)
      .set("width", bar_width)
      .set("height", bar_height)
      .set("fill", front_color)
      .set("stroke", "rgba(0,0,0,0.2)")
      .set("stroke-width", 1);
    group = group.add(front_face);

    // 顶面 / Top face
    let top_points = format!(
      "{},{} {},{} {},{} {},{}",
      x,
      y,
      x + bar_width,
      y,
      x + bar_width + dx,
      y - dy,
      x + dx,
      y - dy
    );
    let top_face = Polygon::new()
      .set("points", top_points)
      .set("fill", top_color)
      .set("stroke", "rgba(0,0,0,0.2)")
      .set("stroke-width", 1);
    group = group.add(top_face);

    // 右侧面 / Right face
    let right_points = format!(
      "{},{} {},{} {},{} {},{}",
      x + bar_width,
      y,
      x + bar_width,
      y + bar_height,
      x + bar_width + dx,
      y + bar_height - dy,
      x + bar_width + dx,
      y - dy
    );
    let right_face = Polygon::new()
      .set("points", right_points)
      .set("fill", right_color)
      .set("stroke", "rgba(0,0,0,0.2)")
      .set("stroke-width", 1);
    group = group.add(right_face);

    // 数值标签描边 / Value label stroke (white outline)
    let value_stroke = Text::new(format!("{count}"))
      .set("x", x + bar_width / 2.0)
      .set("y", y - 10.0)
      .set("text-anchor", "middle")
      .set("font-family", "Arial, sans-serif")
      .set("font-size", 13)
      .set("font-weight", "bold")
      .set("fill", "none")
      .set("stroke", "white")
      .set("stroke-width", 3);
    group = group.add(value_stroke);

    // 数值标签 / Value label (black text)
    let value_text = Text::new(format!("{count}"))
      .set("x", x + bar_width / 2.0)
      .set("y", y - 10.0)
      .set("text-anchor", "middle")
      .set("font-family", "Arial, sans-serif")
      .set("font-size", 13)
      .set("font-weight", "bold")
      .set("fill", "black");
    group = group.add(value_text);

    // 延时标签 / Latency label
    let latency_text = Text::new(format!("{}ms", latencies[i]))
      .set("x", x + bar_width / 2.0)
      .set("y", margin + title_margin + chart_height + 20) // 调整底部标签位置 / Adjust bottom label position
      .set("text-anchor", "middle")
      .set("font-family", "Arial, sans-serif")
      .set("font-size", 12)
      .set("fill", "rgb(71, 85, 105)");
    group = group.add(latency_text);

    // IP地址标签 / IP address label
    let ip_text = Text::new(format!("{}", servers[i].ip))
      .set("x", x + bar_width / 2.0)
      .set("y", margin + title_margin + chart_height + 38) // 调整底部标签位置 / Adjust bottom label position
      .set("text-anchor", "middle")
      .set("font-family", "Arial, sans-serif")
      .set("font-size", 10)
      .set("fill", "rgb(100, 116, 139)");
    group = group.add(ip_text);

    document = document.add(group);
  }

  // Y轴标签 / Y-axis labels
  let y_desc = if is_chinese {
    "选择次数"
  } else {
    "Selection Count"
  };
  let y_label_x = 35; // Y轴标签X位置，更靠近图表 / Y-axis label X position, closer to chart
  let y_label_y = margin + title_margin + chart_height / 2; // Y轴标签Y位置，居中于图表 / Y-axis label Y position, centered on chart
  let y_label = Text::new(y_desc)
    .set("x", y_label_x)
    .set("y", y_label_y)
    .set("text-anchor", "middle")
    .set("font-family", "Arial, sans-serif")
    .set("font-size", 16) // 增大字体 / Increase font size
    .set("fill", "rgb(71, 85, 105)")
    .set("transform", format!("rotate(-90, {y_label_x}, {y_label_y})")); // 使用变量名 / Use variable names
  document = document.add(y_label);

  // 保存文件 / Save file
  fs::write(filename, document.to_string())?;

  OK
}