use std::collections::HashMap;
use std::convert::Infallible;
use std::io;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use avril_caps::{Capabilities, PathTarget};
use avril_core::serde_v8::{JsBuffer, ToJsBuffer};
use avril_core::{Extension, Op};
use avril_net::{Address, Endpoint, Socket};
use serde::Serialize;
use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf};
use tokio::sync::{Mutex as AsyncMutex, oneshot};
use tokio_util::sync::CancellationToken;

use avril_host::io::response;

mod udp;

const CHUNK: usize = 64 * 1024;
const MAX_ID: u64 = 9_007_199_254_740_991;
type ConnectionInfo = (u64, Endpoint, Endpoint);
type ListenerInfo = (u64, Endpoint);

fn closed() -> io::Error {
    io::Error::new(io::ErrorKind::NotConnected, "Socket is closed")
}

#[derive(Clone)]
enum Authority {
    Net {
        caps: Arc<Capabilities>,
        hostname: String,
        port: u16,
    },
    Unix {
        caps: Arc<Capabilities>,
        target: PathTarget,
    },
}
impl Authority {
    fn check(&self) -> io::Result<()> {
        let result = match self {
            Self::Net { caps, hostname, port } => caps.check_net(hostname, *port),
            Self::Unix { caps, target } => caps.check_write(target),
        };
        result.map_err(|error| io::Error::new(io::ErrorKind::PermissionDenied, error))
    }
}
fn authority(
    mut address: Address,
    caps: &Arc<Capabilities>,
    connecting: bool,
) -> io::Result<(Address, Authority)> {
    let authority = match &mut address {
        Address::Tcp { hostname, port } => Authority::Net {
            caps: caps.clone(),
            hostname: hostname.clone(),
            port: *port,
        },
        Address::Unix { path } => {
            if !cfg!(unix) {
                return Err(io::Error::new(
                    io::ErrorKind::Unsupported,
                    "Unix sockets are not supported on this platform",
                ));
            }
            let label = path
                .as_ref()
                .filter(|path| !path.is_empty())
                .ok_or_else(|| {
                    io::Error::new(io::ErrorKind::InvalidInput, "Missing Unix socket path")
                })?;
            let target = if connecting {
                PathTarget::resolve(caps.cwd().join(label))?
            } else {
                PathTarget::for_entry(caps.cwd().join(label))?
            };
            *path = Some(
                target
                    .as_path()
                    .to_str()
                    .ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "Unix socket paths must be UTF-8",
                        )
                    })?
                    .to_owned(),
            );
            Authority::Unix { caps: caps.clone(), target }
        }
    };
    authority.check()?;
    Ok((address, authority))
}
fn hostname(value: &str) -> &str {
    value
        .strip_prefix('[')
        .and_then(|value| value.strip_suffix(']'))
        .unwrap_or(value)
}

