shell-command-generator-cli/app/shellgen/history.py

"""Command history and learning system using SQLite."""

import sqlite3
import os
from pathlib import Path
from typing import List, Optional, Dict, Any
from datetime import datetime


class HistoryManager:
    """Manage command history and user feedback using SQLite."""

    def __init__(self, path: str = "~/.shellgen/history.db"):
        """Initialize history manager.

        Args:
            path: Path to the history database.
        """
        self._path = os.path.expanduser(path)
        self._ensure_db_exists()

    def _ensure_db_exists(self) -> None:
        """Create database and tables if they don't exist."""
        db_dir = os.path.dirname(self._path)
        if db_dir and not os.path.exists(db_dir):
            os.makedirs(db_dir, exist_ok=True)

        conn = sqlite3.connect(self._path)
        cursor = conn.cursor()

        cursor.execute("""
            CREATE TABLE IF NOT EXISTS history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                prompt TEXT NOT NULL,
                command TEXT NOT NULL,
                shell TEXT NOT NULL,
                executed BOOLEAN DEFAULT FALSE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

        cursor.execute("""
            CREATE TABLE IF NOT EXISTS feedback (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                history_id INTEGER NOT NULL,
                corrected_command TEXT,
                user_feedback TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (history_id) REFERENCES history (id)
            )
        """)

        conn.commit()
        conn.close()

    def add_entry(
        self,
        prompt: str,
        command: str,
        shell: str,
        executed: bool = False,
    ) -> int:
        """Add a new entry to history.

        Args:
            prompt: The natural language description.
            command: The generated command.
            shell: The target shell.
            executed: Whether the command was executed.

        Returns:
            ID of the inserted entry.
        """
        conn = sqlite3.connect(self._path)
        cursor = conn.cursor()

        cursor.execute(
            """
            INSERT INTO history (prompt, command, shell, executed)
            VALUES (?, ?, ?, ?)
            """,
            (prompt, command, shell, executed),
        )

        entry_id = cursor.lastrowid or 0
        conn.commit()
        conn.close()

        return entry_id

    def get_history(self, limit: int = 20) -> List[Dict[str, Any]]:
        """Get command history entries.

        Args:
            limit: Maximum number of entries to return.

        Returns:
            List of history entry dictionaries.
        """
        conn = sqlite3.connect(self._path)
        cursor = conn.cursor()

        cursor.execute(
            """
            SELECT id, prompt, command, shell, executed, created_at
            FROM history
            ORDER BY created_at DESC
            LIMIT ?
            """,
            (limit,),
        )

        rows = cursor.fetchall()
        conn.close()

        return [
            {
                "id": row[0],
                "prompt": row[1],
                "command": row[2],
                "shell": row[3],
                "executed": bool(row[4]),
                "created_at": row[5],
            }
            for row in rows
        ]

    def add_feedback(
        self,
        entry_id: int,
        corrected_command: Optional[str] = None,
        feedback: Optional[str] = None,
    ) -> bool:
        """Add feedback for a history entry.

        Args:
            entry_id: ID of the history entry.
            corrected_command: The user's corrected command.
            feedback: Additional feedback text.

        Returns:
            True if feedback was added successfully.
        """
        conn = sqlite3.connect(self._path)
        cursor = conn.cursor()

        cursor.execute(
            """
            INSERT INTO feedback (history_id, corrected_command, user_feedback)
            VALUES (?, ?, ?)
            """,
            (entry_id, corrected_command, feedback),
        )

        success = cursor.rowcount > 0
        conn.commit()
        conn.close()

        return success

    def get_entry(self, entry_id: int) -> Optional[Dict[str, Any]]:
        """Get a specific history entry.

        Args:
            entry_id: ID of the entry.

        Returns:
            Entry dictionary or None if not found.
        """
        conn = sqlite3.connect(self._path)
        cursor = conn.cursor()

        cursor.execute(
            """
            SELECT id, prompt, command, shell, executed, created_at
            FROM history
            WHERE id = ?
            """,
            (entry_id,),
        )

        row = cursor.fetchone()
        conn.close()

        if row:
            return {
                "id": row[0],
                "prompt": row[1],
                "command": row[2],
                "shell": row[3],
                "executed": bool(row[4]),
                "created_at": row[5],
            }

        return None

    def get_corrections(self) -> List[Dict[str, Any]]:
        """Get all user corrections for learning.

        Returns:
            List of correction dictionaries.
        """
        conn = sqlite3.connect(self._path)
        cursor = conn.cursor()

        cursor.execute(
            """
            SELECT h.id, h.prompt, h.command, f.corrected_command, f.user_feedback
            FROM history h
            JOIN feedback f ON h.id = f.history_id
            WHERE f.corrected_command IS NOT NULL
            ORDER BY f.created_at DESC
            """
        )

        rows = cursor.fetchall()
        conn.close()

        return [
            {
                "history_id": row[0],
                "original_prompt": row[1],
                "original_command": row[2],
                "corrected_command": row[3],
                "user_feedback": row[4],
            }
            for row in rows
        ]

    def clear_history(self, keep_count: int = 100) -> int:
        """Clear old history entries.

        Args:
            keep_count: Number of entries to keep.

        Returns:
            Number of deleted entries.
        """
        conn = sqlite3.connect(self._path)
        cursor = conn.cursor()

        cursor.execute(
            """
            DELETE FROM history
            WHERE id NOT IN (
                SELECT id FROM history
                ORDER BY created_at DESC
                LIMIT ?
            )
            """,
            (keep_count,),
        )

        deleted = cursor.rowcount
        conn.commit()
        conn.close()

        return deleted

    def get_stats(self) -> Dict[str, int]:
        """Get statistics about the history.

        Returns:
            Dictionary with history statistics.
        """
        conn = sqlite3.connect(self._path)
        cursor = conn.cursor()

        cursor.execute("SELECT COUNT(*) FROM history")
        total_count = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM history WHERE executed = 1")
        executed_count = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM feedback")
        feedback_count = cursor.fetchone()[0]

        conn.close()

        return {
            "total_entries": total_count,
            "executed": executed_count,
            "feedback": feedback_count,
        }