use crate::protocol::{
    BalancingIpDetails,
    routeplanner::{FailingAddress, IpBlock, RotatingIpDetails, RoutePlannerStatus},
};
use async_trait::async_trait;
use ipnet::IpNet;
use rand::Rng;
use std::{
    collections::HashMap,
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    str::FromStr,
    sync::Mutex,
};
use tracing::{debug, info};
#[async_trait]
pub trait RoutePlanner: Send + Sync {
    fn get_status(&self) -> RoutePlannerStatus;
    fn free_address(&self, address: &str);
    fn free_all_addresses(&self);
    fn mark_failed(&self, address: &str);
    fn get_address(&self) -> Option<std::net::IpAddr>;
}
pub struct BalancingIpRoutePlanner {
    ip_blocks: Vec<IpBlock>,
    parsed_blocks: Vec<IpNet>,
    failing_addresses: Mutex<HashMap<String, u64>>,
    block_index: Mutex<usize>,
    ip_indices: Mutex<Vec<u128>>,
}
impl BalancingIpRoutePlanner {
    pub fn new(cidrs: Vec<String>) -> Result<Self, String> {
        let mut ip_blocks = Vec::with_capacity(cidrs.len());
        let mut parsed_blocks = Vec::with_capacity(cidrs.len());
        let mut total_ips = 0;
        for cidr in cidrs.iter() {
            let parsed = IpNet::from_str(cidr)
                .or_else(|_| {
                    let suffix = if cidr.contains(':') { "/128" } else { "/32" };
                    IpNet::from_str(&format!("{}{}", cidr, suffix))
                })
                .map_err(|e| format!("Invalid CIDR or IP '{}': {}", cidr, e))?;
            let block_type = match parsed {
                IpNet::V4(_) => "Inet4Address",
                IpNet::V6(_) => "Inet6Address",
            }
            .to_string();
            let size = match parsed {
                IpNet::V4(net) => {
                    let prefix_len = net.prefix_len();
                    if prefix_len >= 32 {
                        1
                    } else {
                        2u128.pow(32 - prefix_len as u32)
                    }
                }
                IpNet::V6(net) => {
                    let prefix_len = net.prefix_len();
                    if prefix_len >= 128 {
                        1
                    } else if prefix_len <= 64 {
                        u128::MAX
                    } else {
                        2u128.pow(128 - prefix_len as u32)
                    }
                }
            };
            if size == u128::MAX {
                info!(
                    "Added {} block: {} (virtually unlimited addresses)",
                    block_type, cidr
                );
            } else {
                info!("Added {} block: {} ({} addresses)", block_type, cidr, size);
                total_ips += size;
            }
            ip_blocks.push(IpBlock {
                block_type,
                size: cidr.clone(),
            });
            parsed_blocks.push(parsed);
        }
        if total_ips > 0 && total_ips != u128::MAX {
            info!(
                "Route planner initialized with {} total addresses",
                total_ips
            );
        } else if total_ips == u128::MAX {
            info!("Route planner initialized with virtually unlimited addresses");
        }
        Ok(Self {
            ip_blocks,
            parsed_blocks,
            failing_addresses: Mutex::new(HashMap::new()),
            block_index: Mutex::new(0),
            ip_indices: Mutex::new(vec![0; cidrs.len()]),
        })
    }
    fn calculate_ip(block: &IpNet, index: u128) -> IpAddr {
        let prefix_len = block.prefix_len();
        match block {
            IpNet::V4(net) => {
                let addr_u32 = u32::from(net.addr());
                let offset = if prefix_len >= 32 {
                    0
                } else if prefix_len == 0 {
                    index as u32
                } else {
                    (index as u32) & (!0u32 >> prefix_len)
                };
                IpAddr::V4(Ipv4Addr::from(addr_u32 + offset))
            }
            IpNet::V6(net) => {
                let addr_u128 = u128::from(net.addr());
                let offset = if prefix_len >= 128 {
                    0
                } else if prefix_len == 0 {
                    index
                } else {
                    index & (!0u128 >> prefix_len)
                };
                IpAddr::V6(Ipv6Addr::from(addr_u128 + offset))
            }
        }
    }
    fn next_ip(&self) -> IpAddr {
        let mut b_idx = self.block_index.lock().unwrap_or_else(|e| e.into_inner());
        let mut indices = self.ip_indices.lock().unwrap_or_else(|e| e.into_inner());
        let block_idx = *b_idx % self.parsed_blocks.len();
        let block = &self.parsed_blocks[block_idx];
        let prefix_len = block.prefix_len();
        let max_bits = match block {
            IpNet::V4(_) => 32,
            IpNet::V6(_) => 128,
        };
        let size_bits = max_bits - prefix_len;
        let increment = if size_bits > 7 {
            rand::thread_rng().gen_range(10..20) as u128
        } else {
            1
        };
        let current_index = indices[block_idx];
        indices[block_idx] = current_index.wrapping_add(increment);
        let final_index = indices[block_idx];
        *b_idx = (*b_idx + 1) % self.parsed_blocks.len();
        let ip = Self::calculate_ip(block, final_index);
        debug!("Route planner picked IP: {}", ip);
        ip
    }
    fn get_address_internal(&self) -> Option<IpAddr> {
        for _ in 0..100 {
            let ip = self.next_ip();
            let ip_str = ip.to_string();
            let mut failing = self
                .failing_addresses
                .lock()
                .unwrap_or_else(|e| e.into_inner());
            if let Some(&timestamp) = failing.get(&ip_str) {
                if crate::common::utils::now_ms()
                    > timestamp + crate::audio::constants::ROUTE_PLANNER_FAIL_EXPIRE_MS
                {
                    failing.remove(&ip_str);
                } else {
                    continue;
                }
            }
            return Some(ip);
        }
        None
    }
}
#[async_trait]
impl RoutePlanner for BalancingIpRoutePlanner {
    fn get_status(&self) -> RoutePlannerStatus {
        let failing = self
            .failing_addresses
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        let failing_vec: Vec<FailingAddress> = failing
            .iter()
            .map(|(addr, &ts)| FailingAddress {
                failing_address: addr.clone(),
                failing_timestamp: ts,
                failing_time: String::new(),
            })
            .collect();
        if self.ip_blocks.len() == 1 {
            let index = self.ip_indices.lock().unwrap_or_else(|e| e.into_inner())[0];
            let current_ip = Self::calculate_ip(&self.parsed_blocks[0], index).to_string();
            RoutePlannerStatus::RotatingIpRoutePlanner(RotatingIpDetails {
                ip_block: self.ip_blocks[0].clone(),
                failing_addresses: failing_vec,
                rotate_index: "0".to_string(),
                ip_index: index.to_string(),
                current_address: current_ip,
            })
        } else {
            RoutePlannerStatus::BalancingIpRoutePlanner(BalancingIpDetails {
                ip_block: self.ip_blocks[0].clone(),
                failing_addresses: failing_vec,
            })
        }
    }
    fn free_address(&self, address: &str) {
        self.failing_addresses
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .remove(address);
    }
    fn free_all_addresses(&self) {
        self.failing_addresses
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .clear();
    }
    fn mark_failed(&self, address: &str) {
        self.failing_addresses
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .insert(address.to_string(), crate::common::utils::now_ms());
    }
    fn get_address(&self) -> Option<std::net::IpAddr> {
        self.get_address_internal()
    }
}