struct Connection {
    reader: AsyncMutex<ReadHalf<Socket>>,
    writer: AsyncMutex<WriteHalf<Socket>>,
    write_closed: AtomicBool,
    canceled: CancellationToken,
    authority: Authority,
    local: Endpoint,
    remote: Endpoint,
}
impl Connection {
    fn new(
        socket: Socket,
        local: Endpoint,
        remote: Endpoint,
        authority: Authority,
        canceled: CancellationToken,
    ) -> Self {
        let (reader, writer) = tokio::io::split(socket);
        Self {
            reader: AsyncMutex::new(reader),
            writer: AsyncMutex::new(writer),
            write_closed: AtomicBool::new(false),
            canceled,
            authority,
            local,
            remote,
        }
    }
}
struct Listener {
    socket: avril_net::Listener,
    authority: Authority,
    canceled: CancellationToken,
}
struct Pending {
    receiver: AsyncMutex<Option<oneshot::Receiver<io::Result<Connection>>>>,
    canceled: CancellationToken,
}
enum Resource {
    Connection(Arc<Connection>),
    Listener(Arc<Listener>),
    Pending(Arc<Pending>),
    Datagram(Arc<udp::Datagram>),
}
impl Resource {
    fn cancel(&self) {
        match self {
            Self::Connection(value) => value.canceled.cancel(),
            Self::Listener(value) => value.canceled.cancel(),
            Self::Pending(value) => value.canceled.cancel(),
            Self::Datagram(value) => value.canceled.cancel(),
        }
    }
}
// Spawned I/O must not keep the registry alive: its drop cancels worker resources.
struct Sockets {
    runtime: tokio::runtime::Handle,
    caps: Arc<Capabilities>,
    resources: Mutex<HashMap<u64, Resource>>,
    next: AtomicU64,
    canceled: CancellationToken,
}
impl Drop for Sockets {
    fn drop(&mut self) {
        self.canceled.cancel();
    }
}
impl Sockets {
    fn id(&self) -> io::Result<u64> {
        self.next
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| {
                (id <= MAX_ID).then_some(id + 1)
            })
            .map_err(|_| io::Error::other("Socket IDs exhausted"))
    }
    fn connection(&self, id: u64) -> io::Result<Arc<Connection>> {
        match self.resources.lock().unwrap().get(&id) {
            Some(Resource::Connection(value)) => Ok(value.clone()),
            _ => Err(closed()),
        }
    }
    fn insert(&self, connection: Connection) -> io::Result<ConnectionInfo> {
        let id = self.id()?;
        let row = (id, connection.local.clone(), connection.remote.clone());
        self.resources
            .lock()
            .unwrap()
            .insert(id, Resource::Connection(Arc::new(connection)));
        Ok(row)
    }
    fn listen(&self, target: Address) -> io::Result<ListenerInfo> {
        let (target, mut authority) = authority(target, &self.caps, false)?;
        let (socket, local) =
            avril_net::Listener::bind(&self.runtime, &target).map_err(io::Error::other)?;
        if let (Authority::Net { port, .. }, Endpoint::Tcp(selected)) = (&mut authority, &local) {
            *port = selected.port();
        }
        authority.check()?;
        let id = self.id()?;
        self.resources.lock().unwrap().insert(
            id,
            Resource::Listener(Arc::new(Listener {
                socket,
                authority,
                canceled: self.canceled.child_token(),
            })),
        );
        Ok((id, local))
    }
    fn start(&self, target: Address) -> io::Result<u64> {
        let (target, authority) = authority(target, &self.caps, true)?;
        let id = self.id()?;
        let canceled = self.canceled.child_token();
        let (sender, receiver) = oneshot::channel();
        self.resources.lock().unwrap().insert(
            id,
            Resource::Pending(Arc::new(Pending {
                receiver: AsyncMutex::new(Some(receiver)),
                canceled: canceled.clone(),
            })),
        );
        self.runtime.spawn(async move {
            let result = async {
                let (socket,local,remote): (Socket,Endpoint,Endpoint) = match &target {
                    Address::Tcp { hostname:host,port } => {
                        let socket=tokio::net::TcpStream::connect((hostname(host),*port)).await?;
                        let local=Endpoint::Tcp(socket.local_addr()?);
                        let remote=Endpoint::Tcp(socket.peer_addr()?);
                        (socket.into(),local,remote)
                    }
                    #[cfg(unix)]
                    Address::Unix { path:Some(path) } => {
                        let socket=tokio::net::UnixStream::connect(path).await?;
                        let local=Endpoint::Unix {path:socket.local_addr()?.as_pathname().map(|path|path.to_string_lossy().into_owned())};
                        (socket.into(),local,Endpoint::Unix {path:Some(path.clone())})
                    }
                    _ => return Err(io::Error::new(io::ErrorKind::Unsupported,"Unsupported socket transport")),
                };
                authority.check()?;
                Ok(Connection::new(socket,local,remote,authority,canceled.clone()))
            };
            let value = tokio::select! { biased; _=canceled.cancelled()=>Err(closed()), value=result=>value };
            let _ = sender.send(value);
        });
        Ok(id)
    }
    fn close(&self, id: u64) -> io::Result<()> {
        let value = self
            .resources
            .lock()
            .unwrap()
            .remove(&id)
            .ok_or_else(closed)?;
        value.cancel();
        Ok(())
    }
}

async fn reply<T: Serialize>(
    task: tokio::task::JoinHandle<io::Result<T>>,
) -> Result<impl Serialize, Infallible> {
    Ok(response(task.await.unwrap_or_else(|error| {
        Err(io::Error::other(error.to_string()))
    })))
}

