python-sdk/pkg/hanzo-tools-code/hanzo_tools/code/code_tool.py

"""Unified code semantics tool for HIP-0300 architecture.

This module provides a single unified 'code' tool that handles all semantic code operations:
- parse: Parse source to AST (tree-sitter)
- serialize: AST back to text
- symbols: List symbols in file/scope
- definition: Go to definition (LSP)
- references: Find all references (LSP)
- transform: Pure codemod → Patch (no side effects)
- summarize: Compress Diff/Log/Report to summary

Following Unix philosophy: one tool for the Symbols + Structure axis.
All operations are PURE (no side effects) except where noted.

Effect lattice position: PURE
Representation: Text → AST → Patch
Scope: File → Package → Repo
"""

import difflib
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, ClassVar, Literal

from mcp.server.fastmcp import Context as MCPContext

from hanzo_tools.core import (
    BaseTool,
    InvalidParamsError,
    NotFoundError,
    ToolError,
    content_hash,
)

# Language detection by extension
LANG_MAP = {
    ".py": "python",
    ".js": "javascript",
    ".ts": "typescript",
    ".tsx": "tsx",
    ".jsx": "jsx",
    ".go": "go",
    ".rs": "rust",
    ".c": "c",
    ".cpp": "cpp",
    ".h": "c",
    ".hpp": "cpp",
    ".java": "java",
    ".rb": "ruby",
    ".php": "php",
    ".swift": "swift",
    ".kt": "kotlin",
    ".scala": "scala",
    ".cs": "c_sharp",
    ".lua": "lua",
    ".sh": "bash",
    ".bash": "bash",
    ".zsh": "bash",
    ".json": "json",
    ".yaml": "yaml",
    ".yml": "yaml",
    ".toml": "toml",
    ".md": "markdown",
    ".html": "html",
    ".css": "css",
    ".sql": "sql",
}


@dataclass
class TransformSpec:
    """Specification for code transformation."""

    kind: Literal["rename", "extract", "inline", "move", "codemod", "generate"]

    # For rename
    old_name: str | None = None
    new_name: str | None = None
    scope: Literal["file", "package", "repo"] = "file"

    # For codemod (pattern-based)
    match_pattern: str | None = None
    replace_template: str | None = None

    # For generate
    template: str | None = None
    context: dict[str, Any] = field(default_factory=dict)


