use crate::{BACKEND_VERSION, Result, wire};
use serde::de::DeserializeOwned;
use serde_json::{Value, json};
use std::{
    collections::HashMap,
    path::Path,
    process::Stdio,
    sync::{Arc, Mutex},
    time::Duration,
};
use tokio::{
    io::BufReader,
    process::{Child, Command},
    sync::{mpsc, oneshot},
    task::{JoinHandle, JoinSet},
};

pub trait Callbacks: Send + Sync + 'static {
    fn call(&self, method: &str, params: Value) -> Result<Value>;
}

type Reply = Result<Value>;
#[derive(Default)]
struct State {
    next: u64,
    failed: Option<String>,
    pending: HashMap<u64, oneshot::Sender<Reply>>,
}
impl State {
    fn fail(&mut self, error: String) {
        self.failed.get_or_insert(error.clone());
        for (_, reply) in self.pending.drain() {
            let _ = reply.send(Err(error.clone().into()));
        }
    }
}

const WRITER_CAPACITY: usize = 16;

type Writer = mpsc::Sender<WriteRequest>;
struct WriteRequest {
    value: Value,
    result: oneshot::Sender<Result<()>>,
}

async fn send(writer: &Writer, value: Value) -> Result<()> {
    let (result, rx) = oneshot::channel();
    writer.send(WriteRequest { value, result }).await?;
    rx.await?
}

async fn writer_loop<W: tokio::io::AsyncWrite + Unpin>(
    mut output: W,
    mut requests: mpsc::Receiver<WriteRequest>,
    state: Arc<Mutex<State>>,
) {
    while let Some(request) = requests.recv().await {
        let result = wire::write(&mut output, request.value).await;
        let failed = result.as_ref().err().map(ToString::to_string);
        let _ = request.result.send(result);
        if let Some(error) = failed {
            state.lock().unwrap().fail(error);
            break;
        }
    }
}

pub struct Backend {
    child: Child,
    state: Arc<Mutex<State>>,
    writer: Option<Writer>,
    writer_task: Option<JoinHandle<()>>,
    reader: Option<JoinHandle<()>>,
}

impl Backend {
    pub async fn start(path: &Path, host: Arc<dyn Callbacks>) -> Result<Self> {
        let version = tokio::time::timeout(
            Duration::from_secs(10),
            Command::new(path)
                .env_clear()
                .kill_on_drop(true)
                .arg("--version")
                .output(),
        )
        .await??;
        if !version.status.success()
            || String::from_utf8_lossy(&version.stdout).trim()
                != format!("Version {BACKEND_VERSION}")
        {
            return Err(format!("Expected TypeScript {BACKEND_VERSION} native executable").into());
        }
        let mut child = Command::new(path)
            .env_clear()
            .kill_on_drop(true)
            .args([
                "--api",
                "--async",
                "--cwd",
                "/avril-lsp",
                "--callbacks",
                concat!(
                    "readFile,fileExists,directoryExists,getAccessibleEntries,",
                    "realpath,stat:fakeStat,writeFile:error,removeFile:error"
                ),
            ])
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .spawn()?;
        let stdin = child.stdin.take().ok_or("No backend stdin")?;
        let (writer, requests) = mpsc::channel::<WriteRequest>(WRITER_CAPACITY);
        let mut stdout = BufReader::new(child.stdout.take().ok_or("No backend stdout")?);
        let state = Arc::new(Mutex::new(State::default()));
        let writer_task = tokio::spawn(writer_loop(stdin, requests, state.clone()));
        let reader = tokio::spawn({
            let state = state.clone();
            let writer = writer.clone();
            async move {
                let mut callbacks = JoinSet::new();
                let result: Result<()> = async {
                    loop {
                        tokio::select! {
                            completed = callbacks.join_next(), if !callbacks.is_empty() => { completed.ok_or("Missing callback task")???; }
                            message = wire::read(&mut stdout) => {
                                let message = message?.ok_or("TypeScript backend closed stdout")?;
                                if let Some(method) = message["method"].as_str() {
                                    if let Some(id) = message.get("id").cloned() {
                                        let method = method.to_owned();
                                        let host = host.clone();
                                        let writer = writer.clone();
                                        callbacks.spawn(async move {
                                            let response = match host.call(&method, message["params"].clone()) {
                                                Ok(result) => json!({"id":id,"result":result}),
                                                Err(error) => json!({"id":id,"error":{"code":-32603,"message":error.to_string()}}),
                                            };
                                            send(&writer, response).await
                                        });
                                    }
                                } else if let Some(id) = message["id"].as_u64() {
                                    if let Some(reply) = state.lock().unwrap().pending.remove(&id) {
                                        let result = if let Some(error) = message.get("error") {
                                            Err(serde_json::from_value::<wire::RpcError>(error.clone())?.into())
                                        } else {
                                            Ok(message["result"].clone())
                                        };
                                        let _ = reply.send(result);
                                    }
                                } else { return Err("Invalid backend response ID".into()); }
                            }
                        }
                    }
                }.await;
                callbacks.abort_all();
                state.lock().unwrap().fail(
                    result
                        .err()
                        .map_or("Backend reader stopped".into(), |e| e.to_string()),
                );
            }
        });
        Ok(Self {
            child,
            state,
            writer: Some(writer),
            writer_task: Some(writer_task),
            reader: Some(reader),
        })
    }

