alpine-protocol-sdk 0.2.4

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::time::Duration;

use alpine::attestation::AttesterRegistry;
use thiserror::Error;
use tokio::time::{sleep, Duration as TokioDuration};
use tracing::{info, warn};

use crate::discovery::{
    DiscoveryAttempt, DiscoveryClient, DiscoveryClientOptions, DiscoveryError, DiscoveryOutcome,
    DiscoveryResult,
};
use crate::environment::ensure_supported_environment;
use crate::phase::claim_discovery;
use crate::self_check::run_sdk_self_check;

#[derive(Debug, Clone)]
pub struct DiscoveryRunOptions {
    pub local_addr: Option<SocketAddr>,
    pub prefer_multicast: bool,
    pub allow_broadcast: bool,
    pub cache_first: bool,
    pub cached_targets: Vec<SocketAddr>,
    pub scan_subnets: bool,
    pub scan_rate_per_sec: u32,
    pub scan_timeout_ms: u64,
    pub scan_max_hosts: u32,
    pub attester_registry: Option<AttesterRegistry>,
}

impl Default for DiscoveryRunOptions {
    fn default() -> Self {
        Self {
            local_addr: None,
            prefer_multicast: false,
            allow_broadcast: true,
            cache_first: false,
            cached_targets: Vec::new(),
            scan_subnets: false,
            scan_rate_per_sec: 200,
            scan_timeout_ms: 500,
            scan_max_hosts: 1024,
            attester_registry: None,
        }
    }
}

impl DiscoveryRunOptions {
    /// Sensible defaults for a typical LAN deployment (broadcast on, no subnet scan).
    pub fn defaults_for_lan() -> Self {
        Self::default()
    }

    /// Defaults for lab environments (no multicast, cache-first probing).
    pub fn defaults_for_lab() -> Self {
        Self {
            prefer_multicast: false,
            allow_broadcast: false,
            cache_first: true,
            scan_subnets: false,
            scan_max_hosts: 0,
            ..Self::default()
        }
    }
}

#[derive(Debug, Error)]
pub enum DiscoveryRunError {
    #[error("unicast discovery failed: {0}")]
    Unicast(DiscoveryError),
    #[error("broadcast discovery failed: {0}")]
    Broadcast(DiscoveryError),
    #[error("no viable interfaces for broadcast discovery (need a non-loopback IPv4 address)")]
    NoInterfaces,
    #[error("cached unicast discovery failed")]
    CachedUnicastFailed,
    #[error("subnet scan discovery failed")]
    SubnetScanFailed,
    #[error("subnet scan disabled (scan_max_hosts = 0)")]
    SubnetScanDisabled,
    #[error("unsupported environment: {0}")]
    UnsupportedEnvironment(String),
}

#[derive(Debug)]
pub struct DiscoveryRunReport {
    pub result: Result<DiscoveryOutcome, DiscoveryRunError>,
    pub attempts: Vec<DiscoveryAttempt>,
}

impl DiscoveryRunReport {
    pub fn chosen_attempt(&self) -> Option<DiscoveryAttempt> {
        let outcome = self.result.as_ref().ok()?;
        self.attempts
            .iter()
            .find(|attempt| attempt.target == outcome.peer)
            .cloned()
    }

    pub fn decision_summary(&self) -> String {
        match &self.result {
            Ok(outcome) => {
                if let Some(attempt) = self
                    .attempts
                    .iter()
                    .find(|attempt| attempt.target == outcome.peer)
                {
                    format!(
                        "selected {} via {} (local_bind {}); {} attempts",
                        outcome.peer,
                        attempt.mode.as_str(),
                        attempt.local_bind,
                        self.attempts.len()
                    )
                } else {
                    format!(
                        "selected {} (attempt not recorded); {} attempts",
                        outcome.peer,
                        self.attempts.len()
                    )
                }
            }
            Err(err) => format!(
                "no device selected: {} ({} attempts)",
                err,
                self.attempts.len()
            ),
        }
    }

