mod.rs

use crate::analysis::checker::Checker;
use crate::analysis::class_helpers::{ClassDefHelpers, QualifiedNameHelpers};
use crate::analysis::class_type_detector::ClassTypeDetector;
use crate::analysis::parameter_generator::ParameterGenerator;
use crate::analysis::type_analyzer;
use crate::ast;
use crate::render::renderer::{
    Attribute, ClassNode, CompositionEdge, CompositionKind, Diagram, MethodSignature, RelationType,
    RelationshipEdge, Visibility,
};
use indexmap::IndexSet;
use ruff_linter::source_kind::SourceKind;
use ruff_linter::Locator;
use ruff_python_ast::name::{QualifiedName, UnqualifiedName};
use ruff_python_ast::{Expr, Number, PySourceType};
use ruff_python_codegen::Stylist;
use ruff_python_parser::parse_unchecked_source;
use ruff_python_semantic::analyze::visibility::{
    is_abstract, is_classmethod, is_final, is_overload, is_override, is_property, is_staticmethod,
};
use ruff_python_semantic::{Module, ModuleKind, ModuleSource, SemanticModel};
use ruff_python_stdlib::typing::simple_magic_return_type;
use std::path::Path;

/// Represents a class member (attribute or method) during processing
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum ClassMember {
    Attribute(Attribute),
    Method(MethodSignature),
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum BaseKind {
    Skip,
    InheritanceTarget {
        name: String,
        is_abstract_or_protocol: bool,
    },
}

pub struct ClassDiagram {
    diagram: Diagram,
    options: crate::render::mermaid_renderer::RenderOptions,
    pub path: String,
}

impl Default for ClassDiagram {
    fn default() -> Self {
        Self::new(crate::render::mermaid_renderer::RenderOptions::default())
    }
}

impl ClassDiagram {
    #[must_use]
    pub fn new(options: crate::render::mermaid_renderer::RenderOptions) -> Self {
        Self {
            diagram: Diagram::new(),
            options,
            path: String::new(),
        }
    }

    pub const fn set_hide_private_members(&mut self, hide: bool) {
        self.options.hide_private_members = hide;
    }

    #[must_use]
    pub const fn is_empty(&self) -> bool {
        self.diagram.is_empty()
    }

    #[must_use]
    pub fn render(&self) -> Option<String> {
        if self.is_empty() {
            return None;
        }

        let title = if self.path.is_empty() {
            None
        } else {
            Some(self.path.as_str())
        };

        crate::render::mermaid_renderer::render_diagram(&self.diagram, title, &self.options)
    }

    pub fn add_class(
        &mut self,
        checker: &Checker,
        class: &ast::StmtClassDef,
        _indent_level: usize,
    ) {
        let class_name = class.name.to_string();

        // Find generic type parameters - either from explicit [T] syntax or Generic[T] bases
        let generic_type_var = class.type_params.as_ref().map_or_else(
            || {
                // Check for Generic[T] in bases
                let mut found = None;
                for base in class.bases() {
                    if let Some(type_var) = type_analyzer::extract_generic_params(base, checker) {
                        found = Some(type_var);
                        break;
                    }
                }
                found
            },
            |params| {
                // Explicit type parameters via [T] syntax (Python 3.12+)
                let raw_params = checker.locator().slice(params.as_ref());
                // Remove the brackets to get just the type names
                Some(
                    raw_params
                        .trim_start_matches('[')
                        .trim_end_matches(']')
                        .to_owned(),
                )
            },
        );

        // Detect composition relationships from class attributes
        let mut composition_types: IndexSet<(String, CompositionKind)> = IndexSet::new();
        for stmt in &class.body {
            if let ast::Stmt::AnnAssign(ast::StmtAnnAssign { annotation, .. }) = stmt {
                composition_types.extend(type_analyzer::extract_composition_types(
                    annotation.as_ref(),
                    checker,
                ));
            }
        }

        // Process class body statements
        let mut members: IndexSet<ClassMember> = IndexSet::new();
        for stmt in &class.body {
            if let Some(member) = Self::process_stmt_to_member(checker, stmt) {
                members.insert(member);
            }
        }

        // Detect class type using ClassTypeDetector
        let detector = ClassTypeDetector::new(checker);
        let class_type = detector.detect_type(class);
        let class_is_enum = class.is_enum(checker.semantic());

        // Split members into attributes and methods
        let mut attributes = Vec::new();
        let mut methods = Vec::new();
        for member in members {
            match member {
                ClassMember::Attribute(attr) => attributes.push(attr),
                ClassMember::Method(method) => methods.push(method),
            }
        }

        let class_node = ClassNode {
            name: class_name.clone(),
            type_params: generic_type_var,
            class_type,
            attributes,
            methods,
        };

        self.diagram.add_class(class_node);

        // Handle inheritance relationships
        for base in class.bases() {
            let BaseKind::InheritanceTarget {
                name,
                is_abstract_or_protocol,
            } = self.classify_base(checker, &detector, base, class_is_enum)
            else {
                continue;
            };
            let rel = RelationshipEdge {
                from: class_name.clone(),
                to: name,
                relation_type: if is_abstract_or_protocol {
                    RelationType::Implementation
                } else {
                    RelationType::Inheritance
                },
            };
            self.diagram.add_relationship(rel);
        }

        // Add composition relationships
        for (comp_type, kind) in &composition_types {
            let comp_display = comp_type.split('.').next_back().unwrap_or(comp_type);
            let comp = CompositionEdge {
                container: class_name.clone(),
                contained: comp_display.to_string(),
                kind: *kind,
            };
            self.diagram.add_composition(comp);
        }
    }

    /// Returns true if the function is a property setter or deleter (e.g. @name.setter, @name.deleter).
    /// These are implementation details and should be omitted from the diagram.
    fn is_property_setter_or_deleter(decorator_list: &[ast::Decorator], fn_name: &str) -> bool {
        decorator_list.iter().any(|decorator| {
            UnqualifiedName::from_expr(&decorator.expression).is_some_and(|name| {
                let segs = name.segments();
                (segs == [fn_name, "setter"]) || (segs == [fn_name, "deleter"])
            })
        })
    }

    #[allow(clippy::too_many_lines)]
    fn process_stmt_to_member(checker: &Checker, stmt: &ast::Stmt) -> Option<ClassMember> {
        match stmt {
            ast::Stmt::AnnAssign(ast::StmtAnnAssign {
                target,
                annotation,
                simple,
                ..
            }) => {
                if !simple {
                    return None;
                }

                let Expr::Name(ast::ExprName { id: target, .. }) = target.as_ref() else {
                    return None;
                };

                let target_name = target.to_string();
                let annotation_name = checker.generator().expr(annotation.as_ref());
                let is_dunder = target_name.starts_with("__") && target_name.ends_with("__");
                let is_private = target_name.starts_with('_') && !is_dunder;

                Some(ClassMember::Attribute(Attribute {
                    name: target_name,
                    type_annotation: annotation_name,
                    visibility: if is_private {
                        Visibility::Private
                    } else {
                        Visibility::Public
                    },
                }))
            }

            ast::Stmt::Assign(ast::StmtAssign { targets, value, .. }) => {
                // Handle simple assignments (like enum members)
                let value_type = match value.as_ref() {
                    Expr::BoolOp(_) | Expr::BooleanLiteral(_) => "bool",
                    Expr::BinOp(_) | Expr::UnaryOp(_) => "int",
                    Expr::Lambda(_) => "Callable",
                    Expr::DictComp(_) | Expr::Dict(_) => "dict",
                    Expr::Set(_) | Expr::SetComp(_) => "set",
                    Expr::FString(_) | Expr::StringLiteral(_) => "str",
                    Expr::NoneLiteral(_) => "None",
                    Expr::BytesLiteral(_) => "bytes",
                    Expr::EllipsisLiteral(_) => "...",
                    Expr::ListComp(_) | Expr::List(_) => "list",
                    Expr::Tuple(_) => "tuple",
                    Expr::NumberLiteral(inner) => match inner.value {
                        Number::Int(_) => "int",
                        Number::Float(_) => "float",
                        Number::Complex { .. } => "complex",
                    },
                    _ => "",
                };

                // For now, just handle the first target (typical for enums and simple assignments)
                if let Some(Expr::Name(ast::ExprName { id: target, .. })) = targets.first() {
                    let target_name = target.to_string();
                    let type_annotation = if value_type.is_empty() {
                        "Any"
                    } else {
                        value_type
                    }
                    .to_owned();

                    return Some(ClassMember::Attribute(Attribute {
                        name: target_name,
                        type_annotation,
                        visibility: Visibility::Public, // Simple assignments are always public
                    }));
                }

                None
            }

            ast::Stmt::FunctionDef(ast::StmtFunctionDef {
                name,
                is_async,
                parameters,
                returns,
                decorator_list,
                ..
            }) => {
                // Skip property setters and deleters - they're implementation details of the property
                if Self::is_property_setter_or_deleter(decorator_list, name.as_str()) {
                    return None;
                }

                let is_dunder = name.starts_with("__") && name.ends_with("__");
                let is_private = name.starts_with('_') && !is_dunder;
                let is_static = is_staticmethod(decorator_list, checker.semantic());

                // @property getters: show as attributes (read-only) instead of methods
                if is_property(
                    decorator_list,
                    std::iter::empty::<QualifiedName>(),
                    checker.semantic(),
                ) {
                    let return_type = returns.as_ref().map_or_else(
                        || {
                            simple_magic_return_type(name)
                                .map_or_else(|| "Any".to_owned(), String::from)
                        },
                        |target| checker.generator().expr(target.as_ref()),
                    );
                    return Some(ClassMember::Attribute(Attribute {
                        name: name.to_string(),
                        type_annotation: return_type,
                        visibility: if is_private {
                            Visibility::Private
                        } else {
                            Visibility::Public
                        },
                    }));
                }

                let mut param_gen = ParameterGenerator::new();
                param_gen.unparse_parameters(parameters);
                let params = param_gen.generate();

                let returns = returns
                    .as_ref()
                    .map(|target| checker.generator().expr(target.as_ref()))
                    .or_else(|| simple_magic_return_type(name).map(String::from));

                let mut decorators = vec![];
                if is_final(decorator_list, checker.semantic()) {
                    decorators.push("@final".to_string());
                }
                if is_classmethod(decorator_list, checker.semantic()) {
                    decorators.push("@classmethod".to_string());
                } else if is_static {
                    decorators.push("@staticmethod".to_string());
                }
                if is_overload(decorator_list, checker.semantic()) {
                    decorators.push("@overload".to_string());
                }
                if is_override(decorator_list, checker.semantic()) {
                    decorators.push("@override".to_string());
                }

                Some(ClassMember::Method(MethodSignature {
                    name: name.to_string(),
                    parameters: params,
                    return_type: returns,
                    visibility: if is_private {
                        Visibility::Private
                    } else {
                        Visibility::Public
                    },
                    is_static,
                    is_abstract: is_abstract(decorator_list, checker.semantic()),
                    is_async: *is_async,
                    decorators,
                }))
            }

            _ => None,
        }
    }

    fn classify_base(
        &self,
        checker: &Checker,
        detector: &ClassTypeDetector,
        base: &ast::Expr,
        class_is_enum: bool,
    ) -> BaseKind {
        // Enums are a special case: we don't draw inheritance relationships for enum bases.
        if class_is_enum {
            return BaseKind::Skip;
        }

        // Skip generic parameter carrier bases like Generic[T].
        if type_analyzer::extract_generic_params(base, checker).is_some() {
            return BaseKind::Skip;
        }

        if checker
            .semantic()
            .resolve_qualified_name(base)
            .is_some_and(|name| {
                matches!(name.segments(), ["typing", "Generic"])
                    || matches!(name.segments(), ["" | "builtins", "object"])
                    || matches!(name.segments(), ["abc", "ABC" | "ABCMeta"])
                    || matches!(
                        name.segments(),
                        ["typing" | "typing_extensions", "Protocol"]
                    )
            })
        {
            return BaseKind::Skip;
        }

        let base_name = checker.semantic().resolve_qualified_name(base).map_or_else(
            || {
                let name = checker.locator().slice(base);
                QualifiedName::user_defined(name).normalize_name()
            },
            |base_name| base_name.normalize_name(),
        );

        // Extract just the base class name without the generic specialization.
        let base_display = base_name
            .split('[')
            .next()
            .unwrap_or(&base_name)
            .trim_matches('`')
            .to_string();

        // Check if the base class is abstract or a protocol (either built-in or user-defined).
        let base_is_abstract_or_protocol = self.diagram.is_abstract_or_interface(&base_display)
            || detector.is_stdlib_abstract_or_protocol(base);

        BaseKind::InheritanceTarget {
            name: base_display,
            is_abstract_or_protocol: base_is_abstract_or_protocol,
        }
    }

    /// Add source code to the diagram (for stdin/WASM - uses Python defaults)
    pub fn add_source(&mut self, source: &str) {
        self.add_source_with_options(source, PySourceType::Python, ModuleKind::Module);
    }

    /// Add source code from a file path (infers source type and module kind)
    pub fn add_file(&mut self, source: &str, path: &Path) {
        let source_type = PySourceType::from(path);
        let module_kind = Self::module_kind_for_path(path);
        self.add_source_with_options(source, source_type, module_kind);
    }

    fn add_source_with_options(
        &mut self,
        source: &str,
        source_type: PySourceType,
        module_kind: ModuleKind,
    ) {
        let source_kind = SourceKind::Python {
            code: source.to_owned(),
            is_stub: false,
        };

        let parsed = Self::parse_python(source_kind.source_code(), source_type);
        let mut checker = Self::build_checker(
            &parsed.stylist,
            &parsed.locator,
            &parsed.python_ast,
            module_kind,
        );
        checker.see_imports(&parsed.python_ast);

        self.add_classes_from_ast(&checker, &parsed.python_ast);
    }

    fn add_classes_from_ast(&mut self, checker: &Checker, python_ast: &[ast::Stmt]) {
        for stmt in python_ast {
            if let ast::Stmt::ClassDef(class) = stmt {
                // we only care about class definitions
                self.add_class(checker, class, 1);
            }
        }
    }

    fn module_kind_for_path(path: &Path) -> ModuleKind {
        if path.ends_with("__init__.py") {
            ModuleKind::Package
        } else {
            ModuleKind::Module
        }
    }

    fn parse_python(source: &str, source_type: PySourceType) -> ParsedPython<'_> {
        let parsed = parse_unchecked_source(source, source_type);
        let stylist = Stylist::from_tokens(parsed.tokens(), source);
        let python_ast = parsed.into_suite().to_vec();

        ParsedPython {
            python_ast,
            locator: Locator::new(source),
            stylist,
        }
    }

    fn build_checker<'a>(
        stylist: &'a Stylist<'a>,
        locator: &'a Locator<'a>,
        python_ast: &'a [ast::Stmt],
        module_kind: ModuleKind,
    ) -> Checker<'a> {
        // Use a static dummy path for the semantic model (it's only used for diagnostics)
        static DUMMY_PATH: &str = "";
        let dummy = Path::new(DUMMY_PATH);

        let module = Module {
            kind: module_kind,
            source: ModuleSource::File(dummy),
            python_ast,
            name: None,
        };
        let semantic = SemanticModel::new(&[], dummy, module);
        Checker::new(stylist, locator, semantic)
    }
}

struct ParsedPython<'a> {
    python_ast: Vec<ast::Stmt>,
    locator: Locator<'a>,
    stylist: Stylist<'a>,
}

#[cfg(test)]
mod tests;