database.py

"""SQLite database module for music memories app."""

import sqlite3
from contextlib import contextmanager
from datetime import datetime
from typing import Optional

DATABASE_PATH = "./music_memories.db"


@contextmanager
def get_db_connection():
    """Get a database connection with row factory."""
    conn = sqlite3.connect(DATABASE_PATH)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
    finally:
        conn.close()


def init_database() -> None:
    """Initialize the SQLite database with all tables."""
    with get_db_connection() as conn:
        # Songs table
        conn.execute("""
            CREATE TABLE IF NOT EXISTS songs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title TEXT NOT NULL,
                artist TEXT NOT NULL,
                album TEXT,
                duration INTEGER,
                bpm INTEGER,
                energy_level INTEGER CHECK(energy_level BETWEEN 1 AND 10),
                audio_file_path TEXT,
                file_size INTEGER
            )
        """)
        
        # Users table
        conn.execute("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        
        # Playlists table
        conn.execute("""
            CREATE TABLE IF NOT EXISTS playlists (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                vibe_code TEXT,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        
        # Contexts table
        conn.execute("""
            CREATE TABLE IF NOT EXISTS contexts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                weather TEXT,
                time_of_day TEXT,
                location_type TEXT
            )
        """)
        
        # Memories table
        conn.execute("""
            CREATE TABLE IF NOT EXISTS memories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                description TEXT NOT NULL,
                date TEXT,
                song_id INTEGER,
                FOREIGN KEY (user_id) REFERENCES users(id),
                FOREIGN KEY (song_id) REFERENCES songs(id)
            )
        """)
        
        # Play history table (with composite primary key)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS play_history (
                user_id INTEGER NOT NULL,
                song_id INTEGER NOT NULL,
                played_at TEXT DEFAULT CURRENT_TIMESTAMP,
                context_id INTEGER,
                PRIMARY KEY (user_id, song_id, played_at),
                FOREIGN KEY (user_id) REFERENCES users(id),
                FOREIGN KEY (song_id) REFERENCES songs(id),
                FOREIGN KEY (context_id) REFERENCES contexts(id)
            )
        """)
        
        conn.commit()


# ============== SONGS ==============

def add_song(title: str, artist: str, album: str = None, duration: int = None,
             bpm: int = None, energy_level: int = None) -> dict:
    """Add a new song."""
    with get_db_connection() as conn:
        cursor = conn.execute(
            """INSERT INTO songs (title, artist, album, duration, bpm, energy_level) 
               VALUES (?, ?, ?, ?, ?, ?)""",
            (title, artist, album, duration, bpm, energy_level)
        )
        conn.commit()
        song_id = cursor.lastrowid
        return {"id": song_id, "title": title, "artist": artist, "album": album,
                "duration": duration, "bpm": bpm, "energy_level": energy_level}


def get_all_songs() -> list[dict]:
    """Get all songs."""
    with get_db_connection() as conn:
        cursor = conn.execute("SELECT * FROM songs")
        return [dict(row) for row in cursor.fetchall()]


def get_song_by_id(song_id: int) -> Optional[dict]:
    """Get a song by ID."""
    with get_db_connection() as conn:
        cursor = conn.execute("SELECT * FROM songs WHERE id = ?", (song_id,))
        row = cursor.fetchone()
        return dict(row) if row else None


def delete_song(song_id: int) -> bool:
    """Delete a song."""
    with get_db_connection() as conn:
        cursor = conn.execute("DELETE FROM songs WHERE id = ?", (song_id,))
        conn.commit()
        return cursor.rowcount > 0


def update_song_file(song_id: int, audio_file_path: str, file_size: int) -> bool:
    """Update song with MinIO file path and size."""
    with get_db_connection() as conn:
        cursor = conn.execute(
            """UPDATE songs SET audio_file_path = ?, file_size = ? WHERE id = ?""",
            (audio_file_path, file_size, song_id)
        )
        conn.commit()
        return cursor.rowcount > 0


# ============== USERS ==============

def add_user(name: str) -> dict:
    """Add a new user."""
    with get_db_connection() as conn:
        cursor = conn.execute("INSERT INTO users (name) VALUES (?)", (name,))
        conn.commit()
        user_id = cursor.lastrowid
        cursor = conn.execute("SELECT * FROM users WHERE id = ?", (user_id,))
        return dict(cursor.fetchone())


def get_all_users() -> list[dict]:
    """Get all users."""
    with get_db_connection() as conn:
        cursor = conn.execute("SELECT * FROM users")
        return [dict(row) for row in cursor.fetchall()]


def get_user_by_id(user_id: int) -> Optional[dict]:
    """Get a user by ID."""
    with get_db_connection() as conn:
        cursor = conn.execute("SELECT * FROM users WHERE id = ?", (user_id,))
        row = cursor.fetchone()
        return dict(row) if row else None


