plum

// Completion support. Three tiers, matching what's cheaply available:

//

//   1. Keywords + top-level names (fn/enum/variant/trait/const, own file and

//      already-imported ones) — pure AST, no type inference.

//   2. Member completion after `expr.` — needs to know `expr`'s type. There's

//      no real "infer the type at this cursor" facility in `plum_checker`

//      (it type-checks whole functions, not arbitrary positions), so this

//      reconstructs a best-effort local variable environment by walking the

//      enclosing function's OWN body top-to-bottom (see `localVarTypes`) —

//      correct for the common case (a variable assigned once, in a straight

//      line above the cursor), not control-flow- or shadowing-aware.

//   3. Std lib symbols not yet imported, offered with an `additionalTextEdit`

//      that inserts the `import` line — built once per `lib_path` and cached

//      by the caller (parsing every std file on every keystroke would be

//      wasteful).

#![allow(non_snake_case)]

use std::collections::HashMap;

use std::path::Path;

use plum_checker::types::{PlumType, TypeEnv, TypeScheme};

use plum_checker::{inferExpr, methodReceiverName, plumTypeFromAst, CheckCtx};

use plum_core::ast;

use tower_lsp::lsp_types::*;

const KEYWORDS: &[&str] = &[

    "fun", "enum", "trait", "import", "module", "extern", "if", "else", "while", "for", "match",

    "case", "when", "break", "continue", "return", "todo", "assert", "test", "self", "True",

    "False",

];

pub fn keywordItems() -> Vec<CompletionItem> {

    KEYWORDS

        .iter()

        .map(|k| CompletionItem {

            label: k.to_string(),

            kind: Some(CompletionItemKind::KEYWORD),

            ..Default::default()

        })

        .collect()

}

/// Every top-level name `source` declares: free functions, enums (and their

/// bare variant names), traits, consts. Excludes methods (`type_param.is_some()`)

/// — those only make sense after a `.`, handled by `memberItems` instead.

pub fn topLevelItems(source: &ast::Source) -> Vec<CompletionItem> {

    let mut items = Vec::new();

    for item in &source.items {

        match item {

            ast::Item::Fn(f) if f.type_param.is_none() => {

                items.push(CompletionItem {

                    label: f.name.clone(),

                    kind: Some(CompletionItemKind::FUNCTION),

                    detail: Some("fn".to_string()),

                    ..Default::default()

                });

            }

            ast::Item::Fn(_) => {}

            ast::Item::Enum(e) => {

                items.push(CompletionItem {

                    label: e.name.clone(),

                    kind: Some(CompletionItemKind::ENUM),

                    ..Default::default()

                });

                for v in &e.variants {

                    items.push(CompletionItem {

                        label: v.name.clone(),

                        kind: Some(CompletionItemKind::ENUM_MEMBER),

                        detail: Some(format!("{}.{}", e.name, v.name)),

                        ..Default::default()

                    });

                }

            }

            ast::Item::Trait(t) => {

                items.push(CompletionItem {

                    label: t.name.clone(),

                    kind: Some(CompletionItemKind::INTERFACE),

                    ..Default::default()

                });

            }

            ast::Item::Const(c) => {

                items.push(CompletionItem {

                    label: c.name.clone(),

                    kind: Some(CompletionItemKind::CONSTANT),

                    ..Default::default()

                });

            }

            ast::Item::Test(_) => {}

        }

    }

    items

}

/// Replaces the identifier run ending exactly at `pos` (possibly empty, e.g.

/// right after a bare `.`) with a placeholder identifier, so the document

/// parses cleanly for `mergedSource`'s sake — mid-edit, that run is

/// necessarily incomplete (`self.` or `self.ag` while typing `self.age`),

/// which `tree-sitter` turns into an ERROR node, which `loadAndMergeWithOverrides`

/// treats as a hard parse failure for the WHOLE file (see `parseFile`'s

/// `firstErrorNode` check) — without this, completion would fail on almost

/// every real keystroke, since that's exactly when it's requested.

pub fn sanitizeForParsing(text: &str, pos: Position) -> String {

    let mut out = String::with_capacity(text.len() + 8);

    for (i, line) in text.lines().enumerate() {

        if i > 0 {

            out.push('\n');

        }

        if i as u32 == pos.line {

            let chars: Vec<char> = line.chars().collect();

            let end = (pos.character as usize).min(chars.len());

            let mut start = end;

            while start > 0 && (chars[start - 1].is_alphanumeric() || chars[start - 1] == '_') {

                start -= 1;

            }

            out.extend(&chars[..start]);

            out.push_str("zzzzzz");

            out.extend(&chars[end..]);

        } else {

            out.push_str(line);

        }

    }

    out

}