    pub fn summary(&self) -> String {
        match &self.result {
            Ok(outcome) => format!(
                "discovered {} (device_id={}) after {} attempts",
                outcome.peer,
                outcome.reply.device_id,
                self.attempts.len()
            ),
            Err(err) => {
                let mut hints = Vec::new();
                for attempt in &self.attempts {
                    if let Some(error) = &attempt.error {
                        if let Some(hint) = error.hint {
                            hints.push(hint);
                        }
                    }
                }
                hints.sort();
                hints.dedup();
                if hints.is_empty() {
                    format!(
                        "discovery failed: {} ({} attempts)",
                        err,
                        self.attempts.len()
                    )
                } else {
                    format!(
                        "discovery failed: {} ({} attempts); hints: {}",
                        err,
                        self.attempts.len(),
                        hints.join("; ")
                    )
                }
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct DiscoveryInterface {
    pub name: String,
    pub local_ip: std::net::Ipv4Addr,
    pub netmask: std::net::Ipv4Addr,
    pub broadcast: std::net::Ipv4Addr,
}

#[derive(Debug, Clone)]
pub struct DiscoveryDryRun {
    pub remote_addr: SocketAddr,
    pub unicast_target: Option<SocketAddr>,
    pub prefer_multicast: bool,
    pub allow_broadcast: bool,
    pub cached_targets: Vec<SocketAddr>,
    pub scan_subnets: bool,
    pub scan_rate_per_sec: u32,
    pub scan_timeout_ms: u64,
    pub scan_max_hosts: u32,
    pub interfaces: Vec<DiscoveryInterface>,
    pub broadcast_targets: Vec<SocketAddr>,
}

impl DiscoveryDryRun {
    pub fn print_table(&self) -> String {
        let mut lines = Vec::new();
        lines.push("interface, local_ip, netmask, broadcast".to_string());
        for iface in &self.interfaces {
            lines.push(format!(
                "{}, {}, {}, {}",
                iface.name, iface.local_ip, iface.netmask, iface.broadcast
            ));
        }
        if let Some(target) = self.unicast_target {
            lines.push(format!("unicast_target, {}", target));
        }
        for target in &self.broadcast_targets {
            lines.push(format!("broadcast_target, {}", target));
        }
        if !self.cached_targets.is_empty() {
            for target in &self.cached_targets {
                lines.push(format!("cached_target, {}", target));
            }
        }
        lines.join("\n")
    }
}

#[derive(Debug, Clone)]
pub struct DiscoveryRetryPolicy {
    pub max_attempts: usize,
    pub backoff_base_ms: u64,
    pub backoff_max_ms: u64,
}

impl Default for DiscoveryRetryPolicy {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            backoff_base_ms: 200,
            backoff_max_ms: 2_000,
        }
    }
}

fn default_local_addr() -> SocketAddr {
    SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
}

fn resolve_local_addr(remote_addr: SocketAddr, override_addr: Option<SocketAddr>) -> SocketAddr {
    if let Some(addr) = override_addr {
        return addr;
    }
    if matches!(remote_addr.ip(), IpAddr::V4(v4) if v4.is_broadcast()) {
        return default_local_addr();
    }
    let bind_addr = if remote_addr.is_ipv4() {
        SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
    } else {
        SocketAddr::new(IpAddr::V6(std::net::Ipv6Addr::UNSPECIFIED), 0)
    };
    if let Ok(sock) = std::net::UdpSocket::bind(bind_addr) {
        if sock.connect(remote_addr).is_ok() {
            if let Ok(local) = sock.local_addr() {
                return SocketAddr::new(local.ip(), 0);
            }
        }
    }
    default_local_addr()
}

pub async fn run_discovery(remote_addr: SocketAddr) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    run_discovery_with_options(remote_addr, DiscoveryRunOptions::default()).await
}

/// Runs discovery and returns both the result and per-attempt diagnostics.
pub async fn run_discovery_with_report(remote_addr: SocketAddr) -> DiscoveryRunReport {
    run_discovery_with_options_report(remote_addr, DiscoveryRunOptions::default()).await
}

pub async fn run_discovery_with_retry(
    remote_addr: SocketAddr,
    policy: DiscoveryRetryPolicy,
) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    run_discovery_with_options_retry(remote_addr, DiscoveryRunOptions::default(), policy).await
}

