//! Python import/alias resolution and type-inference helpers.
//!
//! Covers:
//! - `add_python_aliases` — harvest `import` / `from … import` alias map entries
//! - `resolve_relative_module` — convert `.`/`..` relative imports to absolute
//! - Type-inference utilities used by `calls.rs`:
//!   `strip_generics`, `container_element_type`, `collect_python_param_types`,
//!   `collect_python_annotated_types`, `infer_iter_element_type`,
//!   `record_case_pattern_bindings`, `record_assignment_types`, `infer_python_type`

use tree_sitter::Node;

pub(crate) use crate::util::text;

/// `List[int]` / `dict[str, X]` / `Optional[Foo]` → leading head identifier.
/// Used so type-aware dispatch can match the annotated type against the
/// extractor's `name_to_id` (which only sees bare class names).
pub(crate) fn strip_generics(ty: &str) -> String {
    ty.trim()
        .split(['[', ',', '|', ' '])
        .next()
        .unwrap_or(ty)
        .trim()
        .to_string()
}

/// Pull the element type out of `list[Foo]` / `Sequence[Foo]` /
/// `tuple[Foo, ...]`. Returns `None` if the annotation isn't a recognized
/// iterable container or the inner head isn't a class-shaped name.
pub(crate) fn container_element_type(ty: &str) -> Option<String> {
    let ty = ty.trim();
    let head = ty.split('[').next()?.trim().rsplit('.').next()?;
    let known = matches!(
        head,
        "list"
            | "List"
            | "Sequence"
            | "Iterable"
            | "Iterator"
            | "tuple"
            | "Tuple"
            | "set"
            | "Set"
            | "frozenset"
            | "FrozenSet"
    );
    if !known {
        return None;
    }
    let open = ty.find('[')?;
    let close = ty.rfind(']')?;
    let inner = ty[open + 1..close].trim();
    let elem = inner.split(',').next()?.trim();
    if elem
        .chars()
        .next()
        .map(|c| c.is_ascii_uppercase())
        .unwrap_or(false)
    {
        Some(elem.to_string())
    } else {
        None
    }
}

pub(crate) fn collect_python_param_types(
    params: Node<'_>,
    src: &str,
    out: &mut std::collections::HashMap<String, String>,
) {
    let mut c = params.walk();
    for p in params.children(&mut c) {
        if p.kind() == "typed_parameter"
            && let (Some(name), Some(ty)) = (
                p.child_by_field_name("name").or_else(|| {
                    let mut cc = p.walk();
                    p.children(&mut cc).find(|x| x.kind() == "identifier")
                }),
                p.child_by_field_name("type"),
            )
        {
            let var = text(name, src).to_string();
            let raw_ty = text(ty, src);
            let stripped = strip_generics(raw_ty);
            // Stash the un-stripped form under `<var>$iter` so the
            // comprehension iter-var typing can recover the element
            // type from `list[Foo]` / `tuple[Foo, …]`.
            if let Some(elem) = container_element_type(raw_ty) {
                out.insert(format!("{var}$iter"), elem);
            }
            out.insert(var, stripped);
        }
    }
}

pub(crate) fn collect_python_annotated_types(
    body: Node<'_>,
    src: &str,
    out: &mut std::collections::HashMap<String, String>,
) {
    let mut cursor = body.walk();
    let mut stack = vec![body];
    while let Some(n) = stack.pop() {
        match n.kind() {
            "expression_statement" => {
                let mut c = n.walk();
                for inner in n.children(&mut c) {
                    if inner.kind() == "assignment" {
                        record_assignment_types(inner, src, out);
                    }
                }
            }
            // Walrus `(x := Foo())` and the rare annotated form
            // `(x: Foo := …)` / `if (x: int := f()) ...`.
            "named_expression" => {
                if let Some(name) = n.child_by_field_name("name")
                    && name.kind() == "identifier"
                {
                    let var = text(name, src).to_string();
                    let ty = n
                        .child_by_field_name("type")
                        .map(|t| strip_generics(text(t, src)))
                        .or_else(|| {
                            n.child_by_field_name("value")
                                .and_then(|v| infer_python_type(v, src))
                        });
                    if let Some(t) = ty {
                        out.insert(var, t);
                    }
                }
            }
            // Comprehension iteration variable: `for x in items` where
            // `items: List[Foo]` lets us bind `x: Foo`. We do a best-effort
            // single-class heuristic on the iterable RHS.
            "for_in_clause" => {
                if let (Some(left), Some(right)) = (
                    n.child_by_field_name("left"),
                    n.child_by_field_name("right"),
                ) && left.kind() == "identifier"
                    && let Some(t) = infer_iter_element_type(right, src, out)
                {
                    out.insert(text(left, src).to_string(), t);
                }
            }
            // `match obj: case Foo(x=y): …` — bind the keyword pattern
            // target identifier `y` to the field name's annotated type
            // when known. Cheap form: bind to the class name's `field` —
            // we don't track per-field types, so fall back to the class
            // name itself when uppercase.
            "case_clause" => {
                let mut cc = n.walk();
                for child in n.children(&mut cc) {
                    record_case_pattern_bindings(child, src, out);
                }
            }
            _ => {}
        }
        for child in n.children(&mut cursor) {
            stack.push(child);
        }
    }
}