/// If the cursor sits right after `<ident>.<partial>` (or `<ident>.`),

/// returns `<ident>`. `line` is just the cursor's own line text, `char_idx`

/// its column treated as a char index (not real UTF-16 offset math — fine

/// for the ASCII identifiers this looks for).

pub fn receiverBeforeDot(line: &str, char_idx: usize) -> Option<String> {

    let chars: Vec<char> = line.chars().collect();

    let mut i = char_idx.min(chars.len());

    while i > 0 && (chars[i - 1].is_alphanumeric() || chars[i - 1] == '_') {

        i -= 1;

    }

    if i == 0 || chars[i - 1] != '.' {

        return None;

    }

    let dot = i - 1;

    let mut j = dot;

    while j > 0 && (chars[j - 1].is_alphanumeric() || chars[j - 1] == '_') {

        j -= 1;

    }

    if j == dot {

        return None;

    }

    Some(chars[j..dot].iter().collect())

}

/// The `Fn` item (own free functions AND methods) whose declaration is the

/// closest one at or before `cursor_line` (1-based, matching `ast::Fn::line`)

/// — a stand-in for "which function contains the cursor" since blocks don't

/// carry end positions to check containment properly.

pub fn enclosingFn(source: &ast::Source, cursor_line: usize) -> Option<&ast::Fn> {

    source

        .items

        .iter()

        .filter_map(|i| match i {

            ast::Item::Fn(f) if f.line > 0 && f.line <= cursor_line => Some(f),

            _ => None,

        })

        .max_by_key(|f| f.line)

}

/// Best-effort local variable types for `f`'s body, built by walking every

/// `Stmt::Assign` (at any nesting depth) top-to-bottom and re-inferring with

/// `inferExpr` each time — see this module's doc comment for the tradeoffs.

fn localVarTypes(f: &ast::Fn, global: &TypeEnv, ctx: &CheckCtx) -> HashMap<String, PlumType> {

    let mut env = HashMap::new();

    if let Some(recv) = &f.type_param {

        env.insert(

            "self".to_string(),

            plumTypeFromAst(&ast::Type {

                name: recv.clone(),

                generics: vec![],

            }),

        );

    }

    for p in &f.params {

        let ty = match &p.ty {

            ast::ParamType::Type(t) => plumTypeFromAst(t),

            ast::ParamType::Variadic(t) => PlumType::TVariadic(Box::new(plumTypeFromAst(t))),

            ast::ParamType::Fn(params, ret) => PlumType::TFun(

                params.iter().map(plumTypeFromAst).collect(),

                Box::new(

                    ret.as_ref()

                        .map(|r| plumTypeFromAst(r))

                        .unwrap_or(PlumType::TUnit),

                ),

            ),

        };

        env.insert(p.name.clone(), ty);

    }

    if let ast::FnBody::Block(block) = &f.body {

        walkBlock(block, &mut env, global, ctx);

    }

    env

}

fn walkBlock(

    block: &ast::Block,

    env: &mut HashMap<String, PlumType>,

    global: &TypeEnv,

    ctx: &CheckCtx,

) {

    for stmt in &block.stmts {

        walkStmt(stmt, env, global, ctx);

    }

}

fn walkStmt(

    stmt: &ast::Stmt,

    env: &mut HashMap<String, PlumType>,

    global: &TypeEnv,

    ctx: &CheckCtx,

) {

    match stmt {

        ast::Stmt::Assign(a) => {

            for (target, value) in a.targets.iter().zip(a.values.iter()) {

                if let ast::AssignTarget::Var(name) = target {

                    let mut combined = global.clone();

                    for (k, v) in env.iter() {

                        combined.insert(k.clone(), TypeScheme::mono(v.clone()));

                    }

                    if let Ok(t) = inferExpr(value, &combined, ctx) {

                        env.insert(name.clone(), t);

                    }

                }

            }

        }

        ast::Stmt::If(if_) => {

            walkBlock(&if_.body, env, global, ctx);

            for ei in &if_.else_ifs {

                walkBlock(&ei.body, env, global, ctx);

            }

            if let Some(b) = &if_.else_ {

                walkBlock(b, env, global, ctx);

            }

        }

        ast::Stmt::While(w) => walkBlock(&w.body, env, global, ctx),

        ast::Stmt::For(f) => {

            for v in &f.vars {

                env.entry(v.clone()).or_insert(PlumType::TInt);

            }

            walkBlock(&f.body, env, global, ctx);

        }

        ast::Stmt::Match(m) => {

            for case in &m.cases {

                walkBlock(&case.body, env, global, ctx);

            }

        }

        _ => {}

    }

}

/// Completion items for `receiver`'s methods/fields, given the function the

