shell-memory-cli/shell_memory/scripts.py

"""Natural language to shell script generation."""

from typing import List, Optional, Dict
from dataclasses import dataclass

from .models import ScriptTemplate
from .database import Database


@dataclass
class GeneratedScript:
    """Result of script generation."""
    script: str
    confidence: float
    matched_keywords: List[str]
    template_used: Optional[ScriptTemplate]


class ScriptGenerator:
    """Generates shell scripts from natural language descriptions."""

    def __init__(self, db: Database):
        self.db = db
        self._ensure_default_templates()

    def _ensure_default_templates(self):
        existing = self.db.get_script_templates()
        if not existing:
            default_templates = [
                ScriptTemplate(
                    keywords=["deploy", "app", "application", "production"],
                    template="#!/bin/bash\n\n# Deploy application\necho \"Deploying application...\"\nsource venv/bin/activate\npip install -r requirements.txt\npython manage.py migrate\npython manage.py collectstatic --noinput\nsudo systemctl restart gunicorn\necho \"Deployment complete!\"",
                    description="Deploy a Python web application to production",
                ),
                ScriptTemplate(
                    keywords=["git", "commit", "message", "push"],
                    template="#!/bin/bash\n\n# Commit and push changes\necho \"Committing and pushing changes...\"\ngit add .\necho \"Enter commit message:\"\nread -r msg\ngit commit -m \"$msg\"\ngit push origin main\necho \"Changes committed and pushed!\"",
                    description="Commit changes to git with message and push",
                ),
                ScriptTemplate(
                    keywords=["docker", "build", "image", "container"],
                    template="#!/bin/bash\n\n# Build and run Docker container\necho \"Building Docker image...\"\ndocker build -t ${IMAGE_NAME:-myapp} .\necho \"Running container...\"\ndocker run -d -p ${PORT:-8000}:8000 --name ${CONTAINER_NAME:-myapp-run} ${IMAGE_NAME:-myapp}\necho \"Container started!\"",
                    description="Build and run a Docker container",
                ),
                ScriptTemplate(
                    keywords=["backup", "database", "db", "sql"],
                    template="#!/bin/bash\n\n# Backup database\necho \"Creating database backup...\"\ndatetime=$(date +%Y%m%d_%H%M%S)\npg_dump -U ${DB_USER:-postgres} ${DB_NAME:-mydb} > backup_${DB_NAME:-mydb}_$datetime.sql\ngzip backup_${DB_NAME:-mydb}_$datetime.sql\necho \"Backup created: backup_${DB_NAME:-mydb}_$datetime.sql.gz\"",
                    description="Create a compressed database backup",
                ),
                ScriptTemplate(
                    keywords=["test", "pytest", "coverage"],
                    template="#!/bin/bash\n\n# Run tests with coverage\necho \"Running tests...\"\npytest --cov=./ --cov-report=term-missing --cov-report=html\n\necho \"Opening coverage report...\"\nopen htmlcov/index.html 2>/dev/null || xdg-open htmlcov/index.html 2>/dev/null || echo \"Coverage report generated in htmlcov/\"",
                    description="Run pytest with coverage reporting",
                ),
            ]

            for template in default_templates:
                self.db.add_script_template(template)

    def generate(self, description: str, custom_vars: Optional[Dict[str, str]] = None) -> GeneratedScript:
        description_lower = description.lower()
        words = description_lower.split()

        templates = self.db.get_script_templates()

        best_match = None
        best_score = 0
        matched_keywords = []

        for template in templates:
            score = 0
            template_matched = []
            for keyword in template.keywords:
                if keyword in description_lower or keyword in words:
                    score += 1
                    template_matched.append(keyword)

            if score > best_score:
                best_score = score
                best_match = template
                matched_keywords = template_matched

        if best_match and best_score > 0:
            script = best_match.template
            if custom_vars:
                for key, value in custom_vars.items():
                    script = script.replace(f"${{{key}}}", value)
                    script = script.replace(f"${key}", value)

            confidence = best_score / max(len(best_match.keywords), 1)

            return GeneratedScript(
                script=script,
                confidence=confidence,
                matched_keywords=matched_keywords,
                template_used=best_match,
            )

        simple_script = self._generate_simple_script(description)
        return GeneratedScript(
            script=simple_script,
            confidence=0.3,
            matched_keywords=[],
            template_used=None,
        )

    def _generate_simple_script(self, description: str) -> str:
        words = description.lower().split()
        script_lines = ["#!/bin/bash", "", f"# Generated from: {description}", ""]

        if "docker" in words or "container" in words:
            script_lines.extend([
                "# Docker operations",
                "docker ps",
                "docker images",
                "echo 'Docker commands ready'",
            ])
        elif "git" in words or "commit" in words:
            script_lines.extend([
                "# Git operations",
                "git status",
                "git add .",
                "echo 'Enter commit message:'",
                "read -r msg",
                "git commit -m \"$msg\"",
                "git push",
            ])
        elif "test" in words:
            script_lines.extend([
                "# Test operations",
                "pytest -v",
                "echo 'Tests completed'",
            ])
        elif "backup" in words:
            script_lines.extend([
                "# Backup operations",
                "echo 'Creating backup...'",
                f"BACKUP_FILE=backup_$(date +%Y%m%d_%H%M%S).tar.gz",
                "tar -czf $BACKUP_FILE .",
                "echo 'Backup created: $BACKUP_FILE'",
            ])
        else:
            script_lines.extend([
                "# Script generated from your description",
                f"# Description: {description}",
                "",
                "echo 'Running: {description}'",
                "",
                "# TODO: Customize this script based on your needs",
                "# Add your commands here",
            ])

        return "\n".join(script_lines)

    def list_templates(self) -> List[ScriptTemplate]:
        return self.db.get_script_templates()

    def add_template(self, keywords: List[str], template: str, description: str = "") -> Optional[int]:
        script_template = ScriptTemplate(
            keywords=keywords,
            template=template,
            description=description,
        )
        return self.db.add_script_template(script_template)