Part 1: The Death of 'Code Typists' — When Syntax is No Longer an Advantage

#!/usr/bin/env python3
"""
Production Tree-sitter AST Extraction & Complexity Analyzer
Parses Go/Python source code, walks syntax trees, extracts function signatures,
computes cyclomatic complexity, and generates structured prompt context for LLM agents.
"""

from __future__ import annotations

import json
import logging
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

import tree_sitter_go as tsgo
from tree_sitter import Language, Node, Parser

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("ASTAnalyzer")

GO_LANGUAGE = Language(tsgo.language())


@dataclass
class ParameterDef:
    name: str
    param_type: str


@dataclass
class FunctionSignature:
    name: str
    receiver: str | None
    parameters: list[ParameterDef]
    return_types: list[str]
    cyclomatic_complexity: int
    start_line: int
    end_line: int
    is_exported: bool
    docstring: str | None


@dataclass
class StructField:
    name: str
    field_type: str
    tag: str | None


@dataclass
class StructDefinition:
    name: str
    fields: list[StructField]
    start_line: int
    end_line: int
    is_exported: bool


@dataclass
class FileASTContext:
    file_path: str
    package_name: str
    imports: list[str] = field(default_factory=list)
    structs: list[StructDefinition] = field(default_factory=list)
    functions: list[FunctionSignature] = field(default_factory=list)
    total_complexity: int = 0


class GoASTContextExtractor:
    """Extracts architectural signatures and complexity metrics from Go source code."""

    BRANCHING_NODE_TYPES = {
        "if_statement",
        "for_statement",
        "type_switch_statement",
        "expression_switch_statement",
        "communication_case",
        "default_case",
        "expression_case",
    }

    def __init__(self) -> None:
        self.parser = Parser()
        self.parser.language = GO_LANGUAGE

    def parse_source(self, file_path: str, source_bytes: bytes) -> FileASTContext:
        tree = self.parser.parse(source_bytes)
        root = tree.root_node

        package_name = self._extract_package_name(root, source_bytes)
        imports = self._extract_imports(root, source_bytes)
        structs = self._extract_structs(root, source_bytes)
        functions = self._extract_functions(root, source_bytes)

        total_complexity = sum(fn.cyclomatic_complexity for fn in functions)

        return FileASTContext(
            file_path=file_path,
            package_name=package_name,
            imports=imports,
            structs=structs,
            functions=functions,
            total_complexity=total_complexity,
        )

    def _extract_package_name(self, root: Node, source: bytes) -> str:
        for child in root.children:
            if child.type == "package_clause":
                for sub in child.children:
                    if sub.type == "package_identifier":
                        return source[sub.start_byte : sub.end_byte].decode("utf-8")
        return "main"

    def _extract_imports(self, root: Node, source: bytes) -> list[str]:
        imports = []
        for child in root.children:
            if child.type == "import_declaration":
                for spec in child.children:
                    if spec.type == "import_spec":
                        path_node = spec.child_by_field_name("path")
                        if path_node:
                            raw = source[path_node.start_byte : path_node.end_byte].decode("utf-8")
                            imports.append(raw.strip('"'))
                    elif spec.type == "import_spec_list":
                        for sub_spec in spec.children:
                            if sub_spec.type == "import_spec":
                                path_node = sub_spec.child_by_field_name("path")
                                if path_node:
                                    raw = source[path_node.start_byte : path_node.end_byte].decode("utf-8")
                                    imports.append(raw.strip('"'))
        return imports

    def _extract_structs(self, root: Node, source: bytes) -> list[StructDefinition]:
        structs = []
        for child in root.children:
            if child.type == "type_declaration":
                for spec in child.children:
                    if spec.type == "type_spec":
                        name_node = spec.child_by_field_name("name")
                        type_node = spec.child_by_field_name("type")
                        if name_node and type_node and type_node.type == "struct_type":
                            struct_name = source[name_node.start_byte : name_node.end_byte].decode("utf-8")
                            fields = self._parse_struct_fields(type_node, source)
                            structs.append(
                                StructDefinition(
                                    name=struct_name,
                                    fields=fields,
                                    start_line=spec.start_point[0] + 1,
                                    end_line=spec.end_point[0] + 1,
                                    is_exported=struct_name[0].isupper(),
                                )
                            )
        return structs

    def _parse_struct_fields(self, struct_node: Node, source: bytes) -> list[StructField]:
        fields = []
        field_list = struct_node.child_by_field_name("fields")
        if not field_list:
            return fields

        for field_decl in field_list.children:
            if field_decl.type == "field_declaration":
                name_node = field_decl.child_by_field_name("name")
                type_node = field_decl.child_by_field_name("type")
                tag_node = field_decl.child_by_field_name("tag")

                fname = source[name_node.start_byte : name_node.end_byte].decode("utf-8") if name_node else "anonymous"
                ftype = source[type_node.start_byte : type_node.end_byte].decode("utf-8") if type_node else "unknown"
                ftag = source[tag_node.start_byte : tag_node.end_byte].decode("utf-8") if tag_node else None

                fields.append(StructField(name=fname, field_type=ftype, tag=ftag))
        return fields

    def _extract_functions(self, root: Node, source: bytes) -> list[FunctionSignature]:
        functions = []
        for child in root.children:
            if child.type in ("function_declaration", "method_declaration"):
                fn_sig = self._parse_function_node(child, source)
                if fn_sig:
                    functions.append(fn_sig)
        return functions

    def _parse_function_node(self, node: Node, source: bytes) -> FunctionSignature | None:
        name_node = node.child_by_field_name("name")
        if not name_node:
            return None

        fn_name = source[name_node.start_byte : name_node.end_byte].decode("utf-8")
        receiver = None
        recv_node = node.child_by_field_name("receiver")
        if recv_node:
            receiver = source[recv_node.start_byte : recv_node.end_byte].decode("utf-8")

        parameters = []
        params_node = node.child_by_field_name("parameters")
        if params_node:
            for p in params_node.children:
                if p.type == "parameter_declaration":
                    p_name_node = p.child_by_field_name("name")
                    p_type_node = p.child_by_field_name("type")
                    pname = source[p_name_node.start_byte : p_name_node.end_byte].decode("utf-8") if p_name_node else "_"
                    ptype = source[p_type_node.start_byte : p_type_node.end_byte].decode("utf-8") if p_type_node else ""
                    parameters.append(ParameterDef(name=pname, param_type=ptype))

        return_types = []
        result_node = node.child_by_field_name("result")
        if result_node:
            if result_node.type == "parameter_list":
                for r in result_node.children:
                    if r.type == "parameter_declaration":
                        rtype_node = r.child_by_field_name("type")
                        if rtype_node:
                            return_types.append(source[rtype_node.start_byte : rtype_node.end_byte].decode("utf-8"))
            else:
                return_types.append(source[result_node.start_byte : result_node.end_byte].decode("utf-8"))

        complexity = self._compute_cyclomatic_complexity(node, source)

        return FunctionSignature(
            name=fn_name,
            receiver=receiver,
            parameters=parameters,
            return_types=return_types,
            cyclomatic_complexity=complexity,
            start_line=node.start_point[0] + 1,
            end_line=node.end_point[0] + 1,
            is_exported=fn_name[0].isupper(),
            docstring=None,
        )

    def _compute_cyclomatic_complexity(self, node: Node, source: bytes) -> int:
        """Computes McCabe cyclomatic complexity: CC = E - N + 2P (approximated via branches + 1)."""
        complexity = 1

        def walk(n: Node):
            nonlocal complexity
            if n.type in self.BRANCHING_NODE_TYPES:
                complexity += 1
            elif n.type == "binary_expression":
                op_node = n.child_by_field_name("operator")
                if op_node:
                    op_text = source[op_node.start_byte : op_node.end_byte].decode("utf-8")
                    if op_text in ("&&", "||"):
                        complexity += 1
            for child in n.children:
                walk(child)

        walk(node)
        return complexity