/// cursor is inside (`None` if the cursor isn't inside any function — e.g. a

/// top-level `test` body isn't covered, see this module's doc comment).

pub fn memberItems(

    cursor_fn: Option<&ast::Fn>,

    receiver: &str,

    global_env: &TypeEnv,

    ctx: &CheckCtx,

) -> Vec<CompletionItem> {

    let Some(f) = cursor_fn else {

        return Vec::new();

    };

    let receiver_ty = if receiver == "self" {

        f.type_param.as_ref().map(|r| {

            plumTypeFromAst(&ast::Type {

                name: r.clone(),

                generics: vec![],

            })

        })

    } else {

        localVarTypes(f, global_env, ctx).get(receiver).cloned()

    };

    let Some(ty) = receiver_ty else {

        return Vec::new();

    };

    let Some(recv_name) = methodReceiverName(&ty) else {

        return Vec::new();

    };

    let mut items = Vec::new();

    for ((r, name), fn_ty) in ctx.methods.iter() {

        if r != &recv_name {

            continue;

        }

        let detail = match fn_ty {

            PlumType::TFun(params, ret) => Some(format!(

                "({}) -> {:?}",

                params

                    .iter()

                    .map(|p| format!("{p:?}"))

                    .collect::<Vec<_>>()

                    .join(", "),

                ret

            )),

            other => Some(format!("{other:?}")),

        };

        items.push(CompletionItem {

            label: name.clone(),

            kind: Some(CompletionItemKind::METHOD),

            detail,

            ..Default::default()

        });

    }

    for info in ctx.enum_variants.values() {

        if info.enum_name != recv_name {

            continue;

        }

        for field_name in &info.field_names {

            items.push(CompletionItem {

                label: field_name.clone(),

                kind: Some(CompletionItemKind::FIELD),

                ..Default::default()

            });

        }

    }

    items

}

/// Maps every top-level name declared anywhere under `<lib_path>/plum-std/*.plum`

/// to that file's `import` path (`Str.plum` -> `"std/Str"`) — built by parsing

/// each file standalone (no import resolution needed just to list names).

/// Meant to be computed once and cached by the caller.

pub fn stdLibIndex(lib_path: &Path) -> HashMap<String, String> {

    let mut map = HashMap::new();

    let std_dir = lib_path.join("plum-std");

    let Ok(entries) = std::fs::read_dir(&std_dir) else {

        return map;

    };

    for entry in entries.flatten() {

        let path = entry.path();

        if path.extension().and_then(|e| e.to_str()) != Some("plum") {

            continue;

        }

        let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else {

            continue;

        };

        let Ok(text) = std::fs::read_to_string(&path) else {

            continue;

        };

        let mut parser = tree_sitter::Parser::new();

        if parser

            .set_language(&tree_sitter_plum::LANGUAGE.into())

            .is_err()

        {

            continue;

        }

        let Some(tree) = parser.parse(&text, None) else {

            continue;

        };

        let source = plum_core::AstParser::new(&text).parseSource(tree.root_node());

        let import_path = format!("std/{stem}");

        for item in &source.items {

            let name = match item {

                ast::Item::Fn(f) if f.type_param.is_none() => Some(f.name.clone()),

                ast::Item::Enum(e) => {

                    map.entry(e.name.clone())

                        .or_insert_with(|| import_path.clone());

                    for v in &e.variants {

                        map.entry(v.name.clone())

                            .or_insert_with(|| import_path.clone());

                    }

                    None

                }

                ast::Item::Trait(t) => Some(t.name.clone()),

                ast::Item::Const(c) => Some(c.name.clone()),

                _ => None,

            };

            if let Some(name) = name {

                map.entry(name).or_insert_with(|| import_path.clone());

            }

        }

    }

    map

}

/// A `TextEdit` inserting `import <import_path>` right after the last

/// existing `import` line (or after the `module` line, or at the very top of

/// the file if neither exists) — plain line scanning, since `ast::Import`

/// doesn't carry a position to look this up structurally.

pub fn importInsertionEdit(text: &str, import_path: &str) -> TextEdit {

    let mut last_import_line: Option<usize> = None;

    let mut module_line: Option<usize> = None;

    for (i, line) in text.lines().enumerate() {

        let trimmed = line.trim_start();

        if trimmed.starts_with("import ") {

            last_import_line = Some(i);

        } else if trimmed.starts_with("module ") {

            module_line = Some(i);

        }

    }

    let at_line = last_import_line.or(module_line).map(|l| l + 1).unwrap_or(0) as u32;

    TextEdit {

        range: Range {

            start: Position::new(at_line, 0),

            end: Position::new(at_line, 0),

        },

        new_text: format!("import {import_path}\n"),

    }

}