use crate::{
    Documents, ProjectService, Result,
    project::{Project, canonical_uri},
    wire,
};
use serde_json::{Value, json};
use std::{
    collections::{BTreeMap, BTreeSet, HashSet},
    path::PathBuf,
    sync::{Arc, Mutex},
};
use tokio::{
    io::{AsyncRead, AsyncWrite, BufReader},
    sync::{Notify, mpsc},
};
use tokio_util::sync::CancellationToken;

struct Session {
    service: Box<dyn ProjectService>,
    documents: Documents,
    failure: Option<String>,
    binary: PathBuf,
    project: Option<Project>,
    initialized: bool,
    shutdown: bool,
    client: Value,
    preferences: Value,
    capabilities: Value,
    methods: BTreeSet<String>,
    open: BTreeMap<String, i64>,
    published: BTreeSet<String>,
    dirty: bool,
    refresh_request: u64,
}
impl Session {
    async fn rebuild(&mut self) -> Result<()> {
        self.service.changed();
        self.dirty = false;
        if let Some(project) = &mut self.project {
            // Teardown still kills/reaps a failed backend; allow the next generation
            // to recover rather than retaining a permanently failed connection.
            if let Err(error) = project.close().await {
                eprintln!("LSP backend cleanup: {error}");
            }
        }
        self.project = None;
        let result = async {
            let input = self.service.refresh(self.documents.clone()).await?;
            let project = Project::create(&self.binary, input).await?;
            let capabilities = project.capabilities(&self.client).await?;
            let methods: BTreeSet<_> = capabilities.methods.into_iter().collect();
            if self.initialized
                && (capabilities.capabilities != self.capabilities || methods != self.methods)
            {
                return Err("TypeScript backend protocol changed during session".into());
            }
            if !self.initialized {
                self.capabilities = capabilities.capabilities;
                self.methods = methods;
            }
            Ok::<_, crate::Error>(project)
        }
        .await;
        match result {
            Ok(project) => {
                self.project = Some(project);
                self.failure = None;
                Ok(())
            }
            Err(error) => {
                self.failure = Some(error.to_string());
                Err(error)
            }
        }
    }
    async fn publish(&mut self, writer: &mut (impl AsyncWrite + Unpin)) -> Result<()> {
        let mut current: BTreeSet<String> = self.open.keys().cloned().collect();
        if let Some(project) = &self.project {
            current.extend(project.error_uris().cloned());
        }
        for uri in current.union(&self.published) {
            let diagnostics = if !current.contains(uri) {
                json!([])
            } else if let Some(project) = &mut self.project {
                match project
                    .diagnostics(uri, &self.client, &self.preferences)
                    .await
                {
                    Ok(diagnostics) => diagnostics,
                    Err(error) => {
                        self.failure = Some(error.to_string());
                        json!([{"range":{"start":{"line":0,"character":0},"end":{"line":0,"character":0}},
                            "severity":1,"source":"avril","message":error.to_string()}])
                    }
                }
            } else {
                json!([{"range":{"start":{"line":0,"character":0},"end":{"line":0,"character":0}},
                    "severity":1,"source":"avril","message":self.failure.as_deref().unwrap_or("Project unavailable")}])
            };
            let mut params = json!({"uri": uri, "diagnostics": diagnostics});
            if let Some(version) = self.open.get(uri) {
                params["version"] = json!(version);
            }
            wire::write(
                writer,
                json!({"method": "textDocument/publishDiagnostics", "params": params}),
            )
            .await?;
        }
        self.published = current;
        for feature in ["codeLens", "inlayHint", "semanticTokens"] {
            if self
                .client
                .pointer(&format!("/workspace/{feature}/refreshSupport"))
                .and_then(Value::as_bool)
                == Some(true)
            {
                self.refresh_request += 1;
                wire::write(
                    writer,
                    json!({
                        "id":format!("avril-refresh-{}",self.refresh_request),
                        "method":format!("workspace/{feature}/refresh"),
                    }),
                )
                .await?;
            }
        }
        Ok(())
    }
    async fn handle(
        &mut self,
        method: &str,
        params: Value,
        writer: &mut (impl AsyncWrite + Unpin),
    ) -> Result<Value> {
        match method {
            "initialize" => {
                if self.initialized {
                    return Err("Already initialized".into());
                }
                self.client = params.get("capabilities").cloned().unwrap_or(json!({}));
                self.preferences = params
                    .pointer("/initializationOptions/preferences")
                    .cloned()
                    .unwrap_or(json!({}));
                if !self.preferences.is_object() {
                    return Err(wire::RpcError::invalid("preferences must be an object").into());
                }
                // Resolve data is bound independently for every completion item.
                if let Some(list) = self
                    .client
                    .pointer_mut("/textDocument/completion/completionList")
                    .and_then(Value::as_object_mut)
                {
                    list.remove("itemDefaults");
                }
                self.rebuild().await?;
                self.initialized = true;
                let mut capabilities = self.capabilities.clone();
                capabilities["positionEncoding"] = json!("utf-16");
                capabilities["textDocumentSync"] = json!({"openClose":true, "change":1});
                return Ok(json!({
                    "capabilities": capabilities,
                    "serverInfo": {"name": "avril-lsp", "version": env!("CARGO_PKG_VERSION")}
                }));
            }
            "initialized" => return Ok(Value::Null),
            "workspace/didChangeWatchedFiles" | "workspace/didChangeConfiguration" => {
                if method == "workspace/didChangeConfiguration"
                    && let Some(preferences) = params.pointer("/settings/preferences")
                {
                    if !preferences.is_object() {
                        return Err(wire::RpcError::invalid("preferences must be an object").into());
                    }
                    self.preferences = preferences.clone();
                }
                if self.initialized && !self.shutdown {
                    self.dirty = true;
                }
                return Ok(Value::Null);
            }
            "shutdown" => {
                self.shutdown = true;
                let cleanup = if let Some(project) = &mut self.project {
                    project.close().await
                } else {
                    Ok(())
                };
                self.project = None;
                cleanup?;
                return Ok(Value::Null);
            }
            _ => {}
        }
        let uri = params
            .pointer("/textDocument/uri")
            .and_then(Value::as_str)
            .map(canonical_uri)
            .transpose()?;
        match method {
            "textDocument/didOpen" => {
                let uri = uri.ok_or("Missing document URI")?;
                if !uri.starts_with("file:")
                    && !self.project.as_ref().is_some_and(|p| p.contains(&uri))
                {
                    return Ok(Value::Null);
                }
                let text = params["textDocument"]["text"]
                    .as_str()
                    .ok_or("Invalid didOpen text")?;
                let version = params["textDocument"]["version"]
                    .as_i64()
                    .ok_or("Invalid document version")?;
                if self.open.contains_key(&uri) {
                    return Err("Document already open".into());
                }
                self.open.insert(uri.clone(), version);
                if uri.starts_with("file:")
                    && !self.project.as_ref().is_some_and(|p| p.is_read_only(&uri))
                {
                    self.documents.insert(uri, text.into());
                    self.dirty = true;
                }
                Ok(Value::Null)
            }
            "textDocument/didChange" => {
                let uri = uri.ok_or("Missing document URI")?;
                let version = params["textDocument"]["version"]
                    .as_i64()
                    .ok_or("Invalid document version")?;
                if self.open.get(&uri).is_none_or(|old| version <= *old) {
                    return Ok(Value::Null);
                }
                let changes = params["contentChanges"]
                    .as_array()
                    .ok_or("Invalid document changes")?;
                if changes.is_empty()
                    || changes
                        .iter()
                        .any(|c| c.get("range").is_some() || !c["text"].is_string())
                {
                    return Err("Only full-document changes are supported".into());
                }
                let text = changes
                    .last()
                    .and_then(|c| c["text"].as_str())
                    .ok_or("Missing change text")?;
                self.open.insert(uri.clone(), version);
                if uri.starts_with("file:")
                    && !self.project.as_ref().is_some_and(|p| p.is_read_only(&uri))
                {
                    self.documents.insert(uri, text.into());
                    self.dirty = true;
                }
                Ok(Value::Null)
            }
            "textDocument/didClose" => {
                let uri = uri.ok_or("Missing document URI")?;
                if self.open.remove(&uri).is_none() {
                    return Ok(Value::Null);
                }
                wire::write(
                    writer,
                    json!({
                        "method": "textDocument/publishDiagnostics",
                        "params": {"uri": uri, "diagnostics": []}
                    }),
                )
                .await?;
                if self.documents.remove(&uri).is_some() {
                    self.dirty = true;
                }
                Ok(Value::Null)
            }
            "textDocument/diagnostic" => {
                let uri = uri.ok_or("Missing document URI")?;
                let items = self
                    .project
                    .as_mut()
                    .ok_or("Project unavailable")?
                    .diagnostics(&uri, &self.client, &self.preferences)
                    .await?;
                Ok(json!({"kind":"full","items":items}))
            }
            "avril/readDocument" => {
                let project = self.project.as_mut().ok_or("Project unavailable")?;
                let uri = canonical_uri(params["uri"].as_str().ok_or("Missing document URI")?)?;
                if uri.starts_with("avril-builtin:/typescript/") {
                    project
                        .language_service(
                            method,
                            json!({"textDocument":{"uri":uri}}),
                            &self.client,
                            &self.preferences,
                        )
                        .await
                } else {
                    Ok(json!(project.read_document(&uri)?))
                }
            }
            _ if self.methods.contains(method) => {
                self.project
                    .as_mut()
                    .ok_or("Project unavailable")?
                    .language_service(method, params, &self.client, &self.preferences)
                    .await
            }
            _ => Err(format!("Unsupported method: {method}").into()),
        }
    }
}

