// SPDX-FileCopyrightText: 2024 Christoph Sax <c_sax@mailbox.org>
//
// SPDX-License-Identifier: EUPL-1.2

mod bootloop;
mod doc;
mod lsp;
mod mainloop;
mod proc;
mod request;
mod response;
mod semantic;
mod tasks;
mod transport;
mod workspace;

pub use {doc::TextDoc, workspace::FileIndex};

#[cfg(test)]
pub use crate::ls::{doc::TextDocs, doc::read_doc, tasks::TaskSystem, workspace::index_files};

use std::{
    thread,
    time::{Duration, Instant},
};

use tree_sitter::Tree;

use url;

use crate::{
    ReturnCode,
    config::{Config, Workspace},
    ls::lsp::Message,
    ls::transport::StdioChannel,
    ls::{
        proc::{ProcState, proc_alive},
        request::{Notification, Request},
        response::{ErrorResponse, InitializeResponse, ReceiveError, Response},
        tasks::{OngoingTask, Task, TaskDone},
    },
    protocol::{
        ErrorCodes, InitializeError, InitializeParams, InitializeResult, LogTraceParams,
        NumberOrString, ResponseError, SemanticTokenModifiers, SemanticTokenTypes,
        SemanticTokensLegend, ServerCapabilities, TokenFormat, TraceValue,
    },
    t32::LangExpressions,
};

#[cfg(not(test))]
use crate::ls::{doc::TextDocs, tasks::TaskSystem};

struct DutyCycleBackoff {
    level: i32,
    inactive_since: Instant,
}

struct InitializationStatus {
    msg: Message,
    rc: ReturnCode,
}

struct InitState {
    shutdown_request_recv: bool,
    exit_requested: bool,
    heartbeat: ProcHeartbeat,
    backoff: DutyCycleBackoff,
    tasks: Tasks,
}

struct ProcHeartbeat {
    pid: Option<u32>,
    interval: Duration,
    last_beat: Instant,
}

struct RunState {
    shutdown_request_recv: bool,
    exit_requested: bool,
    heartbeat: ProcHeartbeat,
    backoff: DutyCycleBackoff,
    tasks: Tasks,
    docs: TextDocs,
    files: FileIndex,
}

struct TaskCounterInternal {
    counter: i32,
}

struct TaskCounterExternal {
    counter: i32,
}

struct TaskCounters {
    tasks_int: TaskCounterInternal,
    tasks_ext: TaskCounterExternal,
    progress: ProgressTokenGenerator,
}

struct Tasks {
    runner: TaskSystem,
    blocked: Vec<Task>,
    ongoing: Vec<Option<OngoingTask>>,
    completed: Vec<Option<TaskDone>>,
    counters: TaskCounters,
}

type FileData = (TextDoc, Tree, LangExpressions);
type ProgressTokenGenerator = TaskCounterExternal;

trait TaskCounter {
    fn next_id(&mut self) -> NumberOrString;
}

impl DutyCycleBackoff {
    const IDLE: [Duration; 4] = [
        Duration::from_millis(5),
        Duration::from_millis(20),
        Duration::from_millis(50),
        Duration::from_millis(250),
    ];
    const STAGES: [Duration; 3] = [
        Duration::from_secs(1),
        Duration::from_secs(10),
        Duration::from_secs(100),
    ];
    fn new() -> Self {
        debug_assert_eq!(Self::IDLE.len(), Self::STAGES.len() + 1);

        Self {
            level: 0,
            inactive_since: Instant::now(),
        }
    }

    fn idle(&mut self, now: &Instant) {
        debug_assert!(self.level >= -1 && self.level <= (Self::STAGES.len() as i32));
        debug_assert!(*now >= self.inactive_since);

        // Grant one additional uninterrupted cycle if the server was just busy.
        if self.level < 0 {
            self.level += 1;
            return;
        }
        debug_assert!(self.level >= 0);

        if (self.level as usize) < Self::IDLE.len() - 1
            && *now - self.inactive_since > Self::STAGES[self.level as usize]
        {
            self.level += 1;
        }
        thread::sleep(Self::IDLE[self.level as usize]);
    }

    fn clear(&mut self) {
        debug_assert!(self.level >= -1 && self.level <= (Self::STAGES.len() as i32));
        self.level = -1;
    }
}

impl InitState {
    pub fn new(cfg: &Config) -> Self {
        Self {
            shutdown_request_recv: false,
            exit_requested: false,
            heartbeat: ProcHeartbeat::build(&cfg),
            backoff: DutyCycleBackoff::new(),
            tasks: Tasks::new(),
        }
    }
}

