# Copyright (c) Meta Platforms, Inc. and affiliates.

# pyre-strict

"""
MultiFileCallGraphAnalyzer: Orchestrates multi-file static call graph analysis.

This module provides the main analyzer that coordinates traversal across multiple
Python files, following imports to extract all transitively-called functions.
"""

import argparse
import ast
import json
import logging
import tempfile
from pathlib import Path
from typing import List, Optional, Set

from tritonparse.reproducer.ast_analyzer import CallGraph, Edge
from tritonparse.reproducer.consolidated_result import AnalysisStats, ConsolidatedResult
from tritonparse.reproducer.import_info import ImportInfo
from tritonparse.reproducer.import_parser import ImportParser
from tritonparse.reproducer.import_resolver import ImportResolver
from tritonparse.tp_logger import get_logger

logger = get_logger("multi_file_analyzer")


def _auto_detect_code_root(entry_file: str) -> str:
    """
    Auto-detect code root from entry file path.

    Walks up the directory tree from the entry file until it finds a common
    code root indicator. Currently supports:
    - "fbcode" directories (Meta's monorepo structure)
    - Directories containing "setup.py", "pyproject.toml" (Python projects)
    - Git repository roots (directories containing ".git")

    Args:
        entry_file: Absolute path to the entry file

    Returns:
        Absolute path to the code root directory

    Raises:
        ValueError: If code root cannot be detected from the entry file path
    """
    path = Path(entry_file).resolve()
    logger.debug("Auto-detecting code root from: %s", path)

    # Walk up the directory tree
    for parent in path.parents:
        # Check for fbcode (Meta's monorepo)
        if parent.name == "fbcode" or parent.name == "fbsource":
            logger.info("Auto-detected code root (fbcode): %s", parent)
            return str(parent)

        # Check for Python project markers
        if (parent / "setup.py").exists() or (parent / "pyproject.toml").exists():
            logger.info("Auto-detected code root (Python project): %s", parent)
            return str(parent)

        # Check for Git repository root
        if (parent / ".git").exists():
            logger.info("Auto-detected code root (Git repository): %s", parent)
            return str(parent)

    raise ValueError(
        f"Could not auto-detect code root from entry file: {entry_file}. "
        "Please specify --code-root explicitly."
    )