pub async fn run_discovery_with_options(
    remote_addr: SocketAddr,
    opts: DiscoveryRunOptions,
) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    run_discovery_with_options_inner(remote_addr, opts, None).await
}

pub async fn discover_with_cache(
    remote_addr: SocketAddr,
    cached_targets: Vec<SocketAddr>,
) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    let mut opts = DiscoveryRunOptions::defaults_for_lan();
    opts.cached_targets = cached_targets;
    run_discovery_with_options(remote_addr, opts).await
}

pub async fn discover_with_cache_report(
    remote_addr: SocketAddr,
    cached_targets: Vec<SocketAddr>,
) -> DiscoveryRunReport {
    let mut opts = DiscoveryRunOptions::defaults_for_lan();
    opts.cached_targets = cached_targets;
    run_discovery_with_options_report(remote_addr, opts).await
}

pub async fn run_discovery_with_options_retry(
    remote_addr: SocketAddr,
    opts: DiscoveryRunOptions,
    policy: DiscoveryRetryPolicy,
) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    let attempts = policy.max_attempts.max(1);
    let mut attempt = 0usize;
    loop {
        attempt = attempt.saturating_add(1);
        match run_discovery_with_options(remote_addr, opts.clone()).await {
            Ok(outcome) => {
                if attempt > 1 {
                    warn!(
                        "[ALPINE][DISCOVERY][WARN] discovery succeeded after {} retries",
                        attempt - 1
                    );
                }
                return Ok(outcome);
            }
            Err(err) => {
                if attempt >= attempts {
                    return Err(err);
                }
                sleep(TokioDuration::from_millis(discovery_backoff(
                    &policy, attempt,
                )))
                .await;
            }
        }
    }
}

async fn run_discovery_with_options_inner(
    remote_addr: SocketAddr,
    opts: DiscoveryRunOptions,
    mut diagnostics: Option<&mut Vec<DiscoveryAttempt>>,
) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    run_sdk_self_check();
    if let Err(err) = ensure_supported_environment() {
        return Err(DiscoveryRunError::UnsupportedEnvironment(err.to_string()));
    }
    let _phase_guard = claim_discovery()
        .map_err(|_| DiscoveryRunError::Unicast(DiscoveryError::PermissionDenied))?;
    let is_broadcast_target = matches!(remote_addr.ip(), IpAddr::V4(v4) if v4.is_broadcast());

    if opts.cache_first {
        if let Ok(outcome) = attempt_cached_unicast(&opts, &mut diagnostics) {
            return Ok(outcome);
        }
    }

    if !is_broadcast_target {
        let local_addr = resolve_local_addr(remote_addr, opts.local_addr);
        let mut options =
            DiscoveryClientOptions::new(remote_addr, local_addr, Duration::from_secs(3));
        options.attester_registry = opts.attester_registry.clone();
        options = options.disable_multicast().disable_broadcast();
        info!(
            "[ALPINE][DISCOVERY] route_hint remote={} local_bind={}",
            remote_addr, options.local_addr
        );
        let local_bind = options.local_addr;
        let client = DiscoveryClient::new(options).map_err(DiscoveryRunError::Unicast)?;
        let report = client.discover_with_report(&["alpine-control".to_string()]);
        append_attempts(&mut diagnostics, &report);
        match report.into_result() {
            Ok(outcome) => return Ok(outcome),
            Err(err) => {
                if opts.allow_broadcast && remote_addr.is_ipv4() {
                    warn!(
                        "[ALPINE][DISCOVERY][WARN] unicast discovery failed ({}); falling back to broadcast",
                        err
                    );
                    match attempt_broadcast(
                        remote_addr.port(),
                        &opts,
                        Some(local_bind.ip()),
                        &mut diagnostics,
                    ) {
                        Ok(outcome) => return Ok(outcome),
                        Err(broadcast_err) => {
                            warn!(
                                "[ALPINE][DISCOVERY][WARN] falling back to cached unicast targets"
                            );
                            if let Ok(outcome) = attempt_cached_unicast(&opts, &mut diagnostics) {
                                return Ok(outcome);
                            }
                            if opts.scan_subnets {
                                warn!("[ALPINE][DISCOVERY][WARN] falling back to subnet scan");
                                return attempt_subnet_scan(
                                    remote_addr.port(),
                                    &opts,
                                    Some(local_bind.ip()),
                                    &mut diagnostics,
                                )
                                .await;
                            }
                            return Err(broadcast_err);
                        }
                    }
                }
                return Err(DiscoveryRunError::Unicast(err));
            }
        }
    }

    match attempt_broadcast(remote_addr.port(), &opts, None, &mut diagnostics) {
        Ok(outcome) => Ok(outcome),
        Err(broadcast_err) => {
            warn!("[ALPINE][DISCOVERY][WARN] falling back to cached unicast targets");
            if let Ok(outcome) = attempt_cached_unicast(&opts, &mut diagnostics) {
                return Ok(outcome);
            }
            if opts.scan_subnets {
                warn!("[ALPINE][DISCOVERY][WARN] falling back to subnet scan");
                return attempt_subnet_scan(remote_addr.port(), &opts, None, &mut diagnostics)
                    .await;
            }
            Err(broadcast_err)
        }
    }
}

