lince-lingua 0.7.0

use std::{collections::HashMap, io::Read, path::PathBuf, sync::Arc};

use clap::{Parser, Subcommand};
use tokio::sync::Mutex;
use tower_lsp::{
    Client, LanguageServer, LspService, Server,
    jsonrpc::Result as LspResult,
    lsp_types::{
        DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
        DocumentFormattingParams, InitializeParams, InitializeResult, InitializedParams,
        MessageType, OneOf, Position, Range, ServerCapabilities, TextDocumentSyncCapability,
        TextDocumentSyncKind, TextEdit,
    },
};

#[derive(Debug, Parser)]
#[command(name = "lingua")]
#[command(about = "Lingua parser, formatter, and language server")]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Debug, Subcommand)]
enum Command {
    /// Format Lingua source.
    Fmt {
        /// Read source from stdin.
        #[arg(short, long, conflicts_with = "files")]
        stdin: bool,

        /// Rewrite files in place.
        #[arg(short, long, conflicts_with = "stdin")]
        write: bool,

        /// Check that files are already formatted.
        #[arg(long, conflicts_with_all = ["stdin", "write"])]
        check: bool,

        /// Files to format. Use `--stdin` to read from stdin.
        files: Vec<PathBuf>,
    },

    /// Parse Lingua source and print the Rust AST.
    Parse {
        /// Read source from stdin.
        #[arg(short, long, conflicts_with = "file")]
        stdin: bool,

        /// File to parse.
        file: Option<PathBuf>,
    },

    /// Run the Lingua language server over stdio.
    Lsp,
}

#[tokio::main]
async fn main() {
    if let Err(err) = run(Cli::parse()).await {
        eprintln!("{err}");
        std::process::exit(1);
    }
}

async fn run(cli: Cli) -> Result<(), String> {
    match cli.command {
        Command::Fmt {
            stdin,
            write,
            check,
            files,
        } => run_fmt(stdin, write, check, files),
        Command::Parse { stdin, file } => run_parse(stdin, file),
        Command::Lsp => {
            run_lsp().await;
            Ok(())
        }
    }
}

fn run_fmt(stdin: bool, write: bool, check: bool, files: Vec<PathBuf>) -> Result<(), String> {
    if stdin {
        let source = read_stdin()?;
        let formatted = format_source(&source)?;
        print!("{formatted}");
        return Ok(());
    }

    if files.is_empty() {
        return Err("fmt needs files, or --stdin".to_owned());
    }

    let mut failed = false;
    for file in files {
        let source = std::fs::read_to_string(&file)
            .map_err(|err| format!("failed to read {}: {err}", file.display()))?;
        let formatted = format_source(&source)?;

        if check {
            if source != formatted {
                eprintln!("not formatted: {}", file.display());
                failed = true;
            }
        } else if write {
            std::fs::write(&file, formatted)
                .map_err(|err| format!("failed to write {}: {err}", file.display()))?;
        } else {
            print!("{formatted}");
            if !formatted.ends_with('\n') {
                println!();
            }
        }
    }

    if failed {
        Err("format check failed".to_owned())
    } else {
        Ok(())
    }
}

fn run_parse(stdin: bool, file: Option<PathBuf>) -> Result<(), String> {
    let source = if stdin {
        read_stdin()?
    } else {
        let file = file.ok_or_else(|| "parse needs a file, or --stdin".to_owned())?;
        std::fs::read_to_string(&file)
            .map_err(|err| format!("failed to read {}: {err}", file.display()))?
    };

    let program = lingua::parse(&source).map_err(format_parse_errors)?;
    println!("{program:#?}");
    Ok(())
}

fn read_stdin() -> Result<String, String> {
    let mut source = String::new();
    std::io::stdin()
        .read_to_string(&mut source)
        .map_err(|err| format!("failed to read stdin: {err}"))?;
    Ok(source)
}

fn format_source(source: &str) -> Result<String, String> {
    lingua::format(source).map_err(format_parse_errors)
}

fn format_parse_errors(errors: Vec<rust_sitter::errors::ParseError>) -> String {
    format!("failed to parse Lingua source: {errors:#?}")
}

async fn run_lsp() {
    let stdin = tokio::io::stdin();
    let stdout = tokio::io::stdout();
    let (service, socket) = LspService::new(|client| Backend {
        client,
        documents: Arc::default(),
    });

    Server::new(stdin, stdout, socket).serve(service).await;
}

struct Backend {
    client: Client,
    documents: Arc<Mutex<HashMap<String, String>>>,
}

#[tower_lsp::async_trait]
impl LanguageServer for Backend {
    async fn initialize(&self, _: InitializeParams) -> LspResult<InitializeResult> {
        Ok(InitializeResult {
            capabilities: ServerCapabilities {
                text_document_sync: Some(TextDocumentSyncCapability::Kind(
                    TextDocumentSyncKind::FULL,
                )),
                document_formatting_provider: Some(OneOf::Left(true)),
                ..ServerCapabilities::default()
            },
            ..InitializeResult::default()
        })
    }

    async fn initialized(&self, _: InitializedParams) {
        self.client
            .log_message(MessageType::INFO, "Lingua language server initialized")
            .await;
    }

    async fn shutdown(&self) -> LspResult<()> {
        Ok(())
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let uri = params.text_document.uri.to_string();
        let text = params.text_document.text;
        self.documents.lock().await.insert(uri, text);
    }

    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        let Some(change) = params.content_changes.into_iter().last() else {
            return;
        };

        self.documents
            .lock()
            .await
            .insert(params.text_document.uri.to_string(), change.text);
    }

    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        self.documents
            .lock()
            .await
            .remove(&params.text_document.uri.to_string());
    }

    async fn formatting(
        &self,
        params: DocumentFormattingParams,
    ) -> LspResult<Option<Vec<TextEdit>>> {
        let uri = params.text_document.uri.to_string();
        let Some(source) = self.documents.lock().await.get(&uri).cloned() else {
            return Ok(None);
        };

        match lingua::format(&source) {
            Ok(formatted) if formatted == source => Ok(Some(Vec::new())),
            Ok(formatted) => Ok(Some(vec![TextEdit {
                range: whole_document_range(&source),
                new_text: formatted,
            }])),
            Err(errors) => {
                self.client
                    .show_message(
                        MessageType::ERROR,
                        format!("Lingua formatting failed: {errors:#?}"),
                    )
                    .await;
                Ok(None)
            }
        }
    }
}

fn whole_document_range(source: &str) -> Range {
    let mut line = 0;
    let mut character = 0;

    for ch in source.chars() {
        if ch == '\n' {
            line += 1;
            character = 0;
        } else {
            character += ch.len_utf16() as u32;
        }
    }

    Range {
        start: Position::new(0, 0),
        end: Position::new(line, character),
    }
}