use super::LanguageProfile;
use crate::parser::PredicateKey;
use std::collections::HashMap;

/// Creates the profile for the Python language.
pub(super) fn create_python_profile() -> LanguageProfile {
    let language = tree_sitter_python::LANGUAGE.into();
    let mut queries = HashMap::new();

    let class_query = "(class_definition name: (identifier) @match)";
    let func_query = "(function_definition name: (identifier) @match)";

    queries.insert(PredicateKey::Def, [class_query, func_query].join("\n"));
    queries.insert(PredicateKey::Class, class_query.to_string());
    queries.insert(PredicateKey::Func, func_query.to_string());

    // Query for `import` and `from ... import` statements.
    queries.insert(
        PredicateKey::Import,
        "
        [
            (import_statement) @match
            (import_from_statement) @match
        ]
        "
        .to_string(),
    );

    // Query for function and method call sites.
    queries.insert(
        PredicateKey::Call,
        "
       (call
           function: [
               (identifier) @match
               (attribute attribute: (identifier) @match)
           ]
       )
       "
        .to_string(),
    );

    queries.insert(PredicateKey::Comment, "(comment) @match".to_string());
    queries.insert(PredicateKey::Str, "(string) @match".to_string());

    LanguageProfile {
        name: "Python",
        extensions: vec!["py"],
        language,
        queries,
    }
}