pub async fn run_discovery_with_options_report(
    remote_addr: SocketAddr,
    opts: DiscoveryRunOptions,
) -> DiscoveryRunReport {
    let mut attempts = Vec::new();
    let result = run_discovery_with_options_inner(remote_addr, opts, Some(&mut attempts)).await;
    DiscoveryRunReport { result, attempts }
}

/// Returns the computed discovery plan without sending any packets.
pub fn discovery_dry_run(
    remote_addr: SocketAddr,
    opts: DiscoveryRunOptions,
) -> Result<DiscoveryDryRun, DiscoveryRunError> {
    let interfaces = collect_interfaces()?;
    let broadcast_targets = if opts.allow_broadcast && remote_addr.is_ipv4() {
        interfaces
            .iter()
            .map(|iface| SocketAddr::new(IpAddr::V4(iface.broadcast), remote_addr.port()))
            .collect::<Vec<_>>()
    } else {
        Vec::new()
    };
    let unicast_target = if matches!(remote_addr.ip(), IpAddr::V4(v4) if v4.is_broadcast()) {
        None
    } else {
        Some(remote_addr)
    };
    Ok(DiscoveryDryRun {
        remote_addr,
        unicast_target,
        prefer_multicast: opts.prefer_multicast,
        allow_broadcast: opts.allow_broadcast,
        cached_targets: opts.cached_targets,
        scan_subnets: opts.scan_subnets,
        scan_rate_per_sec: opts.scan_rate_per_sec,
        scan_timeout_ms: opts.scan_timeout_ms,
        scan_max_hosts: opts.scan_max_hosts,
        interfaces: interfaces
            .into_iter()
            .map(|iface| DiscoveryInterface {
                name: iface.iface,
                local_ip: iface.local_ip,
                netmask: iface.netmask,
                broadcast: iface.broadcast,
            })
            .collect(),
        broadcast_targets,
    })
}