impl ProcHeartbeat {
    fn build(cfg: &Config) -> Self {
        ProcHeartbeat {
            pid: cfg.parent_pid,
            interval: cfg.pid_check_interval,
            last_beat: Instant::now() - cfg.pid_check_interval,
        }
    }

    fn elapsed(&self, now: &Instant) -> bool {
        *now - self.last_beat >= self.interval
    }

    fn check(&mut self, now: &Instant) -> bool {
        self.last_beat = *now;

        let pid = self
            .pid
            .expect("PID must be specified if check is executed.");
        if pid > 0 {
            ProcState::Alive == proc_alive(self.pid.expect("PID must be specified."))
        } else {
            true
        }
    }
}

impl RunState {
    fn from_init(s: InitState, docs: TextDocs, files: FileIndex) -> Self {
        Self {
            shutdown_request_recv: s.shutdown_request_recv,
            exit_requested: s.exit_requested,
            heartbeat: s.heartbeat,
            backoff: s.backoff,
            tasks: s.tasks,
            docs: docs,
            files: files,
        }
    }
}

impl TaskCounterExternal {
    pub fn new() -> Self {
        Self { counter: 0 }
    }
}

impl TaskCounterInternal {
    const PREFIX: &str = "it#";

    pub fn new() -> Self {
        Self { counter: 0 }
    }
}

impl TaskCounters {
    pub fn new() -> Self {
        Self {
            tasks_int: TaskCounterInternal::new(),
            tasks_ext: TaskCounterExternal::new(),
            progress: TaskCounterExternal::new(),
        }
    }
}

impl Tasks {
    pub fn new() -> Self {
        Self {
            runner: TaskSystem::build(),
            blocked: Vec::new(),
            ongoing: Vec::new(),
            completed: Vec::new(),
            counters: TaskCounters::new(),
        }
    }
}

impl TaskCounter for TaskCounterExternal {
    fn next_id(&mut self) -> NumberOrString {
        let id = self.counter;
        self.counter += 1;

        NumberOrString::Number(id)
    }
}

impl TaskCounter for TaskCounterInternal {
    fn next_id(&mut self) -> NumberOrString {
        let id = format!("{}{}", Self::PREFIX, self.counter);
        self.counter += 1;

        NumberOrString::String(id)
    }
}

pub fn serve(mut cfg: Config) -> ReturnCode {
    let mut channel = match transport::build_channel(&cfg) {
        Ok(c) => c,
        Err(rc) => return rc,
    };
    let heartbeat = ProcHeartbeat::build(&cfg);

    let InitializationStatus { msg, rc } =
        wait_for_initialize_req(cfg.trace_level, &mut channel, heartbeat);
    match msg {
        Message::Request(Request::InitializeRequest { id, params }) => {
            match process_initialize_params(params, &mut cfg) {
                Ok(notifs) => {
                    for msg in notifs {
                        channel.send_msg(Message::Notification(msg));
                    }

                    let result = InitializeResult::build(ServerCapabilities::build(
                        cfg.semantic_tokens.legend.clone(),
                    ));
                    channel.send_msg(Message::Response(Response::InitializeResponse(
                        InitializeResponse { id: id, result },
                    )));
                }
                Err(error) => {
                    channel.send_msg(Message::Response(Response::ErrorResponse(ErrorResponse {
                        id: Some(id),
                        error,
                    })));
                    return ReturnCode::ProtocolErr;
                }
            }
        }
        // No shutdown request was received before
        Message::Notification(Notification::ExitNotification {}) => {
            if cfg.trace_level != TraceValue::Off {
                channel.send_msg(Message::Notification(log_notif(
                    &Notification::ExitNotification {},
                )));
            }
            return shutdown(channel, rc);
        }
        _ => unreachable!(),
    }

    if cfg.trace_level != TraceValue::Off {
        channel.send_msg(Message::Notification(notif_initialized()));
    }

    let (g, incoming) = match bootloop::start(InitState::new(&cfg), &mut channel, &mut cfg) {
        Ok((s, m)) => (s, m),
        Err(rc) => return shutdown(channel, rc),
    };

    match mainloop::handle_requests(g, incoming, &mut channel, cfg) {
        Ok(_) => (),
        Err(rc) => return shutdown(channel, rc),
    }

    shutdown(channel, ReturnCode::OkExit);
    ReturnCode::OkExit
}

