#![feature(yggdrasil_os)]
use std::{
    io,
    net::{Ipv4Addr, SocketAddr, ToSocketAddrs, UdpSocket},
    os::fd::AsRawFd,
    process::ExitCode,
    time::{Duration, SystemTime},
};

use bytemuck::{Pod, Zeroable};
use chrono::{DateTime, Timelike, Utc};
use clap::Parser;
use cross::io::{Poll, TimerFd};

const NTP_TIMESTAMP_EPOCH: u32 = 2208988800;

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("Timed out")]
    TimedOut,
    #[error("Invalid response version")]
    InvalidResponseVersion,
    #[error("Invalid response mode")]
    InvalidResponseMode,
    #[error("Invalid response origin timestamp")]
    InvalidResponseOriginTimestamp,
    #[error("Invalid response timestamp")]
    InvalidResponseTimestamp,
    #[error("Invalid server address")]
    InvalidServerAddress,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct NtpPacketHeader {
    status: u8,
    ignored: [u8; 3],
    root_delay: u32,
    root_dispersion: u32,
    reference_id: u32,
    reference_timestamp: u64,
    origin_timestamp: u64,
    receive_timestamp: u64,
    transmit_timestamp: u64,
}

const NANOSECONDS_IN_SECOND: u64 = 1_000_000_000;
const MICROSECONDS_IN_SECOND: u64 = 1_000_000;

fn timestamp_now() -> u64 {
    let now = SystemTime::now();
    let duration = now.duration_since(SystemTime::UNIX_EPOCH).unwrap();
    let seconds = duration
        .as_secs()
        .saturating_add(NTP_TIMESTAMP_EPOCH as u64);
    let microseconds = duration.subsec_micros() as u64;
    let fraction = microseconds * u32::MAX as u64 / MICROSECONDS_IN_SECOND;
    (seconds << 32) | fraction
}

fn time_from_timestamp(ts: u64) -> Result<DateTime<Utc>, Error> {
    let mut seconds = (ts >> 32).saturating_sub(NTP_TIMESTAMP_EPOCH as u64);
    let mut seconds_fraction = ts & 0xFFFFFFFF;
    seconds += seconds_fraction / u32::MAX as u64;
    seconds_fraction %= u32::MAX as u64;
    let nanoseconds = seconds_fraction * NANOSECONDS_IN_SECOND / (u32::MAX as u64);
    DateTime::from_timestamp(seconds as i64, nanoseconds as u32)
        .ok_or(Error::InvalidResponseTimestamp)
}

fn handle_response(response: NtpPacketHeader, timestamp: u64) -> Result<DateTime<Utc>, Error> {
    let response_version = (response.status >> 3) & 0x7;
    let response_mode = response.status & 0x7;
    if response_version != 4 {
        return Err(Error::InvalidResponseVersion);
    }
    if response_mode != 4 {
        return Err(Error::InvalidResponseMode);
    }
    let origin_timestamp = u64::from_be(response.origin_timestamp);
    if origin_timestamp != timestamp {
        eprintln!("{origin_timestamp} != {timestamp}");
        return Err(Error::InvalidResponseOriginTimestamp);
    }

    let timestamp = u64::from_be(response.transmit_timestamp);
    time_from_timestamp(timestamp)
}

fn send_request(socket: &UdpSocket, server: &SocketAddr) -> Result<u64, Error> {
    let send_timestamp = timestamp_now();
    let client_packet = NtpPacketHeader {
        status: (4u8 << 3) | 3,
        ignored: [0; 3],
        root_delay: 0,
        root_dispersion: 0,
        reference_id: 0,
        reference_timestamp: 0,
        origin_timestamp: 0,
        receive_timestamp: 0,
        transmit_timestamp: send_timestamp.to_be(),
    };
    socket.send_to(bytemuck::bytes_of(&client_packet), server)?;
    Ok(send_timestamp)
}

fn do_request(server: SocketAddr) -> Result<DateTime<Utc>, Error> {
    const ATTEMPTS: usize = 5;
    const TIMEOUT_MAX: u64 = 3000;

    let mut poll = Poll::new()?;
    let mut timer = TimerFd::new()?;
    let mut buffer = NtpPacketHeader::zeroed();

    let socket = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0))?;
    let mut timeout = 500;
    let mut response = None;

    poll.add(&timer)?;
    poll.add(&socket)?;

    for attempt in 0..ATTEMPTS {
        let send_timestamp = send_request(&socket, &server)?;
        timer.start(Duration::from_millis(timeout))?;
        timeout = (timeout * 2).min(TIMEOUT_MAX);

        // Wait for response
        loop {
            let event = poll.wait(None)?;
            match event {
                Some(fd) if fd == socket.as_raw_fd() => {
                    let (len, remote) = match socket.recv_from(bytemuck::bytes_of_mut(&mut buffer))
                    {
                        Ok(result) => result,
                        Err(error) => {
                            eprintln!("Receive error: {error}");
                            continue;
                        }
                    };
                    if remote != server || len < size_of::<NtpPacketHeader>() {
                        eprintln!("Remote != server or len too small");
                        continue;
                    }
                    let t = handle_response(buffer, send_timestamp)
                        .inspect_err(|e| {
                            eprintln!("Response from {server}: {e}");
                        })
                        .ok();
                    if let Some(t) = t {
                        response = Some(t);
                        break;
                    }
                }
                Some(fd) if fd == timer.as_raw_fd() => {
                    break;
                }
                _ => unreachable!(),
            }
        }

        if response.is_some() {
            break;
        }

        eprintln!("[{}/{}] Timed out", attempt + 1, ATTEMPTS);
    }

    response.ok_or(Error::TimedOut)
}

#[derive(Debug, Parser)]
struct Args {
    #[clap(short, long, help = "Run in a loop, fetching new time every N seconds")]
    interval: Option<u64>,
    #[clap(
        short,
        long,
        help = "Only print received time, do not change any settings"
    )]
    dry_run: bool,
    server: String,
}

fn handle_time(time: DateTime<Utc>) -> Result<(), Error> {
    let seconds = time.timestamp() as u64;
    let nanoseconds = time.nanosecond();
    cross::time::set_real_time(seconds, nanoseconds)?;
    Ok(())
}

fn resolve_server(args: &Args) -> Result<SocketAddr, Error> {
    let mut server = args.server.to_socket_addrs()?;
    let server = server.next().ok_or(Error::InvalidServerAddress)?;
    Ok(server)
}

fn run_loop(args: &Args, interval: Duration) -> Result<(), Error> {
    let server = resolve_server(args)?;

    loop {
        let response = do_request(server);

        match response {
            Ok(time) if !args.dry_run => {
                if let Err(error) = handle_time(time) {
                    log::error!("handle_time(): {error}");
                }
            }
            Ok(time) => {
                log::info!("Current time: {time}");
            }
            Err(error) => {
                log::error!("{}: {error}", args.server);
            }
        }

        std::thread::sleep(interval);
    }
}

fn run_once(args: &Args) -> Result<(), Error> {
    let server = resolve_server(args)?;
    let response = do_request(server)?;
    if !args.dry_run {
        handle_time(response)?;
    }
    Ok(())
}

fn run(args: &Args) -> Result<(), Error> {
    match args.interval {
        Some(interval) => {
            let interval = Duration::from_secs(interval);
            run_loop(args, interval)
        }
        None => run_once(args),
    }
}

fn main() -> ExitCode {
    logsink::setup_logging(true);
    let args = Args::parse();
    match run(&args) {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::FAILURE
        }
    }
}
