Build: The Database Class

import sqlite3
from pathlib import Path
from logger import get_logger

logger = get_logger(__name__)


class Database:
    """Manages the SQLite database for the Prompt Vault."""

    def __init__(self, db_path="prompt_vault.db"):
        self.db_path = Path(db_path)
        self.conn = sqlite3.connect(self.db_path)
        # Row factory lets you access columns by name (row['platform'])
        # instead of by index (row[1]).
        self.conn.row_factory = sqlite3.Row
        self._create_tables()
        logger.info(f"Database opened: {self.db_path}")

    def _create_tables(self):
        """Create the prompts table if it doesn't exist."""
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS prompts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                platform TEXT NOT NULL,
                shot TEXT NOT NULL,
                prompt TEXT NOT NULL,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        self.conn.commit()

    def add_prompt(self, platform, shot, prompt_text):
        """Insert a new prompt. Returns the new row's ID."""
        cursor = self.conn.execute(
            "INSERT INTO prompts (platform, shot, prompt) VALUES (?, ?, ?)",
            (platform, shot, prompt_text)
        )
        self.conn.commit()
        logger.info(f"Inserted prompt #{cursor.lastrowid}: [{platform}] {shot}")
        return cursor.lastrowid

    def get_all(self):
        """Return all prompts, newest first."""
        cursor = self.conn.execute(
            "SELECT * FROM prompts ORDER BY created_at DESC"
        )
        return cursor.fetchall()

    def get_by_platform(self, platform):
        """Return prompts for a specific platform."""
        cursor = self.conn.execute(
            "SELECT * FROM prompts WHERE platform = ? ORDER BY created_at DESC",
            (platform,)
        )
        return cursor.fetchall()

    def search(self, query):
        """Search prompts by keyword in shot or prompt text."""
        pattern = f"%{query}%"
        cursor = self.conn.execute(
            "SELECT * FROM prompts WHERE shot LIKE ? OR prompt LIKE ?",
            (pattern, pattern)
        )
        return cursor.fetchall()

    def delete(self, prompt_id):
        """Delete a prompt by ID. Returns True if something was deleted."""
        cursor = self.conn.execute(
            "DELETE FROM prompts WHERE id = ?",
            (prompt_id,)
        )
        self.conn.commit()
        deleted = cursor.rowcount > 0
        if deleted:
            logger.info(f"Deleted prompt #{prompt_id}")
        else:
            logger.warning(f"Delete failed: no prompt with id={prompt_id}")
        return deleted

    def count_by_platform(self):
        """Return a dict of platform → count."""
        cursor = self.conn.execute(
            "SELECT platform, COUNT(*) as count FROM prompts GROUP BY platform"
        )
        return {row["platform"]: row["count"] for row in cursor.fetchall()}

    def close(self):
        """Close the database connection."""
        self.conn.close()
        logger.info("Database closed")

from database import Database

db = Database("test.db")

# Add some prompts
id1 = db.add_prompt("Kling", "Golf swing", "Slow motion golf swing, golden hour")
id2 = db.add_prompt("Runway", "Ocean aerial", "Aerial tracking over turquoise ocean")
id3 = db.add_prompt("Kling", "Golf aerial", "Drone over golf course, cinematic")

print(f"Added IDs: {id1}, {id2}, {id3}")

# Get all
rows = db.get_all()
print(f"Total: {len(rows)}")
for row in rows:
    print(f"  #{row['id']} [{row['platform']}] {row['shot']}")

# Search
results = db.search("golf")
print(f"\nSearch 'golf': {len(results)} results")

# Platform counts
print(f"Counts: {db.count_by_platform()}")  # {'Kling': 2, 'Runway': 1}

# Delete
db.delete(id2)
print(f"\nAfter delete: {len(db.get_all())} rows")

db.close()