# WORKSPACE.rbs
load("@rbs//agent/tools.rbs", "AGENT_TOOLS_LOADED")

native.define_agent_tool(
    name = "list_processes",
    description = "List running processes on the system",
    shell_cmd = "ps aux | head -20",
)

native.define_agent_tool(
    name = "search_code",
    description = "Search for a pattern in source files",
    attrs = {
        "pattern": attr.string(
            doc = "The regex pattern to search for",
            mandatory = True,
        ),
        "file_type": attr.string(
            doc = "File extension to search (e.g., 'py', 'java')",
            default = "",
        ),
    },
    shell_cmd = "rg '{pattern}' --type {file_type}",
)

def _analyze_file_impl(ctx):
    """Analyze a file and return statistics."""
    path = ctx.attr.path

    # Read file content
    content = ctx.file.read(path)
    if content == None:
        return ctx.output.error("File not found: " + path)

    # Calculate statistics
    lines = content.split("\n")
    line_count = len(lines)
    char_count = len(content)

    # Run additional analysis
    result = ctx.shell.capture(command = "wc -w " + path)
    word_count = result.strip().split()[0] if result else "0"

    return ctx.output.success(
        message = "File analyzed",
        data = {
            "path": path,
            "lines": line_count,
            "characters": char_count,
            "words": word_count,
        },
    )

native.define_agent_tool(
    name = "analyze_file",
    description = "Analyze a file and return statistics (line count, word count, etc.)",
    attrs = {
        "path": attr.string(
            doc = "Path to the file to analyze",
            mandatory = True,
        ),
    },
    implementation = _analyze_file_impl,
)

def _my_tool(ctx):
    path = ctx.attr.path           # String attribute
    count = ctx.attr.count         # Integer attribute
    files = ctx.attr.files         # List attribute
    options = ctx.attr.options     # Dict attribute

def _my_tool(ctx):
    content = ctx.file.read("path/to/file.txt")       # Read a file
    ctx.file.write("output.txt", "Hello World")        # Write a file
    exists = ctx.file.exists("config.json")            # Check existence
    files = ctx.file.glob("src/**/*.py")               # List files matching pattern

def _my_tool(ctx):
    entries = ctx.dir.list("src/")                     # List directory
    ctx.dir.create("output/reports")                   # Create directory
    exists = ctx.dir.exists("build/")                  # Check existence
    ctx.dir.delete("tmp/", recursive = True)           # Delete directory

def _my_tool(ctx):
    output = ctx.shell.capture(command = "git status --short")     # Capture output
    exit_code = ctx.shell.run(command = "make build")              # Run without capturing
    output = ctx.shell.capture(command = "npm test", cwd = "frontend/") # With working dir

def _my_tool(ctx):
    full_path = ctx.path.join("src", "main", "app.py")    # Join paths
    dir_name = ctx.path.dir("src/main/app.py")             # Get directory
    base = ctx.path.base("src/main/app.py")                # Get base name
    ext = ctx.path.ext("app.py")                           # Get extension
    abs_path = ctx.path.abs("relative/path")               # Absolute path

def _my_tool(ctx):
    data = ctx.json.parse('{"name": "test", "count": 42}')    # Parse JSON
    json_str = ctx.json.stringify({"result": "success"})       # Serialize
    pretty = ctx.json.stringify(data, indent = 2)              # Pretty print

def _my_tool(ctx):
    # Success with message and data
    return ctx.output.success(
        message = "Found 5 files",
        data = {"files": ["a.py", "b.py"], "count": 5},
    )

    # Error
    return ctx.output.error("File not found: " + path)

native.define_agent_prompt(
    name = "code_reviewer",
    system_prompt = """You are a code review expert. Your role is to:
1. Analyze code for bugs and issues
2. Suggest improvements and best practices
3. Check for security vulnerabilities
4. Ensure code follows project conventions

Always explain your reasoning and provide specific examples.""",
    tools = [
        "read_files",
        "code_search",
        "analyze_file",
    ],
)

native.define_agent_subagent(
    name = "security_auditor",
    description = "Performs security analysis of code changes",
    when_to_use = "When reviewing code for security vulnerabilities",
    system_prompt = """You are a security auditor. Analyze the provided code for:
1. SQL injection vulnerabilities
2. XSS attack vectors
3. Authentication bypass risks
4. Sensitive data exposure
5. Dependency vulnerabilities

Report findings with severity levels and remediation steps.""",
    tools = ["read_files", "code_search", "run_terminal_command"],
    read_only = True,
)

