//! TCP endpoint publication, verified connections, and per-request notifications.
use crate::*;
type Notify = Box<dyn FnMut(&str, Value)>;
use serde_json::Value;
#[cfg(target_os = "wasi")]
use serde_json::json;
use std::{future::IntoFuture, marker::PhantomData, path::Path, rc::Rc};

/// A verified endpoint address and identity.
///
/// Connect verifies and closes one connection. Each call opens a new connection
/// and repeats the identity handshake before its request. The first call costs three round trips.
#[derive(Clone)]
#[cfg_attr(not(target_os = "wasi"), allow(dead_code))]
pub struct Endpoint {
    address: String,
    owner: String,
    interface: String,
}
impl Endpoint {
    /// Require the interface expected by a generated client.
    pub fn require_interface(&self, name: &str) -> Result<()> {
        if self.interface == name {
            Ok(())
        } else {
            Err(Error::new(-32000, "endpoint interface mismatch"))
        }
    }
    /// Build a typed endpoint request.
    pub fn call<T>(&self, method: &str, params: impl serde::Serialize) -> EndpointCall<T> {
        EndpointCall {
            endpoint: self.clone(),
            method: method.into(),
            params: serde_json::to_value(params).expect("serializable params"),
            notify: None,
            result: PhantomData,
        }
    }
}
/// A lazy endpoint request with optional progress callbacks.
#[must_use = "await the endpoint call to send it"]
pub struct EndpointCall<T> {
    endpoint: Endpoint,
    method: String,
    params: Value,
    notify: Option<Notify>,
    result: PhantomData<T>,
}
impl<T> EndpointCall<T> {
    /// Receive progress from this endpoint handler.
    pub fn on_notify(mut self, f: impl FnMut(&str, Value) + 'static) -> Self {
        self.notify = Some(Box::new(f));
        self
    }
}
impl<T: serde::de::DeserializeOwned + 'static> IntoFuture for EndpointCall<T> {
    type Output = Result<T>;
    type IntoFuture = futures::future::LocalBoxFuture<'static, Result<T>>;
    fn into_future(self) -> Self::IntoFuture {
        Box::pin(async move {
            Ok(serde_json::from_value(
                endpoint_request(self.endpoint, &self.method, self.params, self.notify).await?,
            )?)
        })
    }
}

/// Publish a declared interface and keep its listener alive until stdin EOF or [`drain`].
#[cfg(target_os = "wasi")]
pub async fn listen<T: Interface>(path: impl AsRef<Path>, value: Rc<T>) -> Result<()> {
    use wstd::iter::AsyncIterator;
    let name = crate::transport::endpoint_name::<T>()?;
    let listener = wstd::net::TcpListener::bind("127.0.0.1:0").await?;
    write_atomic(
        path,
        format!("tcp:{} {} {}\n", listener.local_addr()?, pid(), name),
    )?;
    crate::transport::endpoint_task(async move {
        let mut incoming = listener.incoming();
        while let Some(Ok(stream)) = incoming.next().await {
            let value = value.clone();
            let name = name.clone();
            background(async move {
                let _ = serve_connection(stream, value, name).await;
            });
        }
    });
    Ok(())
}
/// Publish an endpoint. Native programs require the WASI runtime.
#[cfg(not(target_os = "wasi"))]
pub async fn listen<T: Interface>(_path: impl AsRef<Path>, _value: Rc<T>) -> Result<()> {
    crate::transport::endpoint_name::<T>()?;
    Err(Error::new(-32000, "endpoints require WASI"))
}

/// Read an endpoint file and verify its owner before returning a client.
#[cfg(target_os = "wasi")]
pub async fn connect(path: impl AsRef<Path>) -> Result<Endpoint> {
    let mut attempts = 0;
    let contents = loop {
        match std::fs::read_to_string(path.as_ref()) {
            Ok(s) => break s,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound && attempts < 20 => {
                attempts += 1;
                wstd::task::sleep(wstd::time::Duration::from_millis(25)).await;
            }
            Err(e) => return Err(e.into()),
        }
    };
    let words: Vec<_> = contents.split_whitespace().collect();
    if words.len() != 3 {
        return Err(Error::new(-32602, "invalid endpoint file"));
    }
    let endpoint = Endpoint {
        address: words[0]
            .strip_prefix("tcp:")
            .ok_or_else(|| Error::new(-32602, "invalid endpoint address"))?
            .into(),
        owner: words[1].into(),
        interface: words[2].into(),
    };
    let _ = handshake(&endpoint).await?;
    Ok(endpoint)
}
/// Connect to an endpoint. Native programs require the WASI runtime.
#[cfg(not(target_os = "wasi"))]
pub async fn connect(_path: impl AsRef<Path>) -> Result<Endpoint> {
    Err(Error::new(-32000, "endpoints require WASI"))
}