/// Wait for initialization request. Returns `ServerNotInitialized` error for
/// other types of `RequestMessage`. Notifications are dropped. The only
/// exception is the exit notification after which the server shuts down.
/// Exit notifications without prior shutdown request result should trigger an
/// error exit code. However, sending a shutdown request without prior
/// initialization will return an error response.
fn wait_for_initialize_req(
    trace_level: TraceValue,
    channel: &mut StdioChannel,
    mut heartbeat: ProcHeartbeat,
) -> InitializationStatus {
    let mut backoff = DutyCycleBackoff::new();

    let mut shutdown_request_recv = false;
    loop {
        backoff.idle(&Instant::now());

        let msg = match read_msg(trace_level, channel, &mut heartbeat) {
            Ok(Some(m)) => m,
            Ok(None) => continue,
            Err(rc) => {
                return InitializationStatus {
                    msg: Message::Notification(Notification::ExitNotification {}),
                    rc,
                };
            }
        };

        match msg {
            Message::Request(Request::InitializeRequest { .. })
            | Message::Notification(Notification::ExitNotification { .. }) => {
                return InitializationStatus {
                    msg,
                    rc: if shutdown_request_recv {
                        ReturnCode::OkExit
                    } else {
                        ReturnCode::ErrExit
                    },
                };
            }
            m if m.is_request() => {
                backoff.clear();

                if let Message::Request(Request::ShutdownRequest { .. }) = m {
                    shutdown_request_recv = true;
                }
                let req = m.get_request();
                channel.send_msg(Message::Response(Response::ErrorResponse(ErrorResponse {
                    id: Some(req.get_id().clone()),
                    error: error_not_initialized(),
                })));
            }
            _ => (),
        }
    }
}

fn process_initialize_params(
    params: InitializeParams,
    cfg: &mut Config,
) -> Result<Vec<Notification>, ResponseError> {
    if let Some(pid) = params.process_id {
        let parent_pid = match u32::try_from(pid) {
            Ok(num) => num,
            Err(_) => {
                return Err(error_invalid_pid(pid));
            }
        };

        match cfg.parent_pid {
            Some(ppid) if ppid == parent_pid => (),
            Some(ppid) if ppid > 0 => return Err(error_pid_mismatch(parent_pid, ppid)),
            Some(_) | None => {
                cfg.parent_pid = Some(parent_pid);
            }
        }
    }

    if let Some(level) = &params.trace {
        cfg.trace_level = *level;
    }

    // If the key `workspaceFolders` is present in `InitializeParams`, then we
    // can infer that the client must support the `workspace.workspaceFolders`
    // capability. We don't need to check for it separately.
    if params.workspace_folders.is_some() {
        cfg.workspace = Workspace::Folders(params.workspace_folders);
    } else if params.root_uri.is_some() {
        debug_assert!(url::Url::parse(params.root_uri.as_ref().unwrap()).is_ok());
        cfg.workspace = Workspace::Root(params.root_uri);
    } else if params.root_path.is_some() {
        // This is not guaranteed to be an URI, so we try to convert it into
        // one.
        let dir = match url::Url::from_directory_path(params.root_path.as_ref().unwrap()) {
            Ok(url) => url.to_string(),
            Err(_) => params.root_path.unwrap(),
        };
        cfg.workspace = Workspace::Root(Some(dir));
    }

    // The workspace folder can be `null` if no folder was selected in the client.
    // It is possible to query the current workspace folder selection with a
    // `workspaceFolders` request. The client capabilities tell us if this is
    // supported.
    cfg.workspace_folders_supported = params
        .capabilities
        .workspace
        .as_ref()
        .is_some_and(|ws| ws.workspace_folders.unwrap_or(false));

    // Check whether the client support `LocationLink` in the response results.
    cfg.location_links.definitions_supported = params
        .capabilities
        .text_document
        .as_ref()
        .is_some_and(|td| {
            td.definition
                .as_ref()
                .is_some_and(|def| def.link_support.unwrap_or(false))
        });

    cfg.did_rename_files_supported = params.capabilities.workspace.is_some_and(|ws| {
        ws.file_operations
            .is_some_and(|fop| fop.did_rename.unwrap_or(false))
    });

    // We ignore the `textDocument/semanticTokens/range`,
    // `textDocument/semanticTokens/full`, and
    // `textDocument/semanticTokens/full/delta` capabilitiy indicators of
    // the client. Semantic token requests are sent from client to server, so
    // we try to signal what we support and then handle the requests that we
    // receive.
    let mut notif: Vec<Notification> = Vec::new();
    if let Some(td) = params.capabilities.text_document
        && let Some(st) = td.semantic_tokens
    {
        // Client does not specify a supported token format. See
        // [Note: Semantic Token Format] for more information.
        if !st.formats.iter().any(|f| *f == TokenFormat::Relative) {
            notif.push(warn_unknown_sem_token_format());
        }

        cfg.semantic_tokens.encoding.overlapping_tokens =
            st.overlapping_token_support.is_some_and(|ov| ov);
        cfg.semantic_tokens.encoding.multiline_tokens =
            st.multiline_token_support.is_some_and(|ml| ml);

        cfg.semantic_tokens.legend = {
            let mut types: Vec<SemanticTokenTypes> = Vec::with_capacity(SemanticTokenTypes::num());
            {
                let mut unknown: Vec<&str> = Vec::new();

                for (ii, r#type) in st.token_types.iter().enumerate() {
                    if let Ok(token) = r#type.parse::<SemanticTokenTypes>() {
                        types.push(token);
                    } else {
                        unknown.push(&st.token_types[ii]);
                    }
                }

                if !unknown.is_empty() {
                    notif.push(notif_unknown_client_sem_token_type(&unknown));
                }
            }

            let mut modifiers: Vec<SemanticTokenModifiers> =
                Vec::with_capacity(SemanticTokenModifiers::num());
            {
                let mut unknown: Vec<&str> = Vec::new();

                for (ii, r#modifier) in st.token_modifiers.iter().enumerate() {
                    if let Ok(token) = r#modifier.parse::<SemanticTokenModifiers>() {
                        modifiers.push(token);
                    } else {
                        unknown.push(&st.token_modifiers[ii]);
                    }
                }
                if !unknown.is_empty() {
                    notif.push(notif_unknown_client_sem_token_modifiers(&unknown));
                }
            }
            SemanticTokensLegend::from_attrs(types, modifiers)
        };

        // We ignore the `rangeLimit` hint of the client. The server is always
        // producing all code folds for a file.
        if let Some(cf) = td.folding_range {
            cfg.code_folding.encoding.line_folds_only = cf.line_folding_only.is_some_and(|v| v);
            cfg.code_folding.encoding.collapsed_text_supported = cf
                .folding_range
                .is_some_and(|r| r.collapsed_text.is_some_and(|c| c));

            // If the `folding_range_kind` field is missing, then the client
            // does not guarantee that is accepts any type of code fold. On the
            // other hand, if `folding_range_kind` is present, then the client
            // has to map unknown types to a default value.
            // It does not really make much sense to have a client that signals
            // support for code folds, but does not support any of the range
            // kinds. In this case we can only refrain from using any range
            // kinds in the response.
            debug_assert!(cf.folding_range_kind.is_some());
            if let Some(kinds) = cf.folding_range_kind {
                cfg.code_folding.folding_range_kinds = kinds.value_set;
            } else {
                cfg.code_folding.folding_range_kinds.clear();
            }
        }
    }

    if let Some(win) = params.capabilities.window {
        cfg.server_progress_supported = win.work_done_progress.is_some_and(|p| p);
    }
    Ok(notif)
}