/// Best-effort element-type inference for comprehension iterables.
///
/// Handles three iterable forms:
///   `for x in items`       — look up `items` in the type map.
///   `for x in get_items()` — call site; if the function is `Foo(…)` (a
///                            constructor) bind `x` to `Foo`. Otherwise
///                            inspect the call's `function` annotation.
///   `for x in [Foo(), …]`  — list literal; use the first element's
///                            inferred type.
pub(crate) fn infer_iter_element_type(
    rhs: Node<'_>,
    src: &str,
    types: &std::collections::HashMap<String, String>,
) -> Option<String> {
    fn unwrap_container(ty: &str) -> &str {
        let inner = ty
            .find('[')
            .and_then(|i| ty[i + 1..].rfind(']').map(|j| &ty[i + 1..i + 1 + j]))
            .unwrap_or(ty);
        inner.split(',').next().unwrap_or(inner).trim()
    }
    fn is_class_name(s: &str) -> bool {
        s.chars()
            .next()
            .map(|c| c.is_ascii_uppercase())
            .unwrap_or(false)
    }
    match rhs.kind() {
        "identifier" => {
            let var = text(rhs, src);
            // Prefer the `$iter` sentinel — set when the param/binding
            // had a `list[Foo]`-style annotation. Falls back to the
            // generics-stripped form for legacy paths.
            if let Some(elem) = types.get(&format!("{var}$iter"))
                && is_class_name(elem)
            {
                return Some(elem.clone());
            }
            let ty = types.get(var)?;
            let head = unwrap_container(ty);
            if is_class_name(head) {
                Some(head.to_string())
            } else {
                None
            }
        }
        // `for x in get_items()` — if the called function is itself a
        // class constructor (`Foo()`), the iterable is unusual but bind
        // to `Foo`. More commonly, the call is `obj.fetch_all()` and we
        // can't statically know the return type without annotations.
        "call" => {
            let func = rhs.child_by_field_name("function")?;
            match func.kind() {
                "identifier" => {
                    let head = text(func, src);
                    if is_class_name(head) {
                        Some(head.to_string())
                    } else {
                        None
                    }
                }
                "attribute" => {
                    let attr = func.child_by_field_name("attribute")?;
                    let head = text(attr, src);
                    if is_class_name(head) {
                        Some(head.to_string())
                    } else {
                        None
                    }
                }
                _ => None,
            }
        }
        // `for x in [Foo(), …]` / `(Foo(), Bar())` — first element's type.
        "list" | "tuple" | "set" => {
            let mut c = rhs.walk();
            for child in rhs.children(&mut c) {
                if let Some(t) = infer_python_type(child, src) {
                    return Some(t);
                }
            }
            None
        }
        _ => None,
    }
}

/// `case Foo(x=y, z=w): …` — bind `y` and `w` to `Foo` (best-effort).
pub(crate) fn record_case_pattern_bindings(
    n: Node<'_>,
    src: &str,
    out: &mut std::collections::HashMap<String, String>,
) {
    if n.kind() == "class_pattern" {
        let cls = n
            .named_child(0)
            .map(|c| text(c, src).to_string())
            .unwrap_or_default();
        if cls.is_empty()
            || !cls
                .chars()
                .next()
                .map(|c| c.is_ascii_uppercase())
                .unwrap_or(false)
        {
            return;
        }
        let mut c = n.walk();
        for child in n.children(&mut c) {
            // tree-sitter-python represents keyword patterns as
            // `keyword_pattern` with `name` (field) and `value` (target).
            if child.kind() == "keyword_pattern" {
                if let Some(target) = child.child_by_field_name("value") {
                    if target.kind() == "dotted_name"
                        && let Some(id) = target.named_child(0)
                        && id.kind() == "identifier"
                    {
                        out.insert(text(id, src).to_string(), cls.clone());
                    } else if target.kind() == "class_pattern" {
                        // M10: recurse into nested class patterns within keyword patterns,
                        // e.g. `case Point(coord=Vector(x=x_val)): ...`
                        record_case_pattern_bindings(target, src, out);
                    }
                }
            }
        }
    }
}