def generate_llm_ast_prompt_payload(context: FileASTContext) -> str:
    """Serializes AST context into token-efficient JSON format for LLM agent prompts."""
    payload = {
        "file": context.file_path,
        "package": context.package_name,
        "imports": context.imports,
        "struct_contracts": [
            {
                "struct": s.name,
                "exported": s.is_exported,
                "fields": [f"{f.name}: {f.field_type}" for f in s.fields],
            }
            for s in context.structs
        ],
        "interface_signatures": [
            {
                "function": f.name,
                "receiver": f.receiver,
                "inputs": [f"{p.name}: {p.param_type}" for p in f.parameters],
                "outputs": f.return_types,
                "complexity": f.cyclomatic_complexity,
                "lines": f"{f.start_line}-{f.end_line}",
            }
            for f in context.functions
        ],
        "metrics": {
            "total_functions": len(context.functions),
            "aggregate_cyclomatic_complexity": context.total_complexity,
            "complexity_risk": "HIGH" if context.total_complexity > 20 else "NORMAL",
        },
    }
    return json.dumps(payload, indent=2)


if __name__ == "__main__":
    sample_go_code = b"""
    package paymentservice

    import (
        "context"
        "errors"
        "time"
    )

    type PaymentOrder struct {
        OrderID   string    `json:"order_id"`
        Amount    float64   `json:"amount"`
        Currency  string    `json:"currency"`
        CreatedAt time.Time `json:"created_at"`
    }

    type OrderProcessor struct {
        maxRetries int
    }

    func (p *OrderProcessor) ProcessTransaction(ctx context.Context, order *PaymentOrder) (string, error) {
        if order == nil {
            return "", errors.New("nil order")
        }
        if order.Amount <= 0 || order.Currency == "" {
            return "", errors.New("invalid transaction parameters")
        }

        for attempt := 0; attempt < p.maxRetries; attempt++ {
            select {
            case <-ctx.Done():
                return "", ctx.Err()
            default:
                if order.Amount > 10000.0 {
                    return "FLAGGED_FOR_MANUAL_REVIEW", nil
                }
                return "SETTLED_OK", nil
            }
        }
        return "RETRY_EXHAUSTED", nil
    }
    """

    extractor = GoASTContextExtractor()
    ctx = extractor.parse_source("payment_processor.go", sample_go_code)
    json_output = generate_llm_ast_prompt_payload(ctx)

    logger.info("Successfully extracted AST context with Tree-sitter:")
    print(json_output)