class CodeTool(BaseTool):
    """Unified code semantics tool (HIP-0300).

    Handles all semantic code operations on a single axis:
    - parse: Parse source to AST
    - serialize: AST to text
    - symbols: List symbols
    - definition: Go to definition
    - references: Find references
    - transform: Pure codemod → Patch
    - summarize: Compress to summary

    All operations are PURE (safe to parallelize/cache).
    """

    name: ClassVar[str] = "code"
    VERSION: ClassVar[str] = "0.1.0"

    def __init__(self, cwd: str | None = None):
        super().__init__()
        self.cwd = cwd or os.getcwd()
        self._tree_sitter = None
        self._lsp_clients: dict[str, Any] = {}
        self._register_code_actions()

    @property
    def description(self) -> str:
        return """Unified code semantics tool (HIP-0300).

Actions:
- parse: Parse source code to AST (tree-sitter)
- serialize: Convert AST back to text
- symbols: List symbols in file/scope
- definition: Go to symbol definition (LSP)
- references: Find all references to symbol (LSP)
- transform: Pure codemod producing Patch (no side effects)
- summarize: Compress Diff/Log/Report to summary

All operations are PURE - safe to cache and parallelize.
"""

    def _detect_lang(self, path: str | None, text: str | None = None) -> str:
        """Detect language from path extension or content."""
        if path:
            ext = Path(path).suffix.lower()
            if ext in LANG_MAP:
                return LANG_MAP[ext]

        # Fallback: try to detect from shebang or content
        if text:
            first_line = text.split("\n", 1)[0]
            if first_line.startswith("#!"):
                if "python" in first_line:
                    return "python"
                elif "node" in first_line or "deno" in first_line:
                    return "javascript"
                elif "bash" in first_line or "sh" in first_line:
                    return "bash"

        return "unknown"

    def _get_tree_sitter(self):
        """Lazy-load tree-sitter."""
        if self._tree_sitter is None:
            # Try tree-sitter-language-pack first (works with tree-sitter 0.24+)
            try:
                import tree_sitter_language_pack
                self._tree_sitter = tree_sitter_language_pack
            except ImportError:
                try:
                    import tree_sitter_languages
                    self._tree_sitter = tree_sitter_languages
                except ImportError:
                    raise ToolError(
                        code="INTERNAL_ERROR",
                        message="tree-sitter-languages not installed. Run: pip install tree-sitter-language-pack",
                    )
        return self._tree_sitter

    def _parse_with_tree_sitter(self, text: str, lang: str) -> dict:
        """Parse text to AST using tree-sitter."""
        ts = self._get_tree_sitter()

        try:
            parser = ts.get_parser(lang)
            tree = parser.parse(text.encode())

            def node_to_dict(node) -> dict:
                result = {
                    "type": node.type,
                    "start": {"line": node.start_point[0], "col": node.start_point[1]},
                    "end": {"line": node.end_point[0], "col": node.end_point[1]},
                }
                if node.child_count > 0:
                    result["children"] = [node_to_dict(c) for c in node.children]
                else:
                    result["text"] = node.text.decode() if node.text else ""
                return result

            return {
                "root": node_to_dict(tree.root_node),
                "lang": lang,
                "errors": [
                    {
                        "line": n.start_point[0],
                        "col": n.start_point[1],
                        "message": "syntax error",
                    }
                    for n in tree.root_node.children
                    if n.type == "ERROR"
                ],
            }
        except Exception as e:
            raise ToolError(
                code="INTERNAL_ERROR",
                message=f"Parse error: {e}",
                details={"lang": lang},
            )

    def _extract_symbols(self, ast: dict, lang: str) -> list[dict]:
        """Extract symbols from AST."""
        symbols = []

        # Symbol node types by language
        symbol_types = {
            "python": ["function_definition", "class_definition", "assignment"],
            "javascript": [
                "function_declaration",
                "class_declaration",
                "variable_declaration",
                "lexical_declaration",
            ],
            "typescript": [
                "function_declaration",
                "class_declaration",
                "interface_declaration",
                "type_alias_declaration",
            ],
            "go": ["function_declaration", "method_declaration", "type_declaration"],
            "rust": [
                "function_item",
                "struct_item",
                "enum_item",
                "impl_item",
                "trait_item",
            ],
        }

        types_to_find = symbol_types.get(
            lang, ["function_definition", "class_definition"]
        )

        def walk(node: dict, parent_name: str = ""):
            node_type = node.get("type", "")

            if node_type in types_to_find:
                # Try to extract name from first identifier child
                name = None
                for child in node.get("children", []):
                    if child.get("type") == "identifier":
                        name = child.get("text", "")
                        break
                    elif child.get("type") == "name":
                        name = child.get("text", "")
                        break

                if name:
                    symbols.append(
                        {
                            "name": name,
                            "kind": node_type,
                            "range": {
                                "start": node["start"],
                                "end": node["end"],
                            },
                            "parent": parent_name or None,
                        }
                    )
                    parent_name = name

            for child in node.get("children", []):
                walk(child, parent_name)

        if "root" in ast:
            walk(ast["root"])

        return symbols

    def _generate_patch(self, original: str, modified: str, path: str = "file") -> str:
        """Generate unified diff patch."""
        original_lines = original.splitlines(keepends=True)
        modified_lines = modified.splitlines(keepends=True)

        diff = difflib.unified_diff(
            original_lines,
            modified_lines,
            fromfile=f"a/{path}",
            tofile=f"b/{path}",
        )

        return "".join(diff)

    def _apply_rename(self, text: str, old_name: str, new_name: str) -> tuple[str, int]:
        """Apply simple rename transformation. Returns (new_text, count)."""
        import re

        # Word-boundary rename to avoid partial matches
        pattern = rf"\b{re.escape(old_name)}\b"
        new_text, count = re.subn(pattern, new_name, text)
        return new_text, count

    def _apply_codemod(self, text: str, match: str, replace: str) -> tuple[str, int]:
        """Apply pattern-based codemod. Returns (new_text, count)."""
        import re

        try:
            new_text, count = re.subn(match, replace, text)
            return new_text, count
        except re.error as e:
            raise InvalidParamsError(
                f"Invalid regex pattern: {e}", param="match_pattern"
            )

    def _register_code_actions(self):
        """Register all code actions."""

        @self.action("parse", "Parse source code to AST")
        async def parse(
            ctx: MCPContext,
            path: str | None = None,
            text: str | None = None,
            lang: str | None = None,
        ) -> dict:
            """Parse source to AST using tree-sitter.

            Args:
                path: File path to parse
                text: Raw text to parse (alternative to path)
                lang: Language (auto-detected if not specified)

            Returns:
                AST structure with lang and errors

            Effect: PURE
            Cache: hash(text)
            """
            if path and not text:
                full_path = (
                    Path(path) if Path(path).is_absolute() else Path(self.cwd) / path
                )
                if not full_path.exists():
                    raise NotFoundError(f"File not found: {path}", uri=str(full_path))
                text = full_path.read_text()

            if not text:
                raise InvalidParamsError("Either path or text required", param="text")

            detected_lang = lang or self._detect_lang(path, text)

            ast = self._parse_with_tree_sitter(text, detected_lang)

            return {
                "ast": ast,
                "lang": detected_lang,
                "hash": content_hash(text),
                "source_map": {"path": path} if path else None,
            }

        @self.action("serialize", "Convert AST back to text")
        async def serialize(
            ctx: MCPContext,
            ast: dict,
            lang: str | None = None,
        ) -> dict:
            """Serialize AST back to text.

            Note: This is a simplified implementation that extracts
            leaf node text. Full round-trip requires CST preservation.

            Effect: PURE
            """

            def extract_text(node: dict) -> str:
                if "text" in node:
                    return node["text"]
                children_text = []
                for child in node.get("children", []):
                    children_text.append(extract_text(child))
                return " ".join(children_text)

            root = ast.get("root") or ast
            text = extract_text(root)

            return {
                "text": text,
                "lang": lang or ast.get("lang", "unknown"),
            }

        @self.action("symbols", "List symbols in file")
        async def symbols(
            ctx: MCPContext,
            path: str | None = None,
            text: str | None = None,
            ast: dict | None = None,
            lang: str | None = None,
        ) -> dict:
            """Extract symbols from source code.

            Args:
                path: File path
                text: Raw text (alternative)
                ast: Pre-parsed AST (alternative)
                lang: Language hint

            Returns:
                List of symbols with name, kind, range

            Effect: PURE
            """
            if ast is None:
                if path and not text:
                    full_path = (
                        Path(path)
                        if Path(path).is_absolute()
                        else Path(self.cwd) / path
                    )
                    if not full_path.exists():
                        raise NotFoundError(f"File not found: {path}")
                    text = full_path.read_text()

                if not text:
                    raise InvalidParamsError("path, text, or ast required")

                detected_lang = lang or self._detect_lang(path, text)
                parsed = self._parse_with_tree_sitter(text, detected_lang)
                ast = parsed
                lang = detected_lang

            symbols_list = self._extract_symbols(
                ast, lang or ast.get("lang", "unknown")
            )

            return {
                "symbols": symbols_list,
                "count": len(symbols_list),
                "path": path,
            }

        @self.action("definition", "Go to symbol definition")
        async def definition(
            ctx: MCPContext,
            path: str,
            line: int,
            col: int,
        ) -> dict:
            """Find definition of symbol at position.

            Uses LSP when available, falls back to AST search.

            Effect: PURE
            """
            # TODO: Integrate with LSP (gopls, pyright, etc.)
            # For now, return a stub indicating LSP integration needed
            return {
                "path": path,
                "position": {"line": line, "col": col},
                "definition": None,
                "note": "LSP integration required for full definition lookup",
            }

        @self.action("references", "Find all references to symbol")
        async def references(
            ctx: MCPContext,
            path: str,
            line: int,
            col: int,
            scope: str = "repo",
        ) -> dict:
            """Find all references to symbol at position.

            Uses LSP when available, falls back to text search.

            Effect: PURE
            Scope: file | package | repo
            """
            # TODO: Integrate with LSP
            return {
                "path": path,
                "position": {"line": line, "col": col},
                "references": [],
                "scope": scope,
                "note": "LSP integration required for full reference lookup",
            }

        @self.action("transform", "Generate patch from transformation spec")
        async def transform(
            ctx: MCPContext,
            path: str | None = None,
            text: str | None = None,
            kind: str = "rename",
            old_name: str | None = None,
            new_name: str | None = None,
            match_pattern: str | None = None,
            replace_template: str | None = None,
            scope: str = "file",
        ) -> dict:
            """Apply transformation and produce Patch (PURE - no side effects).

            This is the key "pure middle layer" operator:
            - Input: source + transform spec
            - Output: Patch as value (not applied)

            Kinds:
                rename: Rename symbol (old_name → new_name)
                codemod: Pattern-based replacement

            Effect: PURE
            """
            # Load text if path provided
            if path and not text:
                full_path = (
                    Path(path) if Path(path).is_absolute() else Path(self.cwd) / path
                )
                if not full_path.exists():
                    raise NotFoundError(f"File not found: {path}")
                text = full_path.read_text()

            if not text:
                raise InvalidParamsError("path or text required")

            original_hash = content_hash(text)
            original = text
            changes_count = 0

            if kind == "rename":
                if not old_name or not new_name:
                    raise InvalidParamsError("rename requires old_name and new_name")
                text, changes_count = self._apply_rename(text, old_name, new_name)

            elif kind == "codemod":
                if not match_pattern or replace_template is None:
                    raise InvalidParamsError(
                        "codemod requires match_pattern and replace_template"
                    )
                text, changes_count = self._apply_codemod(
                    text, match_pattern, replace_template
                )

            else:
                raise InvalidParamsError(
                    f"Unknown transform kind: {kind}",
                    param="kind",
                    expected="rename | codemod",
                )

            # Generate patch
            patch = self._generate_patch(original, text, path or "input")

            return {
                "patch": patch,
                "changes_count": changes_count,
                "base_hash": original_hash,
                "new_hash": content_hash(text) if changes_count > 0 else original_hash,
                "kind": kind,
                "report": {
                    "files_affected": 1 if changes_count > 0 else 0,
                    "lines_changed": patch.count("\n") if patch else 0,
                },
            }

        @self.action("summarize", "Compress diff/log/report to summary")
        async def summarize(
            ctx: MCPContext,
            diff: str | None = None,
            log: list[dict] | None = None,
            report: dict | None = None,
            text: str | None = None,
            max_length: int = 500,
        ) -> dict:
            """Summarize code artifacts for review.

            Takes one of: diff, log, report, or text
            Produces: summary, risks, next_actions

            Effect: PURE
            """
            summary_parts = []
            risks = []
            next_actions = []

            if diff:
                # Analyze diff
                lines = diff.split("\n")
                adds = sum(
                    1
                    for line in lines
                    if line.startswith("+") and not line.startswith("+++")
                )
                dels = sum(
                    1
                    for line in lines
                    if line.startswith("-") and not line.startswith("---")
                )
                files = [line[6:] for line in lines if line.startswith("+++ b/")]

                summary_parts.append(
                    f"Patch: +{adds}/-{dels} lines across {len(files)} file(s)"
                )

                if adds + dels > 100:
                    risks.append("Large change - review carefully")
                if any("test" in f.lower() for f in files):
                    next_actions.append("Run tests to verify")
                else:
                    risks.append("No test files modified")
                    next_actions.append("Consider adding tests")

            if log:
                summary_parts.append(f"History: {len(log)} commit(s)")
                if log:
                    latest = (
                        log[0] if isinstance(log[0], dict) else {"message": str(log[0])}
                    )
                    summary_parts.append(f"Latest: {latest.get('message', 'N/A')[:50]}")

            if report:
                if "pass" in report:
                    status = "PASS" if report["pass"] else "FAIL"
                    summary_parts.append(f"Status: {status}")
                    if not report["pass"]:
                        risks.append("Tests failing")
                        next_actions.append("Fix failing tests before merge")

            if text:
                # Simple text summary
                word_count = len(text.split())
                summary_parts.append(f"Text: {word_count} words")
                if word_count > 1000:
                    summary_parts.append("(truncated for summary)")

            return {
                "summary": (
                    " | ".join(summary_parts)
                    if summary_parts
                    else "No content to summarize"
                ),
                "risks": risks,
                "next_actions": next_actions,
            }

        # ── TS-parity actions ──────────────────────────────────────────

        @self.action("search_symbol", "Find symbols across project")
        async def search_symbol(
            ctx: MCPContext,
            query: str,
            path: str | None = None,
            max_results: int = 20,
        ) -> dict:
            import re
            search_dir = Path(path) if path else Path(self.cwd)
            results = []
            skip_dirs = {".git", "node_modules", "target", "dist", "__pycache__", ".venv"}
            for p in search_dir.rglob("*"):
                if len(results) >= max_results:
                    break
                if any(sd in p.parts for sd in skip_dirs) or not p.is_file() or p.suffix not in LANG_MAP:
                    continue
                try:
                    text = p.read_text(errors="ignore")
                    for i, line in enumerate(text.splitlines(), 1):
                        if query in line and re.search(r'\b(def|class|function|fn|struct|enum|interface|type|const|let|var|pub)\b', line):
                            results.append({"uri": str(p), "line": i, "text": line.strip(), "type": "definition"})
                            if len(results) >= max_results:
                                break
                except (OSError, UnicodeDecodeError):
                    continue
            return {"query": query, "results": results, "count": len(results)}

        @self.action("outline", "List symbols with import count and export flag")
        async def outline(ctx: MCPContext, path: str, text: str | None = None) -> dict:
            import re
            full_path = Path(path) if Path(path).is_absolute() else Path(self.cwd) / path
            if text is None:
                if not full_path.exists():
                    raise NotFoundError(f"File not found: {path}")
                text = full_path.read_text()
            lines = text.splitlines()
            syms = []
            imports = 0
            sym_re = re.compile(r'(?:export\s+)?(?:async\s+)?(?:function|class|interface|type|enum|const|let|var|def|fn|pub\s+fn|pub\s+struct|struct|impl)\s+(\w+)')
            imp_re = re.compile(r'^(?:import|from|require|use)\b')
            for i, line in enumerate(lines, 1):
                if imp_re.match(line.strip()):
                    imports += 1
                m = sym_re.search(line)
                if m:
                    kind_m = re.search(r'\b(class|interface|type|enum|function|const|struct|impl|fn|def)\b', line)
                    syms.append({"name": m.group(1), "kind": kind_m.group(1) if kind_m else "symbol", "line": i, "exported": line.strip().startswith(("export ", "pub "))})
            return {"uri": str(full_path), "symbols": syms, "imports": imports, "lines": len(lines)}

        @self.action("metrics", "Count files and lines by extension")
        async def metrics(ctx: MCPContext, path: str | None = None) -> dict:
            search_dir = Path(path) if path else Path(self.cwd)
            skip_dirs = {".git", "node_modules", "target", "dist", "__pycache__", ".venv"}
            by_ext: dict[str, dict[str, int]] = {}
            total_files = 0
            total_lines = 0
            for p in search_dir.rglob("*"):
                if any(sd in p.parts for sd in skip_dirs) or not p.is_file():
                    continue
                ext = p.suffix or "other"
                try:
                    lines = len(p.read_text(errors="ignore").splitlines())
                    if ext not in by_ext:
                        by_ext[ext] = {"files": 0, "lines": 0}
                    by_ext[ext]["files"] += 1
                    by_ext[ext]["lines"] += lines
                    total_files += 1
                    total_lines += lines
                except (OSError, UnicodeDecodeError):
                    continue
            return {"total_files": total_files, "total_lines": total_lines, "by_extension": by_ext}

        @self.action("exports", "Extract public exports from file")
        async def exports(ctx: MCPContext, path: str) -> dict:
            full_path = Path(path) if Path(path).is_absolute() else Path(self.cwd) / path
            if not full_path.exists():
                raise NotFoundError(f"File not found: {path}")
            text = full_path.read_text()
            export_lines = [line.strip() for line in text.splitlines() if line.strip().startswith(("export ", "pub ", "__all__"))]
            return {"uri": str(full_path), "exports": export_lines, "count": len(export_lines)}

        @self.action("types", "Find type definitions in file")
        async def types(ctx: MCPContext, path: str) -> dict:
            import re
            full_path = Path(path) if Path(path).is_absolute() else Path(self.cwd) / path
            if not full_path.exists():
                raise NotFoundError(f"File not found: {path}")
            text = full_path.read_text()
            type_defs = []
            type_re = re.compile(r'(?:interface|type|enum|struct|class)\s+(\w+)')
            for i, line in enumerate(text.splitlines(), 1):
                m = type_re.search(line)
                if m:
                    type_defs.append({"name": m.group(1), "line": i, "text": line.strip()})
            return {"uri": str(full_path), "types": type_defs, "count": len(type_defs)}

        @self.action("hierarchy", "Build class inheritance tree")
        async def hierarchy(ctx: MCPContext, query: str, path: str | None = None) -> dict:
            import re
            search_dir = Path(path) if path else Path(self.cwd)
            skip_dirs = {".git", "node_modules", "target", "dist", "__pycache__", ".venv"}
            classes: dict[str, list[str]] = {}
            class_re = re.compile(r'class\s+(\w+)(?:\s*\(([^)]+)\)|\s+extends\s+(\w+))?')
            for p in search_dir.rglob("*"):
                if any(sd in p.parts for sd in skip_dirs) or not p.is_file() or p.suffix not in LANG_MAP:
                    continue
                try:
                    for m in class_re.finditer(p.read_text(errors="ignore")):
                        name = m.group(1)
                        parent = m.group(3) or (m.group(2).split(",")[0].strip() if m.group(2) else None)
                        if name not in classes:
                            classes[name] = []
                        if parent and parent not in ("object", "Object"):
                            classes.setdefault(parent, []).append(name)
                except (OSError, UnicodeDecodeError):
                    continue
            def build(n: str, d: int = 0) -> str:
                out = "  " * d + n + "\n"
                for c in classes.get(n, []):
                    out += build(c, d + 1)
                return out
            return {"root": query, "tree": build(query), "children": classes.get(query, [])}

        @self.action("rename", "Rename symbol across files")
        async def rename(ctx: MCPContext, query: str, new_name: str, path: str | None = None) -> dict:
            import re
            search_dir = Path(path) if path else Path(self.cwd)
            skip_dirs = {".git", "node_modules", "target", "dist", "__pycache__", ".venv"}
            total_changes = 0
            changed = []
            word_re = re.compile(rf'\b{re.escape(query)}\b')
            for p in search_dir.rglob("*"):
                if any(sd in p.parts for sd in skip_dirs) or not p.is_file() or p.suffix not in LANG_MAP:
                    continue
                try:
                    text = p.read_text(errors="ignore")
                    if word_re.search(text):
                        count = len(word_re.findall(text))
                        p.write_text(word_re.sub(new_name, text))
                        total_changes += count
                        changed.append(f"{p}: {count} replacements")
                except (OSError, UnicodeDecodeError):
                    continue
            return {"old_name": query, "new_name": new_name, "files_changed": len(changed), "total_replacements": total_changes, "changed": changed}

        @self.action("grep_replace", "Pattern replacement across codebase")
        async def grep_replace(ctx: MCPContext, query: str, replacement: str, path: str | None = None) -> dict:
            import re as re_mod
            search_dir = Path(path) if path else Path(self.cwd)
            skip_dirs = {".git", "node_modules", "target", "dist", "__pycache__", ".venv"}
            total_changes = 0
            changed = []
            try:
                regex = re_mod.compile(query)
            except re_mod.error as e:
                raise InvalidParamsError(f"Invalid regex: {e}", param="query")
            for p in search_dir.rglob("*"):
                if any(sd in p.parts for sd in skip_dirs) or not p.is_file() or p.suffix not in LANG_MAP:
                    continue
                try:
                    text = p.read_text(errors="ignore")
                    if regex.search(text):
                        count = len(regex.findall(text))
                        p.write_text(regex.sub(replacement, text))
                        total_changes += count
                        changed.append(f"{p}: {count}")
                except (OSError, UnicodeDecodeError):
                    continue
            return {"pattern": query, "replacement": replacement, "files_changed": len(changed), "total_replacements": total_changes, "changed": changed}


# Singleton instance
code_tool = CodeTool