native.define_agent_teammate_role(
    name = "ml_engineer",
    description = "Machine learning specialist for model training and optimization",
    system_prompt = """You are an ML engineer specializing in:
- PyTorch and distributed training (DDP, FSDP)
- Model architecture design and optimization
- Training pipeline development
- Hyperparameter tuning and experiment tracking

Always validate training code with small samples before full runs.""",
    tools = ["read_files", "write_file", "run_terminal_command", "code_search"],
    capabilities = ["model_training", "data_pipelines", "optimization"],
)

rbs agent "Spawn an ml_engineer teammate to optimize the training pipeline"

native.define_agent_skill(
    name = "deploy",
    description = "Deploy the application to staging or production",
    usage = "/deploy [staging|production]",
    prompt = """Deploy the application. Steps:
1. Run all tests
2. Build the production binary
3. Push to the container registry
4. Apply Kubernetes manifests
5. Verify the deployment is healthy

Target environment: {args}""",
    tools = ["run_terminal_command", "read_files"],
)

Session Start
  ├─→ Load rules (.agent-rules/, .cursorrules, etc.)
  ├─→ Load general memory (.rbs/memory/general.md)
  └─→ Agent loop
        ├─→ read_files("src/api.py")
        │     → File content returned
        │     + .rbs/memory/src/api.py.md (auto-injected)
        │     + .rbs/memory/src/_package.md (auto-injected)
        ├─→ Agent discovers pattern → memory_save(...)
        ├─→ User corrects agent → agent saves correction
        └─→ Session ends → memory persists

.rbs/memory/
├── general.md                   # Project-wide notes
├── src/
│   └── api.py.md                # Notes about src/api.py
└── services/
    └── auth/
        ├── _package.md          # Notes about the auth package
        └── handler.py.md        # Notes about handler.py

# General Project Memory

## [m-abc12345] 2026-02-17 10:30
> Tags: architecture, convention

This project uses dependency injection throughout.
Never create singletons — pass dependencies through constructors.

## [m-def67890] 2026-02-17 14:15
> Tags: user-correction

User correction: Use interfaces for the repository pattern,
not concrete implementations.

## [m-user0001] 2026-02-17 16:00
> Tags: convention

We use dependency injection throughout the project.
Never create singletons or global state.

# WORKSPACE.rbs
load("@rbs//agent/tools.rbs", "AGENT_TOOLS_LOADED")

# Custom tool: Check test coverage
def _coverage_check_impl(ctx):
    result = ctx.shell.capture(command = "rbs test //... --coverage")
    if result == None:
        return ctx.output.error("Failed to run tests")

    return ctx.output.success(
        message = "Coverage report generated",
        data = {"output": result},
    )

native.define_agent_tool(
    name = "check_coverage",
    description = "Run tests and check code coverage",
    implementation = _coverage_check_impl,
)

# Custom tool: Generate changelog
native.define_agent_tool(
    name = "changelog",
    description = "Generate changelog from recent git commits",
    attrs = {
        "since": attr.string(
            doc = "Git ref to start from (e.g., 'v1.0.0', 'HEAD~10')",
            default = "HEAD~10",
        ),
    },
    shell_cmd = "git log --oneline {since}..HEAD",
)

# Custom prompt for this project
native.define_agent_prompt(
    name = "project_assistant",
    system_prompt = """You are an assistant for this codebase.

Key conventions:
- Write tests for all new code
- Use dependency injection
- Update CHANGELOG.md for changes

When making changes, always run tests first.""",
    tools = [
        "read_files",
        "write_file",
        "str_replace",
        "run_terminal_command",
        "check_coverage",
        "changelog",
    ],
)

# Custom teammate role
native.define_agent_teammate_role(
    name = "data_engineer",
    description = "Data pipeline and preprocessing specialist",
    system_prompt = "You build efficient data loading and preprocessing pipelines.",
    tools = ["read_files", "write_file", "run_terminal_command"],
    capabilities = ["data_loading", "preprocessing", "augmentation"],
)

# Custom skill
native.define_agent_skill(
    name = "release",
    description = "Create a release",
    usage = "/release [version]",
    prompt = "Create release {args}: run tests, update changelog, tag, build, and publish.",
    tools = ["run_terminal_command", "read_files", "write_file"],
)