fn recv_incoming(
    trace_level: TraceValue,
    channel: &mut StdioChannel,
    heartbeat: &mut ProcHeartbeat,
    incoming: &mut Vec<Option<Message>>,
) -> Result<bool, ReturnCode> {
    let mut inc_recv: bool = false;
    loop {
        match read_msg(trace_level, channel, heartbeat) {
            Ok(Some(r)) => {
                inc_recv = true;
                incoming.push(Some(r));
            }
            Ok(None) => break,
            Err(rc) => return Err(rc),
        };
    }
    Ok(inc_recv)
}

fn read_msg(
    trace_level: TraceValue,
    channel: &mut StdioChannel,
    heartbeat: &mut ProcHeartbeat,
) -> Result<Option<Message>, ReturnCode> {
    match channel.recv_msg() {
        Ok(Some(r)) => Ok(Some(r)),
        Ok(None) => {
            // The server should shut down, if it detects that its parent
            // process is not alive anymore. No actual shutdown request was
            // received, so we exit with an error code. We only check if we
            // did not receive any new message from the client.
            if let Some(_) = heartbeat.pid {
                let now = Instant::now();
                if heartbeat.elapsed(&now) && !heartbeat.check(&now) {
                    if trace_level != TraceValue::Off {
                        channel.send_msg(Message::Notification(notif_ppid_heartbeat_timeout(
                            heartbeat.pid.unwrap(),
                        )));
                    }
                    return Err(ReturnCode::UnavailableErr);
                }
            }
            Ok(None)
        }
        Err(err) => match err {
            ReceiveError::Response(resp) => {
                if resp.error.code == (ErrorCodes::TransportError as i64) {
                    // The loopback reader thread returns a transport error after
                    // it has terminated. We use this return code as shutdown
                    // trigger.
                    Err(ReturnCode::OkExit)
                } else {
                    // The message could not be parsed, so we have no request ID to
                    // work with.
                    channel.send_msg(Message::Response(Response::ErrorResponse(resp)));

                    Ok(None)
                }
            }
            ReceiveError::Notification(notif) => {
                if trace_level != TraceValue::Off {
                    // The response parsed or malformed. We can only inform the
                    // client via notification.
                    channel.send_msg(Message::Notification(notif));
                }
                Ok(None)
            }
        },
    }
}

