use crate::protocol::eval::{ValueGroup, Value};

use crate::runtime2::*;

use crate::runtime2::component::{CompCtx, CompId, PortInstruction};

use crate::runtime2::stdlib::internet::*;

use crate::runtime2::poll::*;

use super::component::{self, *};

use super::control_layer::*;

use super::consensus::*;

use std::io::ErrorKind as IoErrorKind;

enum SocketState {

    Connected(SocketTcpClient),

    Error,

}

impl SocketState {

    fn get_socket(&self) -> &SocketTcpClient {

        match self {

            SocketState::Connected(v) => v,

            SocketState::Error => unreachable!(),

        }

    }

}

/// States from the point of view of the component that is connecting to this

/// TCP component (i.e. from the point of view of attempting to interface with

/// a socket).

#[derive(PartialEq, Debug)]

    AwaitingCmd,

    Getting,

    Putting,

    FinishSync,

    FinishSyncThenQuit,

}

pub struct ComponentTcpClient {

    // Properties for the tcp socket

    socket_state: SocketState,

    poll_ticket: Option<PollTicket>,

    inbox_main: InboxMain,

    inbox_backup: InboxBackup,

    pdl_input_port_id: PortId, // input from PDL, so transmitted over socket

    pdl_output_port_id: PortId, // output towards PDL, so received over socket

    input_union_send_tag_value: i64,

    input_union_receive_tag_value: i64,

    input_union_finish_tag_value: i64,

    input_union_shutdown_tag_value: i64,

    // Generic component state

    exec_state: CompExecState,

    control: ControlLayer,

    consensus: Consensus,

    // Temporary variables

    byte_buffer: Vec<u8>,

}

impl Component for ComponentTcpClient {

    fn on_creation(&mut self, id: CompId, sched_ctx: &SchedulerCtx) {

        // Retrieve type information for messages we're going to receive

        let pd = &sched_ctx.runtime.protocol;

        let cmd_type = pd.find_type(b"std.internet", b"Cmd")

            .expect("'Cmd' type in the 'std.internet' module");

        let cmd_type = cmd_type

            .as_union();

        self.input_union_send_tag_value = cmd_type.get_variant_tag_value(b"Send").unwrap();

        self.input_union_receive_tag_value = cmd_type.get_variant_tag_value(b"Receive").unwrap();

        self.input_union_finish_tag_value = cmd_type.get_variant_tag_value(b"Finish").unwrap();

        self.input_union_shutdown_tag_value = cmd_type.get_variant_tag_value(b"Shutdown").unwrap();

        // Register socket for async events

        if let SocketState::Connected(socket) = &self.socket_state {

            let self_handle = sched_ctx.runtime.get_component_public(id);

            let poll_ticket = sched_ctx.polling.register(socket, self_handle, true, true)

                .expect("registering tcp component");

            debug_assert!(self.poll_ticket.is_none());

            self.poll_ticket = Some(poll_ticket);

        }

    }

    fn on_shutdown(&mut self, sched_ctx: &SchedulerCtx) {

        if let Some(poll_ticket) = self.poll_ticket.take() {

            sched_ctx.polling.unregister(poll_ticket)

                .expect("unregistering tcp component");

        }

    }

    fn adopt_message(&mut self, _comp_ctx: &mut CompCtx, message: DataMessage) {

        let slot = &mut self.inbox_main[0];

        if slot.is_none() {

            *slot = Some(message);

        } else {

            self.inbox_backup.push(message);

        }

    }

    fn handle_message(&mut self, sched_ctx: &mut SchedulerCtx, comp_ctx: &mut CompCtx, message: Message) {

        match message {

            Message::Data(message) => {

                self.handle_incoming_data_message(sched_ctx, comp_ctx, message);

            },

            Message::Sync(message) => {

                let decision = self.consensus.receive_sync_message(sched_ctx, comp_ctx, message);

                component::default_handle_sync_decision(sched_ctx, &mut self.exec_state, comp_ctx, decision, &mut self.consensus);

            },

            Message::Control(message) => {

                if let Err(location_and_message) = component::default_handle_control_message(

                    &mut self.exec_state, &mut self.control, &mut self.consensus,

                    message, sched_ctx, comp_ctx, &mut self.inbox_main, &mut self.inbox_backup

                ) {

                    component::default_handle_error_for_builtin(&mut self.exec_state, sched_ctx, location_and_message);

                }

            },

            Message::Poll => {

                sched_ctx.info("Received polling event");

            },

        }

    }