#[cfg(target_os = "wasi")]
async fn send(stream: &mut impl wstd::io::AsyncWrite, value: Value) -> Result<()> {
    stream.write_all(format!("{value}\n").as_bytes()).await?;
    stream.flush().await?;
    Ok(())
}
#[cfg(target_os = "wasi")]
async fn handshake(endpoint: &Endpoint) -> Result<wstd::net::TcpStream> {
    let mut stream = wstd::net::TcpStream::connect(endpoint.address.as_str()).await?;
    send(
        &mut stream,
        json!({"jsonrpc":"2.0","id":0,"method":"rpc.endpoint","params":{}}),
    )
    .await?;
    let reply = crate::lines::Lines::default()
        .rpc(&mut stream)
        .await?
        .ok_or_else(|| Error::new(-32000, "endpoint closed"))?;
    let result = &reply["result"];
    let plugin = endpoint
        .owner
        .len()
        .checked_sub(37)
        .and_then(|end| endpoint.owner.get(..end));
    if result["owner"] != endpoint.owner
        || result["interface"] != endpoint.interface
        || result["plugin"].as_str() != plugin
    {
        return Err(Error::new(-32000, "endpoint identity mismatch"));
    }
    Ok(stream)
}
#[cfg(target_os = "wasi")]
async fn endpoint_request(
    endpoint: Endpoint,
    method: &str,
    params: Value,
    mut notify: Option<Notify>,
) -> Result<Value> {
    let mut stream = handshake(&endpoint).await?;
    let mut lines = crate::lines::Lines::default();
    send(
        &mut stream,
        json!({"jsonrpc":"2.0","id":1,"method":method,"params":params}),
    )
    .await?;
    while let Some(value) = lines.rpc(&mut stream).await? {
        if value["method"] == "core.notify" {
            if let Some(f) = &mut notify {
                f(
                    value["params"]["name"].as_str().unwrap_or_default(),
                    value["params"]["params"].clone(),
                );
            }
        } else if value["id"] == 1 {
            if let Some(error) = value.get("error") {
                return Err(serde_json::from_value(error.clone())?);
            }
            return Ok(value["result"].clone());
        }
    }
    Err(Error::new(-32000, "endpoint closed before response"))
}
#[cfg(not(target_os = "wasi"))]
async fn endpoint_request(
    _endpoint: Endpoint,
    _method: &str,
    _params: Value,
    _notify: Option<Notify>,
) -> Result<Value> {
    Err(Error::new(-32000, "endpoints require WASI"))
}

#[cfg(target_os = "wasi")]
async fn serve_connection<T: Interface>(
    stream: wstd::net::TcpStream,
    value: Rc<T>,
    name: String,
) -> Result<()> {
    use futures::{FutureExt, StreamExt};
    let (mut reader, mut writer) = stream.split();
    let mut lines = crate::lines::Lines::default();
    while let Some(request) = crate::transport::until_drained(lines.rpc(&mut reader))
        .await
        .transpose()?
        .flatten()
    {
        let method = request["method"].as_str().unwrap_or_default();
        let (tx, mut rx) = futures::channel::mpsc::unbounded();
        let ctx = Ctx {
            sink: Some(tx),
            ..Ctx::default()
        };
        let handler=async {
            if method=="rpc.endpoint" {Ok(json!({"owner":pid(),"plugin":std::env::var("SAUSAGE_PLUGIN").unwrap_or_default(),"interface":name}))}
            else {value.dispatch(method,request.get("params").cloned().unwrap_or(json!({})),&ctx).await}
        }.fuse();
        futures::pin_mut!(handler);
        let result = loop {
            futures::select! {result=handler=>break result,notification=rx.next().fuse()=>{if let Some(notification)=notification {send(&mut writer,notification).await?;}}}
        };
        while let Ok(notification) = rx.try_recv() {
            send(&mut writer, notification).await?;
        }
        if let Some(id) = request.get("id") {
            let response = match result {
                Ok(result) => json!({"jsonrpc":"2.0","id":id,"result":result}),
                Err(error) => json!({"jsonrpc":"2.0","id":id,"error":error}),
            };
            send(&mut writer, response).await?;
        }
        if ctx.exit.get() {
            crate::transport::exit();
            break;
        }
    }
    Ok(())
}