use super::vec_storage::VecStorage;

use crate::common::*;

use crate::runtime::endpoint::EndpointExt;

use crate::runtime::endpoint::EndpointInfo;

use crate::runtime::endpoint::{Endpoint, Msg, SetupMsg};

use crate::runtime::errors::EndpointErr;

use crate::runtime::errors::MessengerRecvErr;

use crate::runtime::errors::PollDeadlineErr;

use crate::runtime::MessengerState;

use crate::runtime::Messengerlike;

use crate::runtime::ReceivedMsg;

use crate::runtime::{ProtocolD, ProtocolS};

use std::net::SocketAddr;

use std::sync::Arc;

pub enum Coupling {

    Active,

    Passive,

}

#[derive(Debug)]

struct Family {

    parent: Option<Port>,

    children: HashSet<Port>,

}

pub struct Binding {

    pub coupling: Coupling,

    pub polarity: Polarity,

    pub addr: SocketAddr,

}

pub struct InPort(Port); // InPort and OutPort are AFFINE (exposed to Rust API)

pub struct OutPort(Port);

impl From<InPort> for Port {

    fn from(x: InPort) -> Self {

        x.0

    }

}

impl From<OutPort> for Port {

    fn from(x: OutPort) -> Self {

        x.0

    }

}

#[derive(Default, Debug)]

struct ChannelIndexStream {

    next: u32,

}

impl ChannelIndexStream {

    fn next(&mut self) -> u32 {

        self.next += 1;

        self.next - 1

    }

}

enum Connector {

    Connecting(Connecting),

    Connected(Connected),

}

#[derive(Default)]

pub struct Connecting {

    bindings: Vec<Binding>,

}

trait Binds<T> {

    fn bind(&mut self, coupling: Coupling, addr: SocketAddr) -> T;

}

impl Binds<InPort> for Connecting {

    fn bind(&mut self, coupling: Coupling, addr: SocketAddr) -> InPort {

        self.bindings.push(Binding { coupling, polarity: Polarity::Getter, addr });

        InPort(Port(self.bindings.len() - 1))

    }

}

impl Binds<OutPort> for Connecting {

    fn bind(&mut self, coupling: Coupling, addr: SocketAddr) -> OutPort {

        self.bindings.push(Binding { coupling, polarity: Polarity::Putter, addr });

        OutPort(Port(self.bindings.len() - 1))

    }

}

#[derive(Debug, Clone)]

pub enum ConnectErr {

    BindErr(SocketAddr),

    NewSocketErr(SocketAddr),

    AcceptErr(SocketAddr),

    ConnectionShutdown(SocketAddr),

    EndpointErr(Port, EndpointErr),

    PollInitFailed,

    PollingFailed,

    Timeout,

}

#[derive(Debug)]

struct Component {

    protocol: Arc<ProtocolD>,

    port_set: HashSet<Port>,

    identifier: Arc<[u8]>,

    state: ProtocolS,

}

impl From<PollDeadlineErr> for ConnectErr {

    fn from(e: PollDeadlineErr) -> Self {

        use PollDeadlineErr as P;

        match e {

            P::PollingFailed => Self::PollingFailed,

            P::Timeout => Self::Timeout,

        }

    }

}

impl From<MessengerRecvErr> for ConnectErr {

    fn from(e: MessengerRecvErr) -> Self {

        use MessengerRecvErr as M;

        match e {

            M::PollingFailed => Self::PollingFailed,

            M::EndpointErr(port, err) => Self::EndpointErr(port, err),

        }

    }

}

impl Connecting {

    fn random_controller_id() -> ControllerId {

        type Bytes8 = [u8; std::mem::size_of::<ControllerId>()];

        let mut bytes = Bytes8::default();

        getrandom::getrandom(&mut bytes).unwrap();

        unsafe {

            // safe:

            // 1. All random bytes give valid Bytes8

            // 2. Bytes8 and ControllerId have same valid representations

            std::mem::transmute::<Bytes8, ControllerId>(bytes)

        }

    }

