pub mod command;

use crate::scheduler::{CyclicTask, TaskHandle};
use crate::state_vector::{SectionIdentifier, SectionWriter, StateVector};
use anyhow::{Result, ensure};
use log::{error, trace, warn};
use nautilus_common::command::{Command, CommandHeader};
pub(crate) use nautilus_common::telemetry::{CommsState, Telemetry, TelemetryMessage};
use nautilus_common::udp::{UdpRecvPostcardError, UdpSocketExt};
use std::any::type_name;
use std::collections::HashMap;
use std::fmt::{Debug, Formatter};
use std::net::{IpAddr, Ipv4Addr, SocketAddr, ToSocketAddrs, UdpSocket};
use std::sync::mpsc::Receiver;
use std::time::Instant;

pub type TelemetrySender = TaskHandle<Telemetry, SectionIdentifier>;

impl TelemetrySender {
    pub fn send(&self, telemetry_message: TelemetryMessage) {
        trace!("TelemetrySender::send(self: {self:?}, telemetry_message: {telemetry_message:?}");
        // Ignore failure
        let _ = self.sender.send(Telemetry {
            timestamp: chrono::Utc::now(),
            message: telemetry_message,
        });
    }
}

type CommandCallback<'a> = Box<dyn FnMut(&[u8]) -> Result<()> + Send + 'a>;

pub struct CommsTask<'a, A>
where
    A: ToSocketAddrs + Debug,
{
    udp: UdpSocket,
    ground_address: A,
    command_callbacks: HashMap<String, CommandCallback<'a>>,
    section_writer: SectionWriter<'a, CommsState>,
}

impl<A> Debug for CommsTask<'_, A>
where
    A: ToSocketAddrs + Debug,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "CommsTask {{ udp: {:?}, ground_address: {:?} }}",
            self.udp, self.ground_address
        )
    }
}

impl<'a, A> CommsTask<'a, A>
where
    A: ToSocketAddrs + Debug,
{
    pub fn new<'b: 'a>(
        local_port: u16,
        ground_address: A,
        state_vector: &'b StateVector,
    ) -> Result<Self> {
        trace!(
            "CommsTask::new<A={}>(local_port: {local_port}, ground_address: {ground_address:?})",
            type_name::<A>()
        );
        let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), local_port);
        // let bind_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), local_port);
        let udp = UdpSocket::bind(bind_addr)?;
        udp.set_nonblocking(true)?;

        let section_writer = state_vector.create_section(CommsState::default());
        Ok(Self {
            udp,
            ground_address,
            command_callbacks: HashMap::new(),
            section_writer,
        })
    }

    pub fn add_command_handler<T: Command>(
        &mut self,
        command: impl Into<String>,
        mut handler: impl FnMut(T) + Send + 'a,
    ) -> Result<()> {
        let command = command.into();
        ensure!(
            !self.command_callbacks.contains_key(&command),
            "Already Contains Command {command}"
        );
        self.command_callbacks.insert(
            command.clone(),
            Box::new(move |cmd_buf| {
                let (cmd, remainder) = postcard::take_from_bytes::<T>(cmd_buf)?;
                ensure!(
                    remainder.is_empty(),
                    "{command} received {} extra bytes",
                    remainder.len()
                );
                handler(cmd);
                Ok(())
            }),
        );
        Ok(())
    }
}

impl<A> CyclicTask for CommsTask<'_, A>
where
    A: ToSocketAddrs + Debug,
{
    type Message = Telemetry;
    type Data = SectionIdentifier;

    fn get_data(&self) -> Self::Data {
        trace!(
            "CommsTask<A={}>::get_data(self: {self:?})",
            type_name::<A>()
        );

        self.section_writer.get_identifier()
    }

    fn step(&mut self, receiver: &Receiver<Self::Message>, step_time: Instant) {
        trace!(
            "CommsTask<A={}>::step(self: {self:?}, receiver: {receiver:?}, step_time: {step_time:?})",
            type_name::<A>()
        );
        let mut buffer = [0u8; 512];

        match self.udp.recv_postcard::<CommandHeader>(&mut buffer) {
            Ok((cmd, _, size)) => {
                self.section_writer.update(|state| {
                    state.rx_packets = state.rx_packets.wrapping_add(1);
                    state.rx_bytes = state
                        .rx_bytes
                        .wrapping_add(u32::try_from(size % (u32::MAX as usize)).unwrap_or(0));
                });
                match self.command_callbacks.get_mut(cmd.name) {
                    Some(handler) => {
                        if let Err(e) = handler(cmd.data) {
                            error!("Command Error: {e}");
                        }
                    }
                    None => {
                        warn!("Unknown Command: {}", cmd.name);
                    }
                }
            }
            Err(UdpRecvPostcardError::NoData) => {}
            Err(err) => {
                error!("Rx error: {err}");
                self.section_writer.update(|state| {
                    state.rx_errors = state.rx_errors.wrapping_add(1);
                });
            }
        }

        // Intentionally ignore Err case
        while let Ok(tlm) = receiver.try_recv() {
            match self
                .udp
                .send_postcard(&tlm, &mut buffer, &self.ground_address)
            {
                Ok(bytes_sent) => self.section_writer.update(|state| {
                    state.tx_packets = state.tx_packets.wrapping_add(1);
                    state.tx_bytes = state
                        .tx_bytes
                        .wrapping_add(u32::try_from(bytes_sent % (u32::MAX as usize)).unwrap_or(0));
                }),
                Err(err) => {
                    error!("Tx Error: {err}");
                    self.section_writer.update(|state| {
                        state.tx_errors = state.tx_errors.wrapping_add(1);
                    });
                }
            }
        }
    }
}