# ============== PLAYLISTS ==============

def add_playlist(name: str, vibe_code: str = None) -> dict:
    """Add a new playlist."""
    with get_db_connection() as conn:
        cursor = conn.execute(
            "INSERT INTO playlists (name, vibe_code) VALUES (?, ?)",
            (name, vibe_code)
        )
        conn.commit()
        playlist_id = cursor.lastrowid
        cursor = conn.execute("SELECT * FROM playlists WHERE id = ?", (playlist_id,))
        return dict(cursor.fetchone())


def get_all_playlists() -> list[dict]:
    """Get all playlists."""
    with get_db_connection() as conn:
        cursor = conn.execute("SELECT * FROM playlists")
        return [dict(row) for row in cursor.fetchall()]


def delete_playlist(playlist_id: int) -> bool:
    """Delete a playlist."""
    with get_db_connection() as conn:
        cursor = conn.execute("DELETE FROM playlists WHERE id = ?", (playlist_id,))
        conn.commit()
        return cursor.rowcount > 0


# ============== MEMORIES ==============

def add_memory(user_id: int, description: str, date: str = None, song_id: int = None) -> dict:
    """Add a new memory."""
    with get_db_connection() as conn:
        cursor = conn.execute(
            "INSERT INTO memories (user_id, description, date, song_id) VALUES (?, ?, ?, ?)",
            (user_id, description, date, song_id)
        )
        conn.commit()
        memory_id = cursor.lastrowid
        cursor = conn.execute("SELECT * FROM memories WHERE id = ?", (memory_id,))
        return dict(cursor.fetchone())


def get_all_memories() -> list[dict]:
    """Get all memories."""
    with get_db_connection() as conn:
        cursor = conn.execute("SELECT * FROM memories")
        return [dict(row) for row in cursor.fetchall()]


def get_memories_by_user(user_id: int) -> list[dict]:
    """Get memories for a specific user."""
    with get_db_connection() as conn:
        cursor = conn.execute("SELECT * FROM memories WHERE user_id = ?", (user_id,))
        return [dict(row) for row in cursor.fetchall()]


def delete_memory(memory_id: int) -> bool:
    """Delete a memory."""
    with get_db_connection() as conn:
        cursor = conn.execute("DELETE FROM memories WHERE id = ?", (memory_id,))
        conn.commit()
        return cursor.rowcount > 0


# ============== CONTEXTS ==============

def add_context(weather: str = None, time_of_day: str = None, location_type: str = None) -> dict:
    """Add a new context."""
    with get_db_connection() as conn:
        cursor = conn.execute(
            "INSERT INTO contexts (weather, time_of_day, location_type) VALUES (?, ?, ?)",
            (weather, time_of_day, location_type)
        )
        conn.commit()
        context_id = cursor.lastrowid
        cursor = conn.execute("SELECT * FROM contexts WHERE id = ?", (context_id,))
        return dict(cursor.fetchone())


def get_all_contexts() -> list[dict]:
    """Get all contexts."""
    with get_db_connection() as conn:
        cursor = conn.execute("SELECT * FROM contexts")
        return [dict(row) for row in cursor.fetchall()]


def delete_context(context_id: int) -> bool:
    """Delete a context."""
    with get_db_connection() as conn:
        cursor = conn.execute("DELETE FROM contexts WHERE id = ?", (context_id,))
        conn.commit()
        return cursor.rowcount > 0


# ============== PLAY HISTORY ==============

def add_play_history(user_id: int, song_id: int, context_id: int = None) -> dict:
    """Add a play history entry."""
    played_at = datetime.utcnow().isoformat()
    with get_db_connection() as conn:
        conn.execute(
            """INSERT INTO play_history (user_id, song_id, played_at, context_id) 
               VALUES (?, ?, ?, ?)""",
            (user_id, song_id, played_at, context_id)
        )
        conn.commit()
        return {"user_id": user_id, "song_id": song_id, "played_at": played_at, "context_id": context_id}


def get_play_history(user_id: int = None, limit: int = 50) -> list[dict]:
    """Get play history, optionally filtered by user."""
    with get_db_connection() as conn:
        if user_id:
            cursor = conn.execute(
                """SELECT ph.*, s.title, s.artist 
                   FROM play_history ph 
                   JOIN songs s ON ph.song_id = s.id 
                   WHERE ph.user_id = ?
                   ORDER BY ph.played_at DESC LIMIT ?""",
                (user_id, limit)
            )
        else:
            cursor = conn.execute(
                """SELECT ph.*, s.title, s.artist 
                   FROM play_history ph 
                   JOIN songs s ON ph.song_id = s.id 
                   ORDER BY ph.played_at DESC LIMIT ?""",
                (limit,)
            )
        return [dict(row) for row in cursor.fetchall()]