pub fn init(runtime: tokio::runtime::Handle, caps: Arc<Capabilities>) -> Extension {
    let sockets = Arc::new(Sockets {
        runtime: runtime.clone(),
        caps,
        resources: Mutex::new(HashMap::new()),
        next: AtomicU64::new(0),
        canceled: CancellationToken::new(),
    });
    let mut ops = vec![
        Op::new("op_net_listen", {
            let sockets = sockets.clone();
            move |target: Address| Ok::<_, Infallible>(response(sockets.listen(target)))
        }),
        Op::new("op_net_connect_start", {
            let sockets = sockets.clone();
            move |target: Address| Ok::<_, Infallible>(response(sockets.start(target)))
        }),
        Op::asynchronous("op_net_connected", {
            let sockets = sockets.clone();
            move |id: u64| {
                let pending = match sockets.resources.lock().unwrap().get(&id) {
                    Some(Resource::Pending(value)) => Some(value.clone()),
                    _ => None,
                };
                let weak = Arc::downgrade(&sockets);
                reply(runtime.spawn(async move {
                    let pending = pending.ok_or_else(closed)?;
                    let receiver = pending.receiver.lock().await.take().ok_or_else(closed)?;
                    let connection = tokio::select! {
                        biased;
                        _ = pending.canceled.cancelled() => return Err(closed()),
                        value = receiver => value.map_err(|_|closed())??,
                    };
                    connection.authority.check()?;
                    let sockets = weak.upgrade().ok_or_else(closed)?;
                    let row = (id, connection.local.clone(), connection.remote.clone());
                    let mut resources = sockets.resources.lock().unwrap();
                    if pending.canceled.is_cancelled() || !resources.contains_key(&id) {
                        return Err(closed());
                    }
                    resources.insert(id, Resource::Connection(Arc::new(connection)));
                    Ok(row)
                }))
            }
        }),
        Op::asynchronous("op_net_accept", {
            let sockets = sockets.clone();
            let runtime = sockets.runtime.clone();
            move |id: u64| {
                let listener = match sockets.resources.lock().unwrap().get(&id) {
                    Some(Resource::Listener(value)) => Some(value.clone()),
                    _ => None,
                };
                let weak = Arc::downgrade(&sockets);
                let canceled = sockets.canceled.child_token();
                reply(runtime.spawn(async move {
                    let listener = listener.ok_or_else(closed)?;
                    listener.authority.check()?;
                    let accepted = tokio::select! {
                        biased;
                        _ = listener.canceled.cancelled() => return Err(closed()),
                        value = listener.socket.accept() => value?,
                    };
                    listener.authority.check()?;
                    if listener.canceled.is_cancelled() {
                        return Err(closed());
                    }
                    let connection = Connection::new(
                        accepted.socket,
                        accepted.local,
                        accepted.remote,
                        listener.authority.clone(),
                        canceled,
                    );
                    weak.upgrade().ok_or_else(closed)?.insert(connection)
                }))
            }
        }),
        Op::asynchronous("op_net_read", {
            let sockets = sockets.clone();
            let runtime = sockets.runtime.clone();
            move |(id, count): (u64, usize)| {
                let connection = sockets.connection(id);
                reply(runtime.spawn(async move {
                    let connection = connection?;
                    connection.authority.check()?;
                    let work = async {
                        let mut reader = connection.reader.lock().await;
                        connection.authority.check()?;
                        let mut bytes = vec![0; count.min(CHUNK)];
                        if bytes.is_empty() {
                            return Ok(Some(ToJsBuffer::from(bytes)));
                        }
                        let count = reader.read(&mut bytes).await?;
                        connection.authority.check()?;
                        if count == 0 {
                            return Ok(None);
                        }
                        bytes.truncate(count);
                        Ok(Some(ToJsBuffer::from(bytes)))
                    };
                    tokio::select! {
                        biased;
                        _ = connection.canceled.cancelled() => Err(closed()),
                        value = work => value,
                    }
                }))
            }
        }),
        Op::asynchronous("op_net_write", {
            let sockets = sockets.clone();
            let runtime = sockets.runtime.clone();
            move |(id, bytes): (u64, JsBuffer)| {
                let connection = sockets.connection(id);
                let bytes = bytes[..bytes.len().min(CHUNK)].to_vec();
                reply(runtime.spawn(async move {
                    let connection = connection?;
                    let work = async {
                        let mut writer = connection.writer.lock().await;
                        connection.authority.check()?;
                        if connection.write_closed.load(Ordering::Acquire) {
                            return Err(closed());
                        }
                        writer.write(&bytes).await
                    };
                    tokio::select! {
                        biased;
                        _ = connection.canceled.cancelled() => Err(closed()),
                        value = work => value,
                    }
                }))
            }
        }),
        Op::asynchronous("op_net_writev", {
            let sockets = sockets.clone();
            let runtime = sockets.runtime.clone();
            move |(id, buffers): (u64, Vec<JsBuffer>)| {
                let connection = sockets.connection(id);
                let segments = avril_host::io::bound_vectored(buffers);
                reply(runtime.spawn(async move {
                    let connection = connection?;
                    let work = async {
                        let mut writer = connection.writer.lock().await;
                        connection.authority.check()?;
                        if connection.write_closed.load(Ordering::Acquire) {
                            return Err(closed());
                        }
                        if segments.is_empty() {
                            return Ok(0);
                        }
                        let slices = segments
                            .iter()
                            .map(|segment| io::IoSlice::new(segment))
                            .collect::<Vec<_>>();
                        writer.write_vectored(&slices).await
                    };
                    tokio::select! {
                        biased;
                        _ = connection.canceled.cancelled() => Err(closed()),
                        value = work => value,
                    }
                }))
            }
        }),
        Op::asynchronous("op_net_shutdown", {
            let sockets = sockets.clone();
            let runtime = sockets.runtime.clone();
            move |id: u64| {
                let connection = sockets.connection(id);
                reply(runtime.spawn(async move {
                    let connection = connection?;
                    let work = async {
                        let mut writer = connection.writer.lock().await;
                        connection.authority.check()?;
                        if !connection.write_closed.swap(true, Ordering::AcqRel) {
                            writer.shutdown().await?;
                        }
                        Ok(())
                    };
                    tokio::select! {
                        biased;
                        _ = connection.canceled.cancelled() => Err(closed()),
                        value = work => value,
                    }
                }))
            }
        }),
        Op::new("op_net_close", {
            let sockets = sockets.clone();
            move |id: u64| Ok::<_, Infallible>(response(sockets.close(id)))
        }),
    ];
    ops.extend(udp::ops(sockets));
    Extension {
        name: "ext:net",
        dependencies: &["ext:file", "ext:web", "ext:webstreams"],
        ops,
        modules: &[(
            "ext:net/main.js",
            include_str!(concat!(env!("OUT_DIR"), "/main.js")),
        )],
        entry_point: None,
    }
}