fn send_outgoing(channel: &mut StdioChannel, msgs: &mut Vec<Option<Message>>) {
    for msg in msgs {
        let msg = msg.take().expect("No empty slots allowed.");
        channel.send_msg(msg);
    }
}

fn error_not_initialized() -> ResponseError {
    ResponseError {
        code: ErrorCodes::ServerNotInitialized as i64,
        message: "ERROR: Server not initialized. Cannot handle request.".to_string(),
        data: None,
    }
}

fn error_invalid_pid(pid: i64) -> ResponseError {
    ResponseError {
        code: ErrorCodes::InvalidParams as i64,
        message: format!(
            "ERROR: Process ID of the parent process {} is invalid.",
            pid
        ),
        data: Some(
            serde_json::to_value(InitializeError { retry: true }).expect("Must convert to value."),
        ),
    }
}

fn error_pid_mismatch(pid_msg: u32, pid_cli: u32) -> ResponseError {
    ResponseError {
        code: ErrorCodes::InvalidParams as i64,
        message: format!(
            "ERROR: Process ID of the parent process {} is different from the process ID specified by \"--clientProcessId=\" {}.",
            pid_msg, pid_cli
        ),
        data: Some(
            serde_json::to_value(InitializeError { retry: true }).expect("Must convert to value."),
        ),
    }
}

fn error_shutdown_seq(id: NumberOrString) -> Message {
    Message::Response(Response::ErrorResponse(ErrorResponse {
        id: Some(id),
        error: ResponseError {
            code: ErrorCodes::InvalidRequest as i64,
            message: "ERROR: Server has received shutdown request. Cannot handle request."
                .to_string(),
            data: None,
        },
    }))
}

fn error_task_queue_abort() -> ResponseError {
    ResponseError {
        code: ErrorCodes::RequestFailed as i64,
        message: "ERROR: Task queue worker has aborted due to an internal error. Shutting down..."
            .to_string(),
        data: None,
    }
}

fn notif_ppid_heartbeat_timeout(pid: u32) -> Notification {
    Notification::LogTraceNotification {
        params: LogTraceParams {
            message: format!("INFO: Parent process with process ID {} has exited.", pid),
            verbose: Some("Server is shutting down...".to_string()),
        },
    }
}

fn notif_initialized() -> Notification {
    Notification::LogTraceNotification {
        params: LogTraceParams {
            message: "INFO: Server is initialized.".to_string(),
            verbose: None,
        },
    }
}

fn notif_unknown_client_sem_token_type(unknown_toks: &[&str]) -> Notification {
    debug_assert!(!unknown_toks.is_empty());
    Notification::LogTraceNotification {
        params: LogTraceParams {
            message: "INFO: Client supports unknown semantic token types.".to_string(),
            verbose: Some(format!(
                "The types {} are not supported.",
                unknown_toks.join(", ")
            )),
        },
    }
}

fn notif_unknown_client_sem_token_modifiers(unknown_toks: &[&str]) -> Notification {
    debug_assert!(!unknown_toks.is_empty());
    Notification::LogTraceNotification {
        params: LogTraceParams {
            message: "INFO: Client supports unknown semantic token modifiers.".to_string(),
            verbose: Some(format!(
                "The modifiers {} are not supported.",
                unknown_toks.join(", ")
            )),
        },
    }
}

fn warn_unknown_sem_token_format() -> Notification {
    Notification::LogTraceNotification {
        params: LogTraceParams {
            message: "WARNING: Client does not support relative format for semantic tokens."
                .to_string(),
            verbose: None,
        },
    }
}

fn log_notif(msg: &Notification) -> Notification {
    Notification::LogTraceNotification {
        params: LogTraceParams {
            message: format!("INFO: Received notification \"{:}\".", msg),
            verbose: None,
        },
    }
}

/// We ignore the specification here, when it states that the exit code 1 should
/// be returned if no shutdown request has been received. Instead, we try to return
/// a meaningful error code.
fn shutdown(channel: StdioChannel, rc: ReturnCode) -> ReturnCode {
    drop(channel);
    rc
}
