brain/commands/persona.py

"""ever-brain persona — Manage agent identities and system prompts."""

import os
import shutil
import subprocess
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.table import Table

from brain.config import get_active_brain
from brain.paths import get_personas_dir, get_persona_path
from brain.utils.display import error, success, info

app = typer.Typer(help="Manage agent personas (identities).")
console = Console()

# Built-in personas (shipped with the app)
BUILTIN_PERSONAS_DIR = Path(__file__).resolve().parent.parent / "templates" / "personas"


@app.command("list")
def list_personas() -> None:
    """List all available personas (Built-in + Custom)."""
    table = Table(title="Available Personas")
    table.add_column("Type", style="cyan")
    table.add_column("Name", style="bold white")
    table.add_column("Status", style="green")

    # Built-in
    if BUILTIN_PERSONAS_DIR.exists():
        for p in BUILTIN_PERSONAS_DIR.glob("*.md"):
            table.add_row("Built-in", p.stem, "Ready")

    # Custom
    custom_dir = get_personas_dir()
    for p in custom_dir.glob("*.md"):
        table.add_row("Custom", p.stem, "Ready")

    console.print(table)


@app.command("set")
def set_persona(name: str = typer.Argument(..., help="Name of the persona to apply")) -> None:
    """Apply a persona to the currently active Ever Brain."""
    active = get_active_brain()
    if not active or not active.exists():
        error("No active brain set. Use 'ever-brain use <name>' first.")
        raise typer.Exit(1)

    # Resolve persona source
    source = BUILTIN_PERSONAS_DIR / f"{name}.md"
    if not source.exists():
        source = get_persona_path(name)
    
    if not source.exists():
        error(f"Persona '{name}' not found.")
        raise typer.Exit(1)

    # Apply to active brain
    dest = active / "agent-configs" / "persona.md"
    dest.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy(source, dest)

    # Update AGENT.md to point to persona.md if not already there
    agent_md = active / "AGENT.md"
    if agent_md.exists():
        content = agent_md.read_text(encoding="utf-8")
        if "## ACTIVE PERSONA" not in content:
            new_content = content + f"\n\n## ACTIVE PERSONA\nThis brain is currently operating as: **{name.capitalize()}**\nSee `agent-configs/persona.md` for specific instructions.\n"
            agent_md.write_text(new_content, encoding="utf-8")
        else:
            # Update the persona name in AGENT.md
            import re
            content = re.sub(r"operating as: \*\*(.*?)\*\*", f"operating as: **{name.capitalize()}**", content)
            agent_md.write_text(content, encoding="utf-8")

    success(f"Persona for '{active.name}' set to '{name}'.")


@app.command("create")
def create_persona(
    name: str = typer.Argument(..., help="Name of the new persona"),
    ai: bool = typer.Option(False, "--ai", help="Generate persona using AI (Pro feature)")
) -> None:
    """Create a new custom persona."""
    path = get_persona_path(name)
    if path.exists():
        error(f"Persona '{name}' already exists.")
        raise typer.Exit(1)

    if ai:
        info("AI Generation is a Pro feature. Checking status...")
        # Hook for AI generation
        info("Please describe the persona you want to build:")
        desc = typer.prompt("Description")
        info(f"Generating persona '{name}' based on: {desc}")
        # Placeholder for AI call
        content = f"# PERSONA: {name.capitalize()}\n\nGenerated via AI based on: {desc}\n\n## Instructions\n- Add your instructions here."
        path.write_text(content, encoding="utf-8")
        success(f"AI Persona '{name}' generated.")
    else:
        # Manual creation
        content = f"# PERSONA: {name.capitalize()}\n\n## Core Principles\n- Principle 1\n\n## Interaction Style\n- Style 1"
        path.write_text(content, encoding="utf-8")
        info(f"Template created at {path}")
        
        if typer.confirm("Would you like to edit it now?"):
            typer.edit(filename=str(path))
            success(f"Persona '{name}' saved.")


@app.command("edit")
def edit_persona(name: str = typer.Argument(..., help="Name of the persona to edit")) -> None:
    """Edit a custom persona."""
    path = get_persona_path(name)
    if not path.exists():
        error(f"Custom persona '{name}' not found. Note: Built-in personas cannot be edited directly.")
        raise typer.Exit(1)

    typer.edit(filename=str(path))
    success(f"Persona '{name}' updated.")


if __name__ == "__main__":
    app()