fn attempt_broadcast(
    port: u16,
    opts: &DiscoveryRunOptions,
    preferred_ip: Option<IpAddr>,
    diagnostics: &mut Option<&mut Vec<DiscoveryAttempt>>,
) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    let mut attempts = collect_interfaces()?;
    if attempts.is_empty() {
        return Err(DiscoveryRunError::NoInterfaces);
    }

    if let Some(pref) = preferred_ip {
        if let Some(idx) = attempts.iter().position(|a| IpAddr::V4(a.local_ip) == pref) {
            let preferred = attempts.remove(idx);
            attempts.insert(0, preferred);
        }
    }

    let mut last_err: Option<DiscoveryError> = None;
    for attempt in attempts.iter() {
        let mut options = DiscoveryClientOptions::new(
            SocketAddr::new(IpAddr::V4(attempt.broadcast), port),
            SocketAddr::new(IpAddr::V4(attempt.local_ip), 0),
            Duration::from_secs(3),
        );
        options.interface = Some(attempt.iface.clone());
        options.attester_registry = opts.attester_registry.clone();
        if !opts.prefer_multicast {
            options = options.disable_multicast();
        }
        if !opts.allow_broadcast {
            options = options.disable_broadcast();
        }
        info!(
            "[ALPINE][DISCOVERY] iface={} local_ip={} netmask={} broadcast={} bound={}:0 so_broadcast={}",
            attempt.iface,
            attempt.local_ip,
            attempt.netmask,
            attempt.broadcast,
            attempt.local_ip,
            options.allow_broadcast
        );

        match DiscoveryClient::new(options) {
            Ok(client) => {
                let report = client.discover_with_report(&["alpine-control".to_string()]);
                append_attempts(diagnostics, &report);
                match report.into_result() {
                    Ok(outcome) => return Ok(outcome),
                    Err(err) => {
                        warn!(
                            "[ALPINE][DISCOVERY][WARN] iface={} error={}",
                            attempt.iface, err
                        );
                        last_err = Some(err);
                        continue;
                    }
                }
            }
            Err(err) => {
                warn!(
                    "[ALPINE][DISCOVERY][WARN] iface={} error={}",
                    attempt.iface, err
                );
                last_err = Some(err);
                continue;
            }
        }
    }

    Err(DiscoveryRunError::Broadcast(
        last_err.unwrap_or(DiscoveryError::Timeout),
    ))
}

fn attempt_cached_unicast(
    opts: &DiscoveryRunOptions,
    diagnostics: &mut Option<&mut Vec<DiscoveryAttempt>>,
) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    if opts.cached_targets.is_empty() {
        return Err(DiscoveryRunError::CachedUnicastFailed);
    }
    for target in opts.cached_targets.iter() {
        let local_addr = resolve_local_addr(*target, None);
        let mut options = DiscoveryClientOptions::new(
            *target,
            local_addr,
            Duration::from_millis(opts.scan_timeout_ms),
        );
        options.attester_registry = opts.attester_registry.clone();
        options = options.disable_multicast().disable_broadcast();
        info!(
            "[ALPINE][DISCOVERY] cached_unicast target={} local_bind={}",
            target, options.local_addr
        );
        if let Ok(client) = DiscoveryClient::new(options) {
            let report = client.discover_with_report(&["alpine-control".to_string()]);
            append_attempts(diagnostics, &report);
            match report.into_result() {
                Ok(outcome) => return Ok(outcome),
                Err(err) => {
                    warn!(
                        "[ALPINE][DISCOVERY][WARN] cached_unicast target={} error={}",
                        target, err
                    );
                }
            }
        }
    }
    Err(DiscoveryRunError::CachedUnicastFailed)
}