    pub fn failed(&self) -> bool {
        self.state.lock().unwrap().failed.is_some()
    }

    pub async fn call<T: DeserializeOwned>(&self, method: &str, params: Value) -> Result<T> {
        let (tx, rx) = oneshot::channel();
        let id = {
            let mut state = self.state.lock().unwrap();
            if let Some(error) = &state.failed {
                return Err(error.clone().into());
            }
            state.next += 1;
            let id = state.next;
            state.pending.insert(id, tx);
            id
        };
        // Remove registrations even if the editor cancels this future.
        struct Registration(Arc<Mutex<State>>, u64, Writer);
        impl Drop for Registration {
            fn drop(&mut self) {
                if self.0.lock().unwrap().pending.remove(&self.1).is_some() {
                    let (result, _) = oneshot::channel();
                    let _ = self.2.try_send(WriteRequest {
                        value: json!({"method":"$/cancelRequest","params":{"id":self.1}}),
                        result,
                    });
                }
            }
        }
        let writer = self.writer.as_ref().ok_or("Backend stdin closed")?.clone();
        let _registration = Registration(self.state.clone(), id, writer.clone());
        let operation = async {
            send(&writer, json!({"id":id,"method":method,"params":params})).await?;
            let value = rx.await??;
            Ok::<T, crate::Error>(serde_json::from_value(value)?)
        };
        match tokio::time::timeout(Duration::from_secs(30), operation).await {
            Ok(result) => result,
            Err(_) => {
                self.state
                    .lock()
                    .unwrap()
                    .fail(format!("TypeScript API timed out: {method}"));
                Err(format!("TypeScript API timed out: {method}").into())
            }
        }
    }

    pub async fn close(&mut self) -> Result<()> {
        self.state
            .lock()
            .unwrap()
            .fail("Backend shutting down".into());
        if let Some(reader) = self.reader.take() {
            reader.abort();
            let _ = reader.await;
        }
        self.writer.take();
        if let Some(mut writer_task) = self.writer_task.take() {
            match tokio::time::timeout(Duration::from_secs(2), &mut writer_task).await {
                Ok(Ok(())) => {}
                Ok(Err(error)) => return Err(error.into()),
                Err(_) => {
                    writer_task.abort();
                    let _ = writer_task.await;
                    self.child.kill().await?;
                    self.child.wait().await?;
                    return Err("TypeScript backend writer did not shut down".into());
                }
            }
        }
        match tokio::time::timeout(Duration::from_secs(2), self.child.wait()).await {
            Ok(status) => {
                if status?.success() {
                    Ok(())
                } else {
                    Err("TypeScript backend exited unsuccessfully".into())
                }
            }
            Err(_) => {
                self.child.kill().await?;
                self.child.wait().await?;
                Err("TypeScript backend did not shut down".into())
            }
        }
    }
}
impl Drop for Backend {
    fn drop(&mut self) {
        if let Some(reader) = &self.reader {
            reader.abort();
        }
        if let Some(writer) = &self.writer_task {
            writer.abort();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::{AsyncReadExt, BufReader};

    #[tokio::test]
    async fn canceled_frame_does_not_corrupt_next_frame() {
        let state = Arc::new(Mutex::new(State::default()));
        let (writer, requests) = mpsc::channel(WRITER_CAPACITY);
        let (backend, mut output) = tokio::io::duplex(8);
        let task = tokio::spawn(writer_loop(backend, requests, state));

        let large = json!({"id":1,"method":"large","params":{"text":"x".repeat(1024 * 1024)}});
        let first_call = tokio::spawn({
            let writer = writer.clone();
            async move { send(&writer, large).await }
        });

        let mut prefix = [0; 8];
        output.read_exact(&mut prefix).await.unwrap();
        first_call.abort();
        let _ = first_call.await;

        let (second_result, second_rx) = oneshot::channel();
        writer
            .send(WriteRequest {
                value: json!({"id":2,"method":"query","params":{}}),
                result: second_result,
            })
            .await
            .unwrap();
        drop(writer);
        let mut bytes = prefix.to_vec();
        let (second_result, read_result) = tokio::join!(second_rx, output.read_to_end(&mut bytes));
        second_result.unwrap().unwrap();
        read_result.unwrap();
        task.await.unwrap();
        let mut input = BufReader::new(bytes.as_slice());
        let first = wire::read(&mut input).await.unwrap().unwrap();
        let second = wire::read(&mut input).await.unwrap().unwrap();
        assert_eq!(first["id"], 1);
        assert_eq!(second["id"], 2);
        assert_eq!(second["method"], "query");
    }
}