class MultiFileCallGraphAnalyzer:
    """
    Multi-file static call graph analyzer for creating reproducers.

    This analyzer orchestrates the analysis across multiple Python files by:
    1. Starting with an entry file and function
    2. Using CallGraph to extract dependencies within each file
    3. Following imports to analyze dependent files
    4. Consolidating results from all analyzed files
    """

    def __init__(
        self,
        entry_file: str,
        entry_function: str,
        code_roots: Optional[str] = None,
    ) -> None:
        """
        Initialize multi-file analyzer.

        The analyzer automatically computes the qualified backend name from
        entry_file + entry_function, removing the need for manual backend specification.

        Args:
            entry_file: Path to file containing entry function
            entry_function: Name of entry function (short name, e.g., "main_kernel")
            code_roots: Absolute path to default code root directory (auto-detected if None)

        Example:
            >>> analyzer = MultiFileCallGraphAnalyzer(
            ...     entry_file="/path/to/kernel.py",
            ...     entry_function="main_kernel",
            ...     code_roots="/path/to/fbcode",  # Optional
            ... )
        """
        self.entry_file = entry_file
        self.entry_function = entry_function

        # Auto-detect code_roots if not provided
        if code_roots is None:
            code_roots = _auto_detect_code_root(entry_file)
        self.code_roots = code_roots

        # Track multiple code roots for files from different projects
        self.file_to_code_root: dict[str, str] = {entry_file: code_roots}

        entry_module = self._file_to_module_name(entry_file)
        qualified_backend = f"{entry_module}.{entry_function}"
        self.backends = [entry_function]
        self.qualified_backend = qualified_backend

        self.import_resolver = ImportResolver(code_roots)
        self.import_parser = ImportParser(self.import_resolver)

        self.visited_files: Set[str] = set()
        self.pending_files: list[str] = []
        self.file_analyzers: dict[str, CallGraph] = {}
        self.all_imports: dict[str, list[ImportInfo]] = {}
        self.used_imports: dict[str, list[ImportInfo]] = {}

    def analyze(self) -> ConsolidatedResult:
        """
        Perform multi-file analysis.

        This is the main entry point that:
        1. Analyzes entry file
        2. Recursively analyzes imported files (breadth-first)
        3. Consolidates results

        Returns:
            ConsolidatedResult with functions, imports, and statistics
        """
        self._analyze_file(self.entry_file)

        while self.pending_files:
            pending_item = self.pending_files.pop(0)

            if isinstance(pending_item, tuple):
                file_path, backends_for_file = pending_item
            else:
                file_path = pending_item
                backends_for_file = None

            if file_path not in self.visited_files:
                self._analyze_file(file_path, backends_for_file)

        return self._consolidate_results()

    def _analyze_file(
        self, file_path: str, backends_for_file: Optional[list[str]] = None
    ) -> None:
        """
        Analyze a single file with CallGraph.

        This method:
        1. Marks file as visited
        2. Creates CallGraph analyzer for the file
        3. Extracts dependent functions
        4. Parses imports from AST
        5. Identifies which imports are used by dependent functions
        6. Adds imported files to pending_files queue

        Args:
            file_path: Absolute path to the Python file to analyze
            backends_for_file: Specific backends for this file (defaults to self.backends for entry file)
        """
        self.visited_files.add(file_path)
        logger.info("Analyzing %s", file_path)

        with open(file_path) as f:
            source_code = f.read()
        tree = ast.parse(source_code, filename=file_path)

        file_backends = (
            backends_for_file if backends_for_file is not None else self.backends
        )

        module_name = self._file_to_module_name(file_path)
        analyzer = CallGraph(
            filename=file_path,
            module_name=module_name,
            backends=file_backends,
            transitive_closure=True,
        )
        analyzer.visit(tree)

        self.file_analyzers[file_path] = analyzer

        module_name = self._file_to_module_name(file_path)
        package_name = (
            ".".join(module_name.split(".")[:-1]) if "." in module_name else None
        )
        imports = self.import_parser.parse_imports(
            tree, file_path, package=package_name
        )
        self.all_imports[file_path] = imports

        used_imports_list = self._identify_used_imports(imports, analyzer)
        self.used_imports[file_path] = used_imports_list

        imports_by_file: dict[str, list[str]] = {}
        for import_info in used_imports_list:
            if (
                import_info.resolved_path
                and not import_info.is_external
                and import_info.resolved_path not in self.visited_files
            ):
                if import_info.resolved_path not in imports_by_file:
                    imports_by_file[import_info.resolved_path] = []
                imports_by_file[import_info.resolved_path].extend(import_info.names)

        if imports_by_file:
            logger.info(
                "Found %d internal import file(s) to analyze from %s",
                len(imports_by_file),
                file_path,
            )
            for resolved_path, imported_names in imports_by_file.items():
                logger.debug(
                    "  → Will analyze %s for functions: %s",
                    resolved_path,
                    imported_names,
                )
        else:
            logger.debug(
                "No internal imports found in %s (all imports are external or already visited)",
                file_path,
            )

        for resolved_path, imported_names in imports_by_file.items():
            already_pending = any(
                (
                    pending_item[0] == resolved_path
                    if isinstance(pending_item, tuple)
                    else pending_item == resolved_path
                )
                for pending_item in self.pending_files
            )
            if not already_pending:
                unique_names = list(dict.fromkeys(imported_names))
                self.pending_files.append((resolved_path, unique_names))
                logger.info(
                    "Added %s to pending analysis queue with backends: %s",
                    resolved_path,
                    unique_names,
                )

    def _identify_used_imports(
        self,
        imports: list[ImportInfo],
        analyzer: CallGraph,
    ) -> list[ImportInfo]:
        """
        Identify which imports are actually used by dependent functions.

        Strategy:
        1. Get all callees from dependent functions (from call graph edges)
        2. Match callees to import statements
        3. Return only imports that are actually used

        Args:
            imports: All imports in the file
            analyzer: CallGraph analyzer for this file

        Returns:
            List of ImportInfo objects for imports that are used
        """
        used_symbols = self._extract_used_symbols(analyzer)
        logger.debug("Used symbols: %s", used_symbols)

        matching_imports = self._find_matching_imports(imports, used_symbols)
        import_groups = self._group_imports(matching_imports)
        used_imports_list = self._select_best_imports(import_groups)

        return used_imports_list

    def _extract_used_symbols(self, analyzer: CallGraph) -> Set[str]:
        """Extract all symbols used in the call graph."""
        used_symbols: Set[str] = set()
        for edge in analyzer.edges:
            used_symbols.add(edge.callee)
            callee_parts = edge.callee.split(".")
            if callee_parts:
                used_symbols.add(callee_parts[0])
        return used_symbols

    def _find_matching_imports(
        self, imports: list[ImportInfo], used_symbols: Set[str]
    ) -> list[tuple[ImportInfo, bool]]:
        """Find imports that match the used symbols."""
        matching_imports: list[tuple[ImportInfo, bool]] = []

        for import_info in imports:
            logger.debug(
                "Checking import: %s %s -> %s (external: %s)",
                import_info.import_type,
                import_info.module,
                import_info.names,
                import_info.is_external,
            )

            is_used = self._is_import_used(import_info, used_symbols)
            if is_used:
                is_internal_match = not import_info.is_external
                matching_imports.append((import_info, is_internal_match))

        return matching_imports

    def _is_import_used(self, import_info: ImportInfo, used_symbols: Set[str]) -> bool:
        """Check if an import is used based on the used symbols."""
        if self._matches_qualified_name(import_info, used_symbols):
            return True
        if self._matches_short_name(import_info, used_symbols):
            return True
        if self._matches_name(import_info, used_symbols):
            return True
        if self._matches_alias(import_info, used_symbols):
            return True
        if self._matches_module_prefix(import_info, used_symbols):
            return True
        return False

    def _matches_qualified_name(
        self, import_info: ImportInfo, used_symbols: Set[str]
    ) -> bool:
        """Check if import matches on qualified name."""
        if import_info.import_type != "from_import" or not import_info.module:
            return False

        for name in import_info.names:
            qualified_name = f"{import_info.module}.{name}"
            for symbol in used_symbols:
                if symbol == qualified_name or symbol.startswith(qualified_name + "."):
                    logger.debug(
                        "  Matched on qualified name: %s == %s",
                        qualified_name,
                        symbol,
                    )
                    return True
        return False

    def _matches_short_name(
        self, import_info: ImportInfo, used_symbols: Set[str]
    ) -> bool:
        """Check if import matches on short module name."""
        if import_info.import_type != "from_import" or not import_info.module:
            return False

        module_short_name = import_info.module.split(".")[-1]
        for name in import_info.names:
            short_qualified = f"{module_short_name}.{name}"
            if short_qualified in used_symbols:
                logger.debug(
                    "  Matched on short name: %s in used_symbols",
                    short_qualified,
                )
                return True
        return False

    def _matches_name(self, import_info: ImportInfo, used_symbols: Set[str]) -> bool:
        """Check if import matches on name."""
        for name in import_info.names:
            if name in used_symbols:
                logger.debug("  Matched on name: %s in used_symbols", name)
                return True
        return False

    def _matches_alias(self, import_info: ImportInfo, used_symbols: Set[str]) -> bool:
        """Check if import matches on alias."""
        for alias in import_info.aliases:
            if alias in used_symbols:
                logger.debug("  Matched on alias: %s in used_symbols", alias)
                return True
        return False

    def _matches_module_prefix(
        self, import_info: ImportInfo, used_symbols: Set[str]
    ) -> bool:
        """Check if import matches on module prefix."""
        if import_info.import_type != "import":
            return False

        module_prefix = import_info.module + "."
        for symbol in used_symbols:
            if symbol.startswith(module_prefix):
                logger.debug("  Matched on module prefix: %s", module_prefix)
                return True
        return False

    def _group_imports(
        self, matching_imports: list[tuple[ImportInfo, bool]]
    ) -> dict[tuple[str, tuple[str, ...]], list[tuple[ImportInfo, bool]]]:
        """Group imports by module and names."""
        import_groups: dict[
            tuple[str, tuple[str, ...]], list[tuple[ImportInfo, bool]]
        ] = {}

        for import_info, is_internal_match in matching_imports:
            module_short = (
                import_info.module.split(".")[-1] if import_info.module else ""
            )
            names_key = tuple(sorted(import_info.names))
            key = (module_short, names_key)

            if key not in import_groups:
                import_groups[key] = []
            import_groups[key].append((import_info, is_internal_match))

        return import_groups

    def _select_best_imports(
        self,
        import_groups: dict[tuple[str, tuple[str, ...]], list[tuple[ImportInfo, bool]]],
    ) -> list[ImportInfo]:
        """Select the best import from each group."""
        used_imports_list: list[ImportInfo] = []

        for group_imports in import_groups.values():
            group_imports.sort(key=lambda x: (not x[1], x[0].is_external))
            used_imports_list.append(group_imports[0][0])

        return used_imports_list

    def _get_code_root_for_file(self, file_path: str) -> str:
        """
        Get or auto-detect the code root for a specific file.

        This allows handling files from different project roots within the same analysis.

        Args:
            file_path: Absolute path to Python file

        Returns:
            Code root for this file (from cache or auto-detected)
        """
        if file_path in self.file_to_code_root:
            return self.file_to_code_root[file_path]

        # Auto-detect code root for this file
        try:
            code_root = _auto_detect_code_root(file_path)
            self.file_to_code_root[file_path] = code_root
            logger.debug("Detected code root %s for file %s", code_root, file_path)
            return code_root
        except ValueError:
            # Fall back to the analyzer's default code root
            logger.warning(
                "Could not auto-detect code root for %s, using default: %s",
                file_path,
                self.code_roots,
            )
            self.file_to_code_root[file_path] = self.code_roots
            return self.code_roots

    def _file_to_module_name(self, file_path: str) -> str:
        """
        Convert file path to Python module name.

        Example:
            /data/users/wychi/fbsource/fbcode/pytorch/tritonparse/module.py
            -> pytorch.tritonparse.module

        Args:
            file_path: Absolute path to Python file

        Returns:
            Module name as a dotted string
        """
        code_root = self._get_code_root_for_file(file_path)
        code_root_path = Path(code_root)
        file = Path(file_path)

        try:
            rel_path = file.relative_to(code_root_path)
            module_path = str(rel_path).replace("/", ".").removesuffix(".py")
            return module_path
        except ValueError:
            # File is not under the detected code root
            # Fall back to using the file's stem as the module name
            logger.warning(
                "File %s is not under code root %s, using file stem as module name",
                file_path,
                code_root,
            )
            return file.stem

    def _consolidate_results(self) -> ConsolidatedResult:
        """
        Consolidate results from all file analyzers.

        This method:
        1. Collects all functions and their source code from all files
        2. Tracks function locations
        3. Collects and deduplicates imports
        4. Collects all call graph edges
        5. Builds statistics

        Returns:
            ConsolidatedResult with all analysis results
        """
        all_functions: dict[str, str] = {}
        function_to_file: dict[str, str] = {}
        all_edges: List[Edge] = []

        functions_to_extract = self._collect_functions_to_extract(all_edges)
        self._extract_function_sources(
            functions_to_extract, all_functions, function_to_file
        )
        all_imports_list = self._collect_all_imports()
        unique_imports = self._deduplicate_imports(all_imports_list)

        stats = AnalysisStats(
            total_files_analyzed=len(self.visited_files),
            total_functions_found=len(all_functions),
            total_imports=len(unique_imports),
            external_imports=sum(1 for imp in unique_imports if imp.is_external),
            internal_imports=sum(1 for imp in unique_imports if not imp.is_external),
        )

        return ConsolidatedResult(
            functions=all_functions,
            function_locations=function_to_file,
            function_short_names={
                qualified: qualified.split(".")[-1]
                for qualified in all_functions.keys()
            },
            imports=unique_imports,
            edges=all_edges,
            analyzed_files=self.visited_files.copy(),
            stats=stats,
        )

    def _collect_functions_to_extract(self, all_edges: List[Edge]) -> set[str]:
        """Collect all functions that need to be extracted from analyzers."""
        functions_to_extract: set[str] = set()

        for _file_path, analyzer in self.file_analyzers.items():
            dependent_funcs = analyzer.get_dependent_functions()
            functions_to_extract.update(dependent_funcs)

            for backend in analyzer.backends:
                for local_func in analyzer.local_functions:
                    if local_func.split(".")[-1] == backend:
                        # Skip the primary entry function - it's not a dependency
                        if local_func == self.qualified_backend:
                            logger.debug(
                                "Skipping entry function %s (not a dependency)",
                                local_func,
                            )
                            continue

                        logger.debug(
                            "Adding backend function %s from file %s (backend: %s)",
                            local_func,
                            _file_path,
                            backend,
                        )
                        functions_to_extract.add(local_func)

            all_edges.extend(analyzer.edges)

        return functions_to_extract

    def _extract_function_sources(
        self,
        functions_to_extract: set[str],
        all_functions: dict[str, str],
        function_to_file: dict[str, str],
    ) -> None:
        """Extract source code for all functions from file analyzers."""
        for file_path, analyzer in self.file_analyzers.items():
            for func_name in functions_to_extract:
                if func_name in analyzer.func_nodes:
                    self._extract_single_function_source(
                        func_name, file_path, analyzer, all_functions, function_to_file
                    )

    def _extract_single_function_source(
        self,
        func_name: str,
        file_path: str,
        analyzer: CallGraph,
        all_functions: dict[str, str],
        function_to_file: dict[str, str],
    ) -> None:
        """Extract source code for a single function with source location comment."""
        source_code_map = analyzer.get_dependent_functions_source_code()

        if func_name in source_code_map:
            # Source location comment is already added by get_dependent_functions_source_code
            all_functions[func_name] = source_code_map[func_name]
            function_to_file[func_name] = file_path
        else:
            self._extract_function_from_ast(
                func_name, file_path, analyzer, all_functions, function_to_file
            )

    def _extract_function_from_ast(
        self,
        func_name: str,
        file_path: str,
        analyzer: CallGraph,
        all_functions: dict[str, str],
        function_to_file: dict[str, str],
    ) -> None:
        """Extract function source code directly from AST node with source location comment.

        Uses ast.unparse() for proper indentation handling. This is the preferred
        AST-level approach as it automatically handles nested function definitions
        and always outputs code at column 0.

        Note: ast.unparse() does not preserve comments. The original comments
        in the source code will be lost in the extracted version.
        """
        node = analyzer.func_nodes[func_name]

        # Get line numbers for the source location comment
        if node.decorator_list:
            start_line = node.decorator_list[0].lineno
        else:
            start_line = node.lineno
        end_line = node.end_lineno

        # Use ast.unparse() to generate properly indented source code
        # This handles nested functions, class methods, etc. automatically
        func_source = ast.unparse(node)

        # Add source location comment
        if start_line is not None and end_line is not None:
            source_comment = f"# Source: {file_path}:{start_line}-{end_line}\n"
            all_functions[func_name] = source_comment + func_source
        else:
            all_functions[func_name] = func_source

        function_to_file[func_name] = file_path

    def _collect_all_imports(self) -> list[ImportInfo]:
        """Collect all imports from visited files."""
        all_imports_list: list[ImportInfo] = []
        for file_path in self.visited_files:
            if file_path in self.used_imports:
                all_imports_list.extend(self.used_imports[file_path])
        return all_imports_list

    def _deduplicate_imports(self, imports: list[ImportInfo]) -> list[ImportInfo]:
        """
        Deduplicate imports while preserving order.

        Merges duplicate imports:
        - from X import A
        - from X import B
        -> from X import A, B

        Args:
            imports: List of ImportInfo objects (may contain duplicates)

        Returns:
            Deduplicated list of ImportInfo objects
        """
        import_groups: dict[tuple[str, str], ImportInfo] = {}

        for imp in imports:
            key = (imp.import_type, imp.module)
            if key in import_groups:
                existing = import_groups[key]
                for name in imp.names:
                    if name not in existing.names:
                        existing.names.append(name)
                existing.aliases.update(imp.aliases)
            else:
                import_groups[key] = ImportInfo(
                    import_type=imp.import_type,
                    module=imp.module,
                    names=imp.names.copy(),
                    source_file=imp.source_file,
                    resolved_path=imp.resolved_path,
                    is_external=imp.is_external,
                    lineno=imp.lineno,
                    aliases=imp.aliases.copy(),
                    level=imp.level,
                )

        return list(import_groups.values())


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Multi-file call graph analyzer for Python code. "
        "Analyzes a Python function and its dependencies across multiple files, "
        "extracting all transitively-called functions and their imports."
    )
    parser.add_argument(
        "--entry-file",
        "-f",
        required=True,
        help="Path to the entry file containing the function to analyze",
    )
    parser.add_argument(
        "--entry-function",
        "-F",
        required=True,
        help="Name of the entry function to analyze",
    )
    parser.add_argument(
        "--code-roots",
        "-r",
        default="",
        help="Path to default code root directory (default: auto-detect from entry file path)",
    )
    parser.add_argument(
        "--output",
        "-o",
        default=None,
        help="Output JSON file path (default: creates a temp file in /tmp/)",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose logging",
    )

    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s: %(message)s",
    )

    # Auto-detect code_roots if not provided
    code_roots = args.code_roots
    if not code_roots:
        try:
            code_roots = _auto_detect_code_root(args.entry_file)
        except ValueError as e:
            logger.error("%s", e)
            exit(1)

    analyzer = MultiFileCallGraphAnalyzer(
        entry_file=args.entry_file,
        entry_function=args.entry_function,
        code_roots=code_roots,
    )

    logger.info("Starting analysis of %s in %s", args.entry_function, args.entry_file)
    result = analyzer.analyze()

    logger.info("Analysis complete:")
    logger.info("  Files analyzed: %d", result.stats.total_files_analyzed)
    logger.info("  Functions found: %d", result.stats.total_functions_found)
    logger.info("  Total imports: %d", result.stats.total_imports)
    logger.info("  External imports: %d", result.stats.external_imports)
    logger.info("  Internal imports: %d", result.stats.internal_imports)

    logger.info("\nDependent functions (short names):")
    for func_name in sorted(result.function_short_names.keys()):
        short_name = result.function_short_names[func_name]
        if short_name != args.entry_function:
            logger.info(
                "  - %s. code size: %d", short_name, len(result.functions[func_name])
            )

    output_data = {
        "entry_file": args.entry_file,
        "entry_function": args.entry_function,
        "qualified_backend": analyzer.qualified_backend,
        "stats": {
            "total_files_analyzed": result.stats.total_files_analyzed,
            "total_functions_found": result.stats.total_functions_found,
            "total_imports": result.stats.total_imports,
            "external_imports": result.stats.external_imports,
            "internal_imports": result.stats.internal_imports,
        },
        "analyzed_files": sorted(result.analyzed_files),
        "functions": {
            func_name: {
                "source": source,
                "file": result.function_locations.get(func_name, "unknown"),
                "short_name": result.function_short_names.get(func_name, func_name),
            }
            for func_name, source in result.functions.items()
        },
        "imports": [
            {
                "import_type": imp.import_type,
                "module": imp.module,
                "names": imp.names,
                "source_file": imp.source_file,
                "resolved_path": imp.resolved_path,
                "is_external": imp.is_external,
                "lineno": imp.lineno,
                "aliases": imp.aliases,
                "level": imp.level,
            }
            for imp in result.imports
        ],
        "edges": [
            {"caller": edge.caller, "callee": edge.callee} for edge in result.edges
        ],
    }

    if args.output:
        output_path = args.output
    else:
        # Create a temp file with a descriptive name
        import os

        temp_fd, temp_path = tempfile.mkstemp(
            suffix=".json",
            prefix=f"multi_file_analysis_{args.entry_function}_",
            dir="/tmp",
            text=True,
        )
        os.close(temp_fd)  # Close the file descriptor
        output_path = temp_path

    with open(output_path, "w") as f:
        json.dump(output_data, f, indent=2)
    logger.info("Detailed results written to: %s", output_path)