    fn test_stream_connectivity(stream: &mut TcpStream) -> bool {

        use std::io::Write;

        stream.write(&[]).is_ok()

    }

    fn new_connected(

        &self,

        controller_id: ControllerId,

        timeout: Option<Duration>,

    ) -> Result<Connected, ConnectErr> {

        use ConnectErr::*;

        ///////////////////////////////////////////////////////

        // 1. bindings correspond with ports 0..bindings.len(). For each:

        //    - reserve a slot in endpoint_exts.

        //    - store the port in `native_ports' set.

        let mut endpoint_exts = VecStorage::<EndpointExt>::with_reserved_range(self.bindings.len());

        let native_ports = (0..self.bindings.len()).map(Port).collect();

        // 2. create MessengerState structure for polling channels

        let edge = PollOpt::edge();

        let [ready_r, ready_w] = [Ready::readable(), Ready::writable()];

        let mut ms =

            MessengerState::with_event_capacity(self.bindings.len()).map_err(|_| PollInitFailed)?;

        // 3. create one TODO task per (port,binding) as a vector with indices in lockstep.

        //    we will drain it gradually so we store elements of type Option<Todo> where all are initially Some(_)

        enum Todo {

            PassiveAccepting { listener: TcpListener, channel_id: ChannelId },

            ActiveConnecting { stream: TcpStream },

            PassiveConnecting { stream: TcpStream, channel_id: ChannelId },

            ActiveRecving { endpoint: Endpoint },

        }

        let mut channel_index_stream = ChannelIndexStream::default();

        let mut todos = self

            .bindings

            .iter()

            .enumerate()

            .map(|(index, binding)| {

                Ok(Some(match binding.coupling {

                    Coupling::Passive => {

                        let channel_index = channel_index_stream.next();

                        let channel_id = ChannelId { controller_id, channel_index };

                        let listener =

                            TcpListener::bind(&binding.addr).map_err(|_| BindErr(binding.addr))?;

                        ms.poll.register(&listener, Token(index), ready_r, edge).unwrap(); // registration unique

                        Todo::PassiveAccepting { listener, channel_id }

                    }

                    Coupling::Active => {

                        let stream = TcpStream::connect(&binding.addr)

                            .map_err(|_| NewSocketErr(binding.addr))?;

                        ms.poll.register(&stream, Token(index), ready_w, edge).unwrap(); // registration unique

                        Todo::ActiveConnecting { stream }

                    }

                }))

            })

            .collect::<Result<Vec<Option<Todo>>, ConnectErr>>()?;

        let mut num_todos_remaining = todos.len();

        // 4. handle incoming events until all TODOs are completed OR we timeout

        let deadline = timeout.map(|t| Instant::now() + t);

        let mut polled_undrained_later = IndexSet::<_>::default();

        let mut backoff_millis = 10;

        while num_todos_remaining > 0 {

            ms.poll_events_until(deadline)?;

            for event in ms.events.iter() {

                let token = event.token();

                let index = token.0;

                let binding = &self.bindings[index];

                match todos[index].take() {

                    None => {

                        polled_undrained_later.insert(index);

                    }

                    Some(Todo::PassiveAccepting { listener, channel_id }) => {

                        let (stream, _peer_addr) =

                            listener.accept().map_err(|_| AcceptErr(binding.addr))?;

                        ms.poll.deregister(&listener).expect("wer");

                        ms.poll.register(&stream, token, ready_w, edge).expect("3y5");

                        todos[index] = Some(Todo::PassiveConnecting { stream, channel_id });

                    }

                    Some(Todo::ActiveConnecting { mut stream }) => {

                        let todo = if Self::test_stream_connectivity(&mut stream) {

                            ms.poll.reregister(&stream, token, ready_r, edge).expect("52");

                            let endpoint = Endpoint::from_fresh_stream(stream);

                            Todo::ActiveRecving { endpoint }

                        } else {

                            ms.poll.deregister(&stream).expect("wt");

                            std::thread::sleep(Duration::from_millis(backoff_millis));

                            backoff_millis = ((backoff_millis as f32) * 1.2) as u64 + 3;

                            let stream = TcpStream::connect(&binding.addr).unwrap();

                            ms.poll.register(&stream, token, ready_w, edge).expect("PAC 3");

                            Todo::ActiveConnecting { stream }

                        };

                        todos[index] = Some(todo);

                    }

                    Some(Todo::PassiveConnecting { mut stream, channel_id }) => {

                        if !Self::test_stream_connectivity(&mut stream) {

                            return Err(ConnectionShutdown(binding.addr));

                        }

                        ms.poll.reregister(&stream, token, ready_r, edge).expect("55");

                        let polarity = binding.polarity;

                        let info = EndpointInfo { polarity, channel_id };

                        let msg = Msg::SetupMsg(SetupMsg::ChannelSetup { info });

                        let mut endpoint = Endpoint::from_fresh_stream(stream);

                        endpoint.send(msg).map_err(|e| EndpointErr(Port(index), e))?;

                        let endpoint_ext = EndpointExt { endpoint, info };

                        endpoint_exts.occupy_reserved(index, endpoint_ext);

                        num_todos_remaining -= 1;

                    }

                    Some(Todo::ActiveRecving { mut endpoint }) => {

                        'recv_loop: while let Some(msg) =

                        {

                            if let Msg::SetupMsg(SetupMsg::ChannelSetup { info }) = msg {

                                if info.polarity == binding.polarity {

                                }

                                let channel_id = info.channel_id;

                                let info = EndpointInfo { polarity: binding.polarity, channel_id };

                                ms.polled_undrained.insert(ekey);

                                let endpoint_ext = EndpointExt { endpoint, info };

                                endpoint_exts.occupy_reserved(index, endpoint_ext);

                                num_todos_remaining -= 1;

                                break 'recv_loop;

                            } else {

                                ms.delayed.push(ReceivedMsg { recipient: ekey, msg });

                            }

                        }

                    }

                }

            }

        }

        assert_eq!(None, endpoint_exts.iter_reserved().next());

        drop(todos);

        ///////////////////////////////////////////////////////

        // 1. construct `family', i.e. perform the sink tree setup procedure

        use {Msg::SetupMsg as S, SetupMsg::*};

        let mut messenger = (&mut ms, &mut endpoint_exts);

        impl Messengerlike for (&mut MessengerState, &mut VecStorage<EndpointExt>) {

            fn get_state_mut(&mut self) -> &mut MessengerState {

                self.0

            }

            fn get_endpoint_mut(&mut self, ekey: Key) -> &mut Endpoint {

                &mut self

                    .1

                    .get_occupied_mut(ekey.to_raw() as usize)

                    .expect("OUT OF BOUNDS")

                    .endpoint

            }

        }

        // 1. broadcast my ID as the first echo. await reply from all in net_keylist

        let neighbors = (0..self.bindings.len()).map(Port);

        let echo = S(LeaderEcho { maybe_leader: controller_id });

        let mut awaiting = IndexSet::<Port>::with_capacity(neighbors.len());

        for n in neighbors.clone() {

            messenger.send(n, echo.clone()).map_err(|e| EndpointErr(n, e))?;

            awaiting.insert(n);

        }

        // 2. Receive incoming replies. whenever a higher-id echo arrives,

        //    adopt it as leader, sender as parent, and reset the await set.

        let mut parent: Option<Port> = None;

        let mut my_leader = controller_id;

        messenger.undelay_all();

        'echo_loop: while !awaiting.is_empty() || parent.is_some() {

            let ReceivedMsg { recipient, msg } = messenger.recv_until(deadline)?.ok_or(Timeout)?;

            match msg {

                S(LeaderAnnounce { leader }) => {

                    // someone else completed the echo and became leader first!

                    // the sender is my parent

                    parent = Some(recipient);

                    my_leader = leader;

                    awaiting.clear();

                    break 'echo_loop;

                }

                S(LeaderEcho { maybe_leader }) => {

                    use Ordering::*;

                    match maybe_leader.cmp(&my_leader) {

                        Less => { /* ignore */ }

                        Equal => {

                            awaiting.remove(&recipient);

                            if awaiting.is_empty() {

                                if let Some(p) = parent {

                                    // return the echo to my parent

                                    messenger

                                        .send(p, S(LeaderEcho { maybe_leader }))

                                        .map_err(|e| EndpointErr(p, e))?;

                                } else {

                                    // DECIDE!

                                    break 'echo_loop;

                                }

                            }

                        }

                        Greater => {

                            // join new echo

                            parent = Some(recipient);

                            my_leader = maybe_leader;

                            let echo = S(LeaderEcho { maybe_leader: my_leader });

                            awaiting.clear();

                            if neighbors.len() == 1 {

                                // immediately reply to parent

                                messenger

                                    .send(recipient, echo.clone())

                                    .map_err(|e| EndpointErr(recipient, e))?;

                            } else {

                                for n in neighbors.clone() {

                                    if n != recipient {

                                        messenger

                                            .send(n, echo.clone())

                                            .map_err(|e| EndpointErr(n, e))?;

                                        awaiting.insert(n);

                                    }

                                }

                            }

                        }

                    }

                }

                msg => messenger.delay(ReceivedMsg { recipient, msg }),

            }

        }