/// Owns a workspace session and its native backend. Acquisition and resolution
/// remain in the host's project service; no guest code is executed here.
pub async fn serve<R, W>(
    reader: R,
    mut writer: W,
    binary: PathBuf,
    service: Box<dyn ProjectService>,
) -> Result<()>
where
    R: AsyncRead + Unpin + Send + 'static,
    W: AsyncWrite + Unpin,
{
    let mut session = Session {
        service,
        documents: Documents::new(),
        failure: None,
        binary,
        project: None,
        initialized: false,
        shutdown: false,
        client: json!({}),
        preferences: json!({}),
        capabilities: json!({}),
        methods: BTreeSet::new(),
        open: BTreeMap::new(),
        published: BTreeSet::new(),
        dirty: false,
        refresh_request: 0,
    };
    #[cfg(unix)]
    let mut terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
    let interrupt = async {
        #[cfg(unix)]
        {
            tokio::select! {
                result = tokio::signal::ctrl_c() => result,
                _ = terminate.recv() => Ok(()),
            }
        }
        #[cfg(not(unix))]
        {
            tokio::signal::ctrl_c().await
        }
    };
    tokio::pin!(interrupt);
    let disconnected = CancellationToken::new();
    let pending = Arc::new(Mutex::new(HashSet::<String>::new()));
    let cancelled = Arc::new(Mutex::new(HashSet::<String>::new()));
    let cancellation = Arc::new(Notify::new());
    let (tx, mut rx) = mpsc::channel(64);
    let reader_task = tokio::spawn({
        let disconnected = disconnected.clone();
        let pending = pending.clone();
        let cancelled = cancelled.clone();
        let cancellation = cancellation.clone();
        async move {
            let mut reader = BufReader::new(reader);
            loop {
                let message = match wire::read(&mut reader).await {
                    Ok(Some(message)) => message,
                    Ok(None) => {
                        disconnected.cancel();
                        return;
                    }
                    Err(error) => {
                        let _ = tx.send(Err(error)).await;
                        disconnected.cancel();
                        return;
                    }
                };
                if message["method"] == "$/cancelRequest" {
                    let key = message["params"]["id"].to_string();
                    if pending.lock().unwrap().contains(&key) {
                        cancelled.lock().unwrap().insert(key);
                        cancellation.notify_one();
                    }
                    continue;
                }
                if message["method"].is_string()
                    && let Some(id) = message.get("id")
                {
                    pending.lock().unwrap().insert(id.to_string());
                }
                let exit = message["method"] == "exit";
                if tx.send(Ok(message)).await.is_err() {
                    return;
                }
                // Do not start a blocking stdin read after an explicit exit.
                if exit {
                    return;
                }
            }
        }
    });
    let mut polling = tokio::time::interval(std::time::Duration::from_millis(100));
    polling.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
    let debounce = tokio::time::sleep(std::time::Duration::from_millis(100));
    tokio::pin!(debounce);
    let result: Result<()> = async {
        loop {
            let message = tokio::select! {
                biased;
                message = rx.recv() => match message {
                    Some(message) => message?,
                    None => return if session.shutdown { Ok(()) } else {
                        Err("Editor connection closed without shutdown".into())
                    }
                },
                _ = &mut debounce, if session.dirty && session.initialized && !session.shutdown => {
                    let rebuilt = session.rebuild().await;
                    session.publish(&mut writer).await?;
                    if let Err(error) = rebuilt {
                        wire::write(&mut writer, json!({"method":"window/logMessage","params":{"type":1,"message":error.to_string()}})).await?;
                    }
                    continue;
                },
                _ = polling.tick(), if session.initialized && !session.shutdown => {
                    if session.service.changed() {
                        session.dirty = true;
                        debounce.as_mut().reset(tokio::time::Instant::now() + std::time::Duration::from_millis(100));
                    }
                    continue;
                },
                _ = disconnected.cancelled() => return Err("Editor disconnected".into()),
                _ = &mut interrupt => return Err("LSP interrupted".into()),
            };
            let Some(method) = message["method"].as_str() else {
                continue;
            };
            if message.get("id").is_none() && matches!(method, "$/setTrace" | "textDocument/didSave") {
                continue;
            }
            if method == "exit" {
                return if session.shutdown {
                    Ok(())
                } else {
                    Err("LSP exit without shutdown".into())
                };
            }
            let id = message.get("id").cloned();
            let key = id.as_ref().map(Value::to_string);
            let cancel = || {
                key.as_ref()
                    .is_some_and(|k| cancelled.lock().unwrap().contains(k))
            };
            let mut error_code = -32603;
            let query = id.is_some() && session.initialized && !session.shutdown && session.methods.contains(method);
            let recovery = if query
                && !cancel()
                && (session.dirty
                    || session.project.as_ref().is_none_or(Project::failed))
            {
                let result = session.rebuild().await;
                if result.is_ok() {
                    session.publish(&mut writer).await?;
                }
                Some(result)
            } else {
                None
            };
            let response = if let Some(Err(error)) = recovery {
                Err(error)
            } else if cancel() {
                error_code = -32800;
                Err("Request cancelled".into())
            } else if (!session.initialized && method != "initialize") || session.shutdown {
                error_code = if session.shutdown { -32600 } else { -32002 };
                Err("Server not initialized or has shut down".into())
            } else {
                if !matches!(
                    method,
                    "initialize"
                        | "initialized"
                        | "shutdown"
                        | "textDocument/didOpen"
                        | "textDocument/didChange"
                        | "textDocument/didClose"
                        | "avril/readDocument"
                        | "workspace/didChangeWatchedFiles"
                        | "workspace/didChangeConfiguration"
                ) && !session.methods.contains(method)
                {
                    error_code = -32601;
                }
                tokio::select! {
                    result = session.handle(method,message["params"].clone(),&mut writer) => result,
                    _ = async {
                        loop {
                            let notified = cancellation.notified();
                            if cancel() { break; }
                            notified.await;
                        }
                    }, if session.methods.contains(method) => {
                        error_code = -32800;
                        Err("Request cancelled".into())
                    },
                    _ = disconnected.cancelled() => return Err("Editor disconnected".into()),
                    _ = &mut interrupt => return Err("LSP interrupted".into()),
                }
            };
            if let Some(id) = id {
                let response = if cancel() {
                    json!({"id": id, "error": {"code": -32800, "message": "Request cancelled"}})
                } else {
                    match response {
                        Ok(value) => json!({"id": id, "result": value}),
                        Err(error) => {
                            let native = error.downcast_ref::<wire::RpcError>();
                            json!({"id":id, "error": {
                                "code": native.map_or(error_code, |e| e.code),
                                "message": error.to_string(),
                                "data": native.and_then(|e| e.data.as_ref()),
                            }})
                        }
                    }
                };
                wire::write(&mut writer, response).await?;
            } else if let Err(error) = response {
                wire::write(
                    &mut writer,
                    json!({
                        "method": "window/logMessage",
                        "params": {"type": 1, "message": error.to_string()}
                    }),
                )
                .await?;
            }
            if session.dirty {
                debounce.as_mut().reset(tokio::time::Instant::now() + std::time::Duration::from_millis(100));
            }
            if let Some(key) = key {
                pending.lock().unwrap().remove(&key);
                cancelled.lock().unwrap().remove(&key);
            }
        }
    }
    .await;
    reader_task.abort();
    let _ = reader_task.await;
    let cleanup = if let Some(mut project) = session.project {
        if result.is_err() {
            project.terminate().await
        } else {
            project.close().await
        }
    } else {
        Ok(())
    };
    result.and(cleanup)
}