    fn run(&mut self, sched_ctx: &mut SchedulerCtx, comp_ctx: &mut CompCtx) -> CompScheduling {

        sched_ctx.info(&format!("Running component ComponentTcpClient (mode: {:?}, sync state: {:?})", self.exec_state.mode, self.sync_state));

        match self.exec_state.mode {

            CompMode::BlockedSelect |

            CompMode::PutPortsBlockedTransferredPorts |

            CompMode::PutPortsBlockedAwaitingAcks |

            CompMode::PutPortsBlockedSendingPort |

            CompMode::NewComponentBlocked => {

                // Not possible: we never enter this state

                unreachable!();

            },

            CompMode::NonSync => {

                // When in non-sync mode

                match &mut self.socket_state {

                    SocketState::Connected(_socket) => {

                            // Previous request was to let the component shut down

                            self.exec_state.set_as_start_exit(ExitReason::Termination);

                        } else {

                            // Reset for a new request

                            component::default_handle_sync_start(

                                &mut self.exec_state, &mut self.inbox_main, sched_ctx, comp_ctx, &mut self.consensus

                            );

                        }

                        return CompScheduling::Immediate;

                    },

                    SocketState::Error => {

                        // Could potentially send an error message to the

                        // connected component.

                        self.exec_state.set_as_start_exit(ExitReason::ErrorNonSync);

                        return CompScheduling::Immediate;

                    }

                }

            },

            CompMode::Sync => {

                // When in sync mode: wait for a command to come in

                match self.sync_state {

                        match component::default_attempt_get(

                            &mut self.exec_state, self.pdl_input_port_id, PortInstruction::NoSource,

                            &mut self.inbox_main, &mut self.inbox_backup, sched_ctx, comp_ctx,

                            &mut self.control, &mut self.consensus

                        ) {

                            GetResult::Received(message) => {

                                let (tag_value, embedded_heap_pos) = message.content.values[0].as_union();

                                if tag_value == self.input_union_send_tag_value {

                                    // Retrieve bytes from the message

                                    self.byte_buffer.clear();

                                    let union_content = &message.content.regions[embedded_heap_pos as usize];

                                    debug_assert_eq!(union_content.len(), 1);

                                    let array_heap_pos = union_content[0].as_array();

                                    let array_values = &message.content.regions[array_heap_pos as usize];

                                    self.byte_buffer.reserve(array_values.len());

                                    for value in array_values {

                                        self.byte_buffer.push(value.as_uint8());

                                    }

                                } else if tag_value == self.input_union_receive_tag_value {

                                    // Component requires a `recv`

                                } else if tag_value == self.input_union_finish_tag_value {

                                    // Component requires us to end the sync round

                                } else if tag_value == self.input_union_shutdown_tag_value {

                                    // Component wants to close the connection

                                } else {

                                    unreachable!("got tag_value {}", tag_value)

                                }

                                return CompScheduling::Immediate;

                            },

                            GetResult::NoMessage => {

                                return CompScheduling::Sleep;

                            },

                            GetResult::Error(location_and_message) => {

                                component::default_handle_error_for_builtin(&mut self.exec_state, sched_ctx, location_and_message);

                                return CompScheduling::Immediate;

                            }

                        }

                    },

                        // We're supposed to send a user-supplied message fully

                        // over the socket. But we might end up blocking. In

                        // that case the component goes to sleep until it is

                        // polled.

                        let socket = self.socket_state.get_socket();

                        while !self.byte_buffer.is_empty() {

                            match socket.send(&self.byte_buffer) {

                                Ok(bytes_sent) => {

                                    self.byte_buffer.drain(..bytes_sent);

                                },

                                Err(err) => {

                                    if err.kind() == IoErrorKind::WouldBlock {

                                        return CompScheduling::Sleep; // wait until notified

                                    } else {

                                        todo!("handle socket.send error {:?}", err)

                                    }

                                }

                            }

                        }

                        // If here then we're done putting the data, we can

                        // finish the sync round

                        component::default_handle_sync_end(&mut self.exec_state, sched_ctx, comp_ctx, &mut self.consensus);

                        return CompScheduling::Requeue;

                    },

                        // We're going to try and receive a single message. If

                        // this causes us to end up blocking the component

                        // goes to sleep until it is polled.

                        const BUFFER_SIZE: usize = 1024; // TODO: Move to config

                        let socket = self.socket_state.get_socket();

                        self.byte_buffer.resize(BUFFER_SIZE, 0);

                        match socket.receive(&mut self.byte_buffer) {

                            Ok(num_received) => {

                                self.byte_buffer.resize(num_received, 0);

                                let message_content = self.bytes_to_data_message_content(&self.byte_buffer);

                                let send_result = component::default_send_data_message(

                                    &mut self.exec_state, self.pdl_output_port_id, PortInstruction::NoSource,

                                    message_content, sched_ctx, &mut self.consensus, &mut self.control, comp_ctx

                                );

                                if let Err(location_and_message) = send_result {

                                    component::default_handle_error_for_builtin(&mut self.exec_state, sched_ctx, location_and_message);

                                    return CompScheduling::Immediate;

                                } else {

                                    let scheduling = send_result.unwrap();

                                    return scheduling;

                                }

                            },

                            Err(err) => {

                                if err.kind() == IoErrorKind::WouldBlock {

                                    return CompScheduling::Sleep; // wait until polled

                                } else {

                                    todo!("handle socket.receive error {:?}", err)

                                }

                            }

                        }

                    },

