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"),
    }
}