async fn attempt_subnet_scan(
    port: u16,
    opts: &DiscoveryRunOptions,
    preferred_ip: Option<IpAddr>,
    diagnostics: &mut Option<&mut Vec<DiscoveryAttempt>>,
) -> Result<DiscoveryOutcome, DiscoveryRunError> {
    if opts.scan_max_hosts == 0 {
        return Err(DiscoveryRunError::SubnetScanDisabled);
    }
    let mut attempts = collect_interfaces()?;
    if attempts.is_empty() {
        return Err(DiscoveryRunError::NoInterfaces);
    }
    if let Some(pref) = preferred_ip {
        if let Some(idx) = attempts.iter().position(|a| IpAddr::V4(a.local_ip) == pref) {
            let preferred = attempts.remove(idx);
            attempts.insert(0, preferred);
        }
    }

    let rate = opts.scan_rate_per_sec.max(1);
    let delay_ms = (1000u64 / rate as u64).max(1);
    let timeout_ms = opts.scan_timeout_ms.max(100);

    for attempt in attempts.iter() {
        let mut scanned = 0u32;
        let (network, broadcast) = network_bounds(attempt.local_ip, attempt.netmask);
        let start = network.saturating_add(1);
        let end = broadcast.saturating_sub(1);
        info!(
            "[ALPINE][DISCOVERY] subnet_scan iface={} local_ip={} netmask={} range={}.{} timeout_ms={} rate_per_sec={} max_hosts={}",
            attempt.iface,
            attempt.local_ip,
            attempt.netmask,
            std::net::Ipv4Addr::from(network),
            std::net::Ipv4Addr::from(broadcast),
            timeout_ms,
            rate,
            opts.scan_max_hosts
        );
        let mut ip = start;
        while ip <= end && scanned < opts.scan_max_hosts {
            let target_ip = std::net::Ipv4Addr::from(ip);
            if target_ip != attempt.local_ip {
                let target = SocketAddr::new(IpAddr::V4(target_ip), port);
                let local_addr = SocketAddr::new(IpAddr::V4(attempt.local_ip), 0);
                let mut options = DiscoveryClientOptions::new(
                    target,
                    local_addr,
                    Duration::from_millis(timeout_ms),
                );
                options.attester_registry = opts.attester_registry.clone();
                options = options.disable_multicast().disable_broadcast();
                if let Ok(client) = DiscoveryClient::new(options) {
                    let report = client.discover_with_report(&["alpine-control".to_string()]);
                    append_attempts(diagnostics, &report);
                    if let Ok(outcome) = report.into_result() {
                        return Ok(outcome);
                    }
                }
                scanned += 1;
                sleep(TokioDuration::from_millis(delay_ms)).await;
            }
            ip = ip.saturating_add(1);
        }
    }
    Err(DiscoveryRunError::SubnetScanFailed)
}

struct IfaceAttempt {
    iface: String,
    local_ip: std::net::Ipv4Addr,
    netmask: std::net::Ipv4Addr,
    broadcast: std::net::Ipv4Addr,
}

fn collect_interfaces() -> Result<Vec<IfaceAttempt>, DiscoveryRunError> {
    let mut attempts = Vec::new();
    let ifaces = get_if_addrs::get_if_addrs().map_err(|_| DiscoveryRunError::NoInterfaces)?;
    for iface in ifaces {
        if iface.is_loopback() {
            continue;
        }
        if let get_if_addrs::IfAddr::V4(v4) = iface.addr {
            let ipv4 = v4.ip;
            let maskv4 = v4.netmask;
            let ip_u32 = u32::from_be_bytes(ipv4.octets());
            let mask_u32 = u32::from_be_bytes(maskv4.octets());
            let bcast = ip_u32 | (!mask_u32);
            let bcast_ip = std::net::Ipv4Addr::from(bcast.to_be_bytes());
            attempts.push(IfaceAttempt {
                iface: iface.name,
                local_ip: ipv4,
                netmask: maskv4,
                broadcast: bcast_ip,
            });
        }
    }
    if attempts.len() > 1 {
        let names = attempts
            .iter()
            .map(|iface| iface.iface.as_str())
            .collect::<Vec<_>>()
            .join(", ");
        warn!(
            "[ALPINE][DISCOVERY][WARN] multiple NICs detected (heuristic selection): {}",
            names
        );
    }
    Ok(attempts)
}

fn network_bounds(ip: std::net::Ipv4Addr, netmask: std::net::Ipv4Addr) -> (u32, u32) {
    let ip_u32 = u32::from_be_bytes(ip.octets());
    let mask_u32 = u32::from_be_bytes(netmask.octets());
    let network = ip_u32 & mask_u32;
    let broadcast = network | (!mask_u32);
    (network, broadcast)
}

fn append_attempts(diagnostics: &mut Option<&mut Vec<DiscoveryAttempt>>, report: &DiscoveryResult) {
    if let Some(targets) = diagnostics.as_mut() {
        targets.extend(report.diagnostics.attempts.clone());
    }
}

fn discovery_backoff(policy: &DiscoveryRetryPolicy, attempt: usize) -> u64 {
    let exponent = attempt.saturating_sub(1) as u32;
    let factor = 1u64.checked_shl(exponent).unwrap_or(u64::MAX);
    let delay = policy.backoff_base_ms.saturating_mul(factor);
    delay.min(policy.backoff_max_ms)
}