                        component::default_handle_sync_end(&mut self.exec_state, sched_ctx, comp_ctx, &mut self.consensus);

                        return CompScheduling::Requeue;

                    },

                }

            },

            CompMode::BlockedGet => {

                // Entered when awaiting a new command

                return CompScheduling::Sleep;

            },

            CompMode::SyncEnd | CompMode::BlockedPut =>

                return CompScheduling::Sleep,

            CompMode::StartExit =>

                return component::default_handle_start_exit(&mut self.exec_state, &mut self.control, sched_ctx, comp_ctx, &mut self.consensus),

            CompMode::BusyExit =>

                return component::default_handle_busy_exit(&mut self.exec_state, &mut self.control, sched_ctx),

            CompMode::Exit =>

                return component::default_handle_exit(&self.exec_state),

        }

    }

}

impl ComponentTcpClient {

    pub(crate) fn new(arguments: ValueGroup) -> Self {

        use std::net::{IpAddr, Ipv4Addr};

        debug_assert_eq!(arguments.values.len(), 4);

        // Parsing arguments

        let ip_heap_pos = arguments.values[0].as_array();

        let ip_elements = &arguments.regions[ip_heap_pos as usize];

        if ip_elements.len() != 4 {

            todo!("friendly error reporting: ip contains 4 octects");

        }

        let ip_address = IpAddr::V4(Ipv4Addr::new(

            ip_elements[0].as_uint8(), ip_elements[1].as_uint8(),

            ip_elements[2].as_uint8(), ip_elements[3].as_uint8()

        ));

        let port = arguments.values[1].as_uint16();

        let input_port = component::port_id_from_eval(arguments.values[2].as_input());

        let output_port = component::port_id_from_eval(arguments.values[3].as_output());

        let socket = SocketTcpClient::new(ip_address, port);

        if let Err(socket) = socket {

            todo!("friendly error reporting: failed to open socket (reason: {:?})", socket);

        }

        return Self{

            socket_state: SocketState::Connected(socket.unwrap()),

            poll_ticket: None,

            inbox_main: vec![None],

            inbox_backup: Vec::new(),

            input_union_send_tag_value: -1,

            input_union_receive_tag_value: -1,

            input_union_finish_tag_value: -1,

            input_union_shutdown_tag_value: -1,

            pdl_input_port_id: input_port,

            pdl_output_port_id: output_port,

            exec_state: CompExecState::new(),

            control: ControlLayer::default(),

            consensus: Consensus::new(),

            byte_buffer: Vec::new(),

        }

    }

    // Handles incoming data from the PDL side (hence, going into the socket)

    fn handle_incoming_data_message(&mut self, sched_ctx: &SchedulerCtx, comp_ctx: &mut CompCtx, message: DataMessage) {

        if self.exec_state.mode.is_in_sync_block() {

            self.consensus.handle_incoming_data_message(comp_ctx, &message);

        }

        match component::default_handle_incoming_data_message(

            &mut self.exec_state, &mut self.inbox_main, comp_ctx, message, sched_ctx, &mut self.control

        ) {

            IncomingData::PlacedInSlot => {},

            IncomingData::SlotFull(message) => {

                self.inbox_backup.push(message);

            }

        }

    }