/// Walk a (possibly chained) `assignment` and bind each LHS identifier to:
///   - its annotated type (`x: T = …`)  → `T` (head, generics stripped)
///   - or the inferred class from RHS    → `T` if `T(...)` / `T()` shape
///
/// Handles `a = b = c = Foo()` by descending into nested `assignment` RHS
/// and binding all left-hand identifiers to the same root type.
pub(crate) fn record_assignment_types(
    a: Node<'_>,
    src: &str,
    out: &mut std::collections::HashMap<String, String>,
) {
    fn collect_lefts<'a>(n: Node<'a>, src: &str, lefts: &mut Vec<String>) -> Option<Node<'a>> {
        let left = n.child_by_field_name("left")?;
        match left.kind() {
            "identifier" => lefts.push(text(left, src).to_string()),
            "pattern_list" => {
                let mut c = left.walk();
                for p in left.children(&mut c) {
                    match p.kind() {
                        "identifier" => lefts.push(text(p, src).to_string()),
                        "list_splat_pattern" => {
                            let mut sc = p.walk();
                            for s in p.children(&mut sc) {
                                if s.kind() == "identifier" {
                                    lefts.push(text(s, src).to_string());
                                }
                            }
                        }
                        _ => {}
                    }
                }
            }
            _ => {}
        }
        let right = n.child_by_field_name("right")?;
        if right.kind() == "assignment" {
            collect_lefts(right, src, lefts)
        } else {
            Some(right)
        }
    }
    let mut lefts: Vec<String> = Vec::new();
    let root_rhs = collect_lefts(a, src, &mut lefts);
    if lefts.is_empty() {
        return;
    }
    // Annotated form takes precedence.
    let ty = a
        .child_by_field_name("type")
        .map(|t| strip_generics(text(t, src)))
        .or_else(|| root_rhs.and_then(|r| infer_python_type(r, src)));
    if let Some(t) = ty {
        for name in lefts {
            out.insert(name, t.clone());
        }
    }
}

/// Infer a class name from RHS expressions like `Foo()` / `Foo(...)`.
pub(crate) fn infer_python_type(val: Node<'_>, src: &str) -> Option<String> {
    if val.kind() == "call"
        && let Some(func) = val.child_by_field_name("function")
        && matches!(func.kind(), "identifier" | "attribute")
    {
        let raw = text(func, src);
        let head = raw.split('.').next_back().unwrap_or(raw);
        if head
            .chars()
            .next()
            .map(|c| c.is_ascii_uppercase())
            .unwrap_or(false)
        {
            return Some(head.to_string());
        }
    }
    None
}

/// Record `import` aliases so call resolution can rewrite bare names.
///
/// `caller_module` (e.g. `pkg.sub.foo`) is used to resolve relative imports:
/// `from . import bar` (from `pkg.sub.foo`) → `bar` → `pkg.sub.bar`;
/// `from ..util import y` → `y` → `pkg.util.y`.
pub(crate) fn add_python_aliases(
    out: &mut std::collections::HashMap<String, String>,
    node: Node<'_>,
    src: &str,
    caller_module: &str,
) {
    let body = text(node, src).trim();
    if let Some(rest) = body.strip_prefix("import ") {
        for part in rest.split(',') {
            let part = part.trim();
            if let Some(idx) = part.find(" as ") {
                let path = part[..idx].trim();
                let alias = part[idx + 4..].trim();
                out.insert(alias.to_string(), path.to_string());
            } else {
                let alias = part.split('.').next().unwrap_or(part);
                out.insert(alias.to_string(), part.to_string());
            }
        }
        return;
    }
    if let Some(rest) = body.strip_prefix("from ")
        && let Some(import_idx) = rest.find(" import ")
    {
        let raw_module = rest[..import_idx].trim();
        let module = resolve_relative_module(raw_module, caller_module);
        let imports = rest[import_idx + 8..]
            .trim()
            .trim_start_matches('(')
            .trim_end_matches(')');
        for part in imports.split(',') {
            let part = part.trim();
            if part.is_empty() || part == "*" {
                continue;
            }
            let (name, alias) = if let Some(idx) = part.find(" as ") {
                (part[..idx].trim(), part[idx + 4..].trim())
            } else {
                (part, part)
            };
            let full = if module.is_empty() {
                name.to_string()
            } else {
                format!("{module}.{name}")
            };
            out.insert(alias.to_string(), full);
        }
    }
}

/// Convert relative import (`.`, `..util`) to absolute given the caller's
/// dotted module path. Non-relative paths pass through unchanged.
pub(crate) fn resolve_relative_module(raw: &str, caller: &str) -> String {
    let dots = raw.chars().take_while(|c| *c == '.').count();
    if dots == 0 {
        return raw.to_string();
    }
    let suffix = raw[dots..].trim();
    // Caller drops `dots` trailing segments to find the package root.
    let parts: Vec<&str> = caller.split('.').collect();
    // The caller module's last segment is the file itself (e.g. `pkg.sub.foo`
    // → file `foo`); a single-dot import resolves relative to its package
    // (`pkg.sub`), so we drop `dots` segments off `parts` (always dropping at
    // least the file segment).
    let take = parts.len().saturating_sub(dots);
    let prefix = parts[..take].join(".");
    match (prefix.is_empty(), suffix.is_empty()) {
        (true, true) => String::new(),
        (true, false) => suffix.to_string(),
        (false, true) => prefix,
        (false, false) => format!("{prefix}.{suffix}"),
    }
}