        match parent {

            None => assert_eq!(

                my_leader, controller_id,

                "I've got no parent, but I consider {:?} the leader?",

                my_leader

            ),

            Some(parent) => assert_ne!(

                my_leader, controller_id,

                "I have {:?} as parent, but I consider myself ({:?}) the leader?",

                parent, controller_id

            ),

        }

        // 3. broadcast leader announcement (except to parent: confirm they are your parent)

        //    in this loop, every node sends 1 message to each neighbor

        let msg_for_non_parents = S(LeaderAnnounce { leader: my_leader });

        for n in neighbors.clone() {

            let msg =

                if Some(n) == parent { S(YouAreMyParent) } else { msg_for_non_parents.clone() };

            messenger.send(n, msg).map_err(|e| EndpointErr(n, e))?;

        }

        // await 1 message from all non-parents

        for n in neighbors.clone() {

            if Some(n) != parent {

                awaiting.insert(n);

            }

        }

        let mut children = HashSet::default();

        messenger.undelay_all();

        while !awaiting.is_empty() {

            let ReceivedMsg { recipient, msg } = messenger.recv_until(deadline)?.ok_or(Timeout)?;

            let recipient = recipient;

            match msg {

                S(YouAreMyParent) => {

                    assert!(awaiting.remove(&recipient));

                    children.insert(recipient);

                }

                S(SetupMsg::LeaderAnnounce { leader }) => {

                    assert!(awaiting.remove(&recipient));

                    assert!(leader == my_leader);

                    assert!(Some(recipient) != parent);

                    // they wouldn't send me this if they considered me their parent

                }

                _ => messenger.delay(ReceivedMsg { recipient, msg }),

            }

        }