    fn data_message_to_bytes(&self, message: DataMessage, bytes: &mut Vec<u8>) {

        debug_assert_eq!(message.data_header.target_port, self.pdl_input_port_id);

        debug_assert_eq!(message.content.values.len(), 1);

        if let Value::Array(array_pos) = message.content.values[0] {

            let region = &message.content.regions[array_pos as usize];

            bytes.reserve(region.len());

            for value in region {

                bytes.push(value.as_uint8());

            }

        } else {

            unreachable!();

        }

    }

    fn bytes_to_data_message_content(&self, buffer: &[u8]) -> ValueGroup {

        // Turn bytes into silly executor-style array

        let mut values = Vec::with_capacity(buffer.len());

        for byte in buffer.iter().copied() {

            values.push(Value::UInt8(byte));

        }

        // Put in a value group

        let mut value_group = ValueGroup::default();

        value_group.regions.push(values);

        value_group.values.push(Value::Array(0));

        return value_group;

    }

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

use std::mem::size_of;

use std::io::Error as IoError;

use libc::{

    c_int,

    sockaddr_in, sockaddr_in6, in_addr, in6_addr,

    socket, bind, listen, accept, connect, close,

};

use crate::runtime2::poll::{AsFileDescriptor, FileDescriptor};

#[derive(Debug)]

pub enum SocketError {

    Opening,

    Modifying,

    Binding,

    Listening,

    Connecting,

    Accepted,

    Accepting,

}

enum SocketState {

    Opened,

    Listening,

}

pub struct SocketTcpClient {

    socket_handle: libc::c_int,

    is_blocking: bool,

}

impl SocketTcpClient {

    pub fn new(ip: IpAddr, port: u16) -> Result<Self, SocketError> {

        let socket_handle = create_and_connect_socket(

            libc::SOCK_STREAM, libc::IPPROTO_TCP, ip, port

        )?;

            unsafe{ libc::close(socket_handle); }

            return Err(SocketError::Modifying);

        }

        return Ok(SocketTcpClient{

            socket_handle,

        })

    }

    pub fn send(&self, message: &[u8]) -> Result<usize, IoError> {

        let result = unsafe{

            let message_pointer = message.as_ptr().cast();

            libc::send(self.socket_handle, message_pointer, message.len() as libc::size_t, 0)

        };

        if result < 0 {

            return Err(IoError::last_os_error());

        }

        return Ok(result as usize);

    }

    /// Receives data from the TCP socket. Returns the number of bytes received.

    /// More bytes may be present even thought `used < buffer.len()`.

    pub fn receive(&self, buffer: &mut [u8]) -> Result<usize, IoError> {

        let result = unsafe {

            let message_pointer = buffer.as_mut_ptr().cast();

            libc::recv(self.socket_handle, message_pointer, buffer.len(), 0)

        };

        if result < 0 {

            return Err(IoError::last_os_error());

        }

        return Ok(result as usize);

    }

}

impl Drop for SocketTcpClient {

    fn drop(&mut self) {

        debug_assert!(self.socket_handle >= 0);

        unsafe{ close(self.socket_handle) };

    }

}

impl AsFileDescriptor for SocketTcpClient {

    fn as_file_descriptor(&self) -> FileDescriptor {

        return self.socket_handle;

    }

}

/// Raw socket receiver. Essentially a listener that accepts a single connection

struct SocketRawRx {

    listen_handle: c_int,

    accepted_handle: c_int,

}

impl SocketRawRx {

    pub fn new(ip: Option<Ipv4Addr>, port: u16) -> Result<Self, SocketError> {

        let ip = ip.unwrap_or(Ipv4Addr::UNSPECIFIED); // unspecified is the same as INADDR_ANY

        let address = unsafe{ in_addr{

            s_addr: std::mem::transmute(ip.octets()),

        }};

        let socket_address = sockaddr_in{

            sin_family: libc::AF_INET as libc::sa_family_t,

            sin_port: htons(port),

            sin_addr: address,

            sin_zero: [0; 8],

        };

        unsafe {

            let socket_handle = create_and_bind_socket(libc::SOCK_RAW, 0, IpAddr::V4(ip), port)?;

            let result = listen(socket_handle, 3);

            if result < 0 { return Err(SocketError::Listening); }

            return Ok(SocketRawRx{

                listen_handle: socket_handle,

                accepted_handle: -1,

            });

        }

    }

    // pub fn try_accept(&mut self, timeout_ms: u32) -> Result<(), SocketError> {

    //     if self.accepted_handle >= 0 {

    //         // Already accepted a connection

    //         return Err(SocketError::Accepted);

    //     }

    //

    //     let mut socket_address = sockaddr_in{

    //         sin_family: 0,

    //         sin_port: 0,

    //         sin_addr: in_addr{ s_addr: 0 },

    //         sin_zero: [0; 8]

    //     };

    //     let mut size = size_of::<sockaddr_in>() as u32;

    //     unsafe {

    //         let result = accept(self.listen_handle, &mut socket_address as *mut _, &mut size as *mut _);

    //         if result < 0 {

                let socket_pointer = &socket_address as *const sockaddr_in;

                bind(socket_handle, socket_pointer.cast(), address_size)

            },

            IpAddr::V6(ip) => {

                let (socket_address, address_size) = create_sockaddr_in_v6(ip, port);

                let socket_pointer= &socket_address as *const sockaddr_in6;

                bind(socket_handle, socket_pointer.cast(), address_size)

            }

        };

        if result < 0 {

            close(socket_handle);

            return Err(SocketError::Binding);

        }

        return Ok(socket_handle);

    }

}

/// Performs the `socket` and `connect` calls

fn create_and_connect_socket(socket_type: libc::c_int, protocol: libc::c_int, ip: IpAddr, port: u16) -> Result<libc::c_int, SocketError> {

    let family = socket_family_from_ip(ip);

    unsafe {

        let socket_handle = socket(family, socket_type, protocol);

        if socket_handle < 0 {

            return Err(SocketError::Opening);

        }

        let result = match ip {

            IpAddr::V4(ip) => {

                let (socket_address, address_size) = create_sockaddr_in_v4(ip, port);

                let socket_pointer = &socket_address as *const sockaddr_in;

                connect(socket_handle, socket_pointer.cast(), address_size)

            },

            IpAddr::V6(ip) => {

                let (socket_address, address_size) = create_sockaddr_in_v6(ip, port);

                let socket_pointer= &socket_address as *const sockaddr_in6;

                connect(socket_handle, socket_pointer.cast(), address_size)

            }

        };

        if result < 0 {

            close(socket_handle);

            return Err(SocketError::Connecting);

        }

        return Ok(socket_handle);

    }

}

#[inline]

fn create_sockaddr_in_v4(ip: Ipv4Addr, port: u16) -> (sockaddr_in, libc::socklen_t) {

    let address = unsafe{

        in_addr{

            s_addr: std::mem::transmute(ip.octets())

        }

    };

    let socket_address = sockaddr_in{

        sin_family: libc::AF_INET as libc::sa_family_t,

        sin_port: htons(port),

        sin_addr: address,

        sin_zero: [0; 8]

    };

    let address_size = size_of::<sockaddr_in>();

    return (socket_address, address_size as _);

}

#[inline]

fn create_sockaddr_in_v6(ip: Ipv6Addr, port: u16) -> (sockaddr_in6, libc::socklen_t) {

    // flow label is advised to be, according to RFC6437 a (somewhat

    // secure) random number taken from a uniform distribution

    let flow_info = rand::random();

    let address = unsafe{

        in6_addr{

            s6_addr: ip.octets()

        }

    };

    let socket_address = sockaddr_in6{

        sin6_family: libc::AF_INET6 as libc::sa_family_t,

        sin6_port: htons(port),

        sin6_flowinfo: flow_info,

        sin6_addr: address,

        sin6_scope_id: 0, // incorrect in case of loopback address

    };

    let address_size = size_of::<sockaddr_in6>();

    return (socket_address, address_size as _);

}

#[inline]

fn set_socket_blocking(handle: libc::c_int, blocking: bool) -> bool {

    if handle < 0 {

        return false;

    }

    unsafe{

        let mut flags = libc::fcntl(handle, libc::F_GETFL, 0);

        if flags < 0 {

            return false;

        }

        if blocking {

            flags &= !libc::O_NONBLOCK;

        } else {

            flags |= libc::O_NONBLOCK;

        }

        let result = libc::fcntl(handle, libc::F_SETFL, flags);

        if result < 0 {

            return false;

        }

    }

    return true;

}

#[inline]

fn socket_family_from_ip(ip: IpAddr) -> libc::c_int {

    return match ip {

        IpAddr::V4(_) => libc::AF_INET,

        IpAddr::V6(_) => libc::AF_INET6,

    };

}

#[inline]

fn htons(port: u16) -> u16 {

    return port.to_be();