        let family = Family { parent, children };

        // done!

        Ok(Connected {

            components: Default::default(),

            controller_id,

            channel_index_stream,

            endpoint_exts,

            native_ports,

            family,

        })

    }

    /////////

    pub fn connect_using_id(

        &mut self,

        controller_id: ControllerId,

        timeout: Option<Duration>,

    ) -> Result<Connected, ConnectErr> {

        // 1. try and create a connection from these bindings with self immutable.

        let connected = self.new_connected(controller_id, timeout)?;

        // 2. success! drain self and return

        self.bindings.clear();

        Ok(connected)

    }

    pub fn connect(&mut self, timeout: Option<Duration>) -> Result<Connected, ConnectErr> {

        self.connect_using_id(Self::random_controller_id(), timeout)

    }

}

#[derive(Debug)]

pub struct Connected {

    native_ports: HashSet<Port>,

    controller_id: ControllerId,

    channel_index_stream: ChannelIndexStream,

    endpoint_exts: VecStorage<EndpointExt>,

    components: VecStorage<Component>,

    family: Family,

}

impl Connected {

    pub fn new_component(

        &mut self,

        protocol: &Arc<ProtocolD>,

        identifier: &Arc<[u8]>,

        moved_port_list: &[Port],

    ) -> Result<(), MainComponentErr> {

        //////////////////////////////////////////

        // 1. try and create a new component (without mutating self)

        use MainComponentErr::*;

        let moved_port_set = {

            let mut set: HashSet<Port> = Default::default();

            for &port in moved_port_list.iter() {

                if !self.native_ports.contains(&port) {

                    return Err(CannotMovePort(port));

                }

                if !set.insert(port) {

                    return Err(DuplicateMovedPort(port));

                }

            }

            set

        };

        // moved_port_set is disjoint to native_ports

        let expected_polarities = protocol.component_polarities(identifier)?;

        if moved_port_list.len() != expected_polarities.len() {

            return Err(WrongNumberOfParamaters { expected: expected_polarities.len() });

        }

        // correct polarity list

        for (param_index, (&port, &expected_polarity)) in

            moved_port_list.iter().zip(expected_polarities.iter()).enumerate()

        {

            let polarity =

                self.endpoint_exts.get_occupied(port.0).ok_or(UnknownPort(port))?.info.polarity;

            if polarity != expected_polarity {

                return Err(WrongPortPolarity { param_index, port });

            }

        }

        let state = protocol.new_main_component(identifier, &moved_port_list);

        let component = Component {

            port_set: moved_port_set,

            protocol: protocol.clone(),

            identifier: identifier.clone(),

            state,

        };

        //////////////////////////////

        // success! mutate self and return Ok

        self.native_ports.retain(|e| !component.port_set.contains(e));

        self.components.new_occupied(component);

        Ok(())

    }

    pub fn new_channel(&mut self) -> (OutPort, InPort) {

        assert!(self.endpoint_exts.len() <= std::u32::MAX as usize - 2);

        let channel_id = ChannelId {

            controller_id: self.controller_id,

            channel_index: self.channel_index_stream.next(),

        };

        let [e0, e1] = Endpoint::new_memory_pair();

        let kp = self.endpoint_exts.new_occupied(EndpointExt {

            info: EndpointInfo { channel_id, polarity: Putter },

            endpoint: e0,

        });

        let kg = self.endpoint_exts.new_occupied(EndpointExt {

            info: EndpointInfo { channel_id, polarity: Getter },

            endpoint: e1,

        });

        (OutPort(Port(kp)), InPort(Port(kg)))

    }

    pub fn sync_set(&mut self, _inbuf: &mut [u8], _ops: &mut [PortOpRs]) -> Result<(), ()> {

        Ok(())

    }

    pub fn sync_subsets(

        &mut self,

        _inbuf: &mut [u8],

        _ops: &mut [PortOpRs],

        bit_subsets: &[&[usize]],

    ) -> Result<usize, ()> {

        for (batch_index, bit_subset) in bit_subsets.iter().enumerate() {

            println!("batch_index {:?}", batch_index);

            let chunk_iter = bit_subset.iter().copied();

            for index in BitChunkIter::new(chunk_iter) {

                println!("  index {:?}", index);

            }

        }

        Ok(0)

    }

}

macro_rules! bitslice {

    ($( $num:expr  ),*) => {{

        &[0 $( | (1usize << $num)  )*]

    }};

}

#[test]

fn api_new_test() {

    let mut c = Connecting::default();

    let net_out: OutPort = c.bind(Coupling::Active, "127.0.0.1:8000".parse().unwrap());

    let net_in: InPort = c.bind(Coupling::Active, "127.0.0.1:8001".parse().unwrap());

    let proto_0 = Arc::new(ProtocolD::parse(b"").unwrap());

    let mut c = c.connect(None).unwrap();

    let (mem_out, mem_in) = c.new_channel();

    let mut inbuf = [0u8; 64];

    let identifier: Arc<[u8]> = b"sync".to_vec().into();

    c.new_component(&proto_0, &identifier, &[net_in.into(), mem_out.into()]).unwrap();

    let mut ops = [

        PortOpRs::In { msg_range: None, port: &mem_in },

        PortOpRs::Out { msg: b"hey", port: &net_out, optional: false },

        PortOpRs::Out { msg: b"hi?", port: &net_out, optional: true },

        PortOpRs::Out { msg: b"yo!", port: &net_out, optional: false },

    ];

    c.sync_set(&mut inbuf, &mut ops).unwrap();

    c.sync_subsets(&mut inbuf, &mut ops, &[bitslice! {0,1,2}]).unwrap();

}

#[repr(C)]

pub struct PortOp {

    msgptr: *mut u8, // read if OUT, field written if IN, will point into buf

    msglen: usize,   // read if OUT, written if IN, won't exceed buf

    port: Port,

    optional: bool, // no meaning if

}

pub enum PortOpRs<'a> {

    In { msg_range: Option<Range<usize>>, port: &'a InPort },

    Out { msg: &'a [u8], port: &'a OutPort, optional: bool },

}

unsafe fn c_sync_set(

    connected: &mut Connected,

    inbuflen: usize,

    inbufptr: *mut u8,

    opslen: usize,

    opsptr: *mut PortOp,

) -> i32 {

    let buf = as_mut_slice(inbuflen, inbufptr);

    let ops = as_mut_slice(opslen, opsptr);

    let (subset_index, wrote) = sync_inner(connected, buf);

    assert_eq!(0, subset_index);

    for op in ops {

        if let Some(range) = wrote.get(&op.port) {

            op.msgptr = inbufptr.add(range.start);

            op.msglen = range.end - range.start;

        }

    }

    0

}

use super::bits::{usizes_for_bits, BitChunkIter};

unsafe fn c_sync_subset(

    connected: &mut Connected,

    inbuflen: usize,

    inbufptr: *mut u8,

    opslen: usize,

    opsptr: *mut PortOp,

    subsetslen: usize,

    subsetsptr: *const *const usize,

) -> i32 {

    let buf: &mut [u8] = as_mut_slice(inbuflen, inbufptr);

    let ops: &mut [PortOp] = as_mut_slice(opslen, opsptr);

    let subsets: &[*const usize] = as_const_slice(subsetslen, subsetsptr);

    let subsetlen = usizes_for_bits(opslen);

    // don't yet know subsetptr; which subset fires unknown!

    let (subset_index, wrote) = sync_inner(connected, buf);

    let subsetptr: *const usize = subsets[subset_index];

    let subset: &[usize] = as_const_slice(subsetlen, subsetptr);

    for index in BitChunkIter::new(subset.iter().copied()) {

        let op = &mut ops[index as usize];

        if let Some(range) = wrote.get(&op.port) {

            op.msgptr = inbufptr.add(range.start);

            op.msglen = range.end - range.start;

        }

    }

    subset_index as i32

}

// dummy fn for the actual synchronous round

fn sync_inner<'c, 'b>(

    _connected: &'c mut Connected,

    _buf: &'b mut [u8],

) -> (usize, &'b HashMap<Port, Range<usize>>) {

    todo!()

}

unsafe fn as_mut_slice<'a, T>(len: usize, ptr: *mut T) -> &'a mut [T] {

    std::slice::from_raw_parts_mut(ptr, len)

}

unsafe fn as_const_slice<'a, T>(len: usize, ptr: *const T) -> &'a [T] {

    std::slice::from_raw_parts(ptr, len)

}

#[test]

fn api_connecting() {

    let addrs: [SocketAddr; 3] = [

        "127.0.0.1:8888".parse().unwrap(),

        "127.0.0.1:8889".parse().unwrap(),

        "127.0.0.1:8890".parse().unwrap(),

    ];

    let handles = vec![

        std::thread::spawn(move || {

            let mut connecting = Connecting::default();

            println!("A: {:#?}", connected);

        }),

        std::thread::spawn(move || {

            let mut connecting = Connecting::default();

            println!("B: {:#?}", connected);

        }),

    ];

    for h in handles {

        h.join().unwrap();

    }

}