nfoguard/core/database.py

#!/usr/bin/env python3
"""
Database operations module for NFOGuard
"""
import sqlite3
from pathlib import Path
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple, Any

from core.logging import _log


class NFOGuardDatabase:
    """Database operations for TV series, episodes, movies, and their metadata"""
    
    # Database schemas
    TV_SCHEMA = """
    PRAGMA journal_mode=WAL;

    CREATE TABLE IF NOT EXISTS meta (
      key TEXT PRIMARY KEY,
      value TEXT
    );

    CREATE TABLE IF NOT EXISTS series (
      imdb_id TEXT PRIMARY KEY,
      series_path TEXT NOT NULL,
      last_applied TEXT
    );

    CREATE TABLE IF NOT EXISTS episode_dates (
      imdb_id TEXT NOT NULL,
      season INTEGER NOT NULL,
      episode INTEGER NOT NULL,
      aired TEXT,
      dateadded TEXT,
      source TEXT,
      has_video_file INTEGER DEFAULT 0,
      PRIMARY KEY (imdb_id, season, episode)
    );
    """
    
    MOVIE_SCHEMA = """
    CREATE TABLE IF NOT EXISTS movies (
      imdb_id TEXT PRIMARY KEY,
      movie_path TEXT NOT NULL,
      last_applied TEXT
    );

    CREATE TABLE IF NOT EXISTS movie_dates (
      imdb_id TEXT PRIMARY KEY,
      released TEXT,
      dateadded TEXT,
      source TEXT,
      has_video_file INTEGER DEFAULT 0
    );
    """

    def __init__(self, db_path: Path):
        self.db_path = Path(db_path)
        self._ensure_database()

    def _ensure_database(self):
        """Create database and tables if they don't exist"""
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        
        with self.get_connection() as conn:
            # Check and create TV schema
            try:
                result = conn.execute("PRAGMA table_info(episode_dates);").fetchall()
                columns = [row[1] for row in result]
                required_columns = ['imdb_id', 'season', 'episode', 'aired', 'dateadded', 'source']
                
                if not all(col in columns for col in required_columns):
                    _log("WARNING", "TV database schema outdated, recreating tables...")
                    conn.execute("DROP TABLE IF EXISTS episode_dates;")
                    conn.execute("DROP TABLE IF EXISTS series;")
                    conn.execute("DROP TABLE IF EXISTS meta;")
                elif 'has_video_file' not in columns:
                    _log("INFO", "Adding has_video_file column to episode_dates table")
                    conn.execute("ALTER TABLE episode_dates ADD COLUMN has_video_file INTEGER DEFAULT 0;")
                    
            except sqlite3.OperationalError:
                # Tables don't exist, which is fine
                pass
            
            # Create TV schema
            conn.executescript(self.TV_SCHEMA)
            
            # Check and create movie schema
            try:
                result = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='movies'").fetchone()
                if not result:
                    _log("INFO", "Creating movie tables in database")
                    conn.executescript(self.MOVIE_SCHEMA)
            except sqlite3.Error as e:
                _log("ERROR", f"Failed to create movie tables: {e}")
            
            conn.commit()
            _log("DEBUG", "Database schema verified and up to date")

    def get_connection(self) -> sqlite3.Connection:
        """Get database connection with proper settings"""
        conn = sqlite3.Connection(str(self.db_path))
        conn.execute("PRAGMA foreign_keys=ON;")
        conn.execute("PRAGMA busy_timeout = 30000;")  # 30 second timeout
        return conn

    # TV Series Operations
    def upsert_series(self, imdb_id: str, series_path: str) -> None:
        """Insert or update series record"""
        with self.get_connection() as conn:
            now = datetime.now(timezone.utc).isoformat(timespec="seconds")
            conn.execute(
                "INSERT INTO series(imdb_id, series_path, last_applied) VALUES(?,?,?) "
                "ON CONFLICT(imdb_id) DO UPDATE SET series_path=excluded.series_path, last_applied=excluded.last_applied",
                (imdb_id, series_path, now)
            )
            conn.commit()

    def get_series_episodes(self, imdb_id: str, has_video_file_only: bool = False) -> List[Dict[str, Any]]:
        """Get all episodes for a series"""
        with self.get_connection() as conn:
            if has_video_file_only:
                query = "SELECT season, episode, aired, dateadded, source FROM episode_dates WHERE imdb_id = ? AND has_video_file = 1"
            else:
                query = "SELECT season, episode, aired, dateadded, source FROM episode_dates WHERE imdb_id = ?"
            
            rows = conn.execute(query, (imdb_id,)).fetchall()
            return [
                {
                    "season": row[0],
                    "episode": row[1], 
                    "aired": row[2],
                    "dateadded": row[3],
                    "source": row[4]
                }
                for row in rows
            ]

    def get_episode_date(self, imdb_id: str, season: int, episode: int) -> Optional[Dict[str, Any]]:
        """Get episode date information for a specific episode"""
        with self.get_connection() as conn:
            row = conn.execute(
                "SELECT aired, dateadded, source, has_video_file FROM episode_dates WHERE imdb_id = ? AND season = ? AND episode = ?",
                (imdb_id, season, episode)
            ).fetchone()
            
            if row:
                return {
                    "aired": row[0],
                    "dateadded": row[1],
                    "source": row[2],
                    "has_video_file": bool(row[3])
                }
            return None

    def upsert_episode_date(
        self, 
        imdb_id: str, 
        season: int, 
        episode: int, 
        aired: Optional[str] = None,
        dateadded: Optional[str] = None, 
        source: Optional[str] = None,
        has_video_file: bool = False
    ) -> None:
        """Insert or update episode date record"""
        with self.get_connection() as conn:
            conn.execute(
                "INSERT INTO episode_dates(imdb_id, season, episode, aired, dateadded, source, has_video_file) "
                "VALUES(?,?,?,?,?,?,?) "
                "ON CONFLICT(imdb_id, season, episode) DO UPDATE SET "
                "aired=excluded.aired, dateadded=excluded.dateadded, source=excluded.source, has_video_file=excluded.has_video_file",
                (imdb_id, season, episode, aired, dateadded, source, 1 if has_video_file else 0)
            )
            conn.commit()

    def cleanup_orphaned_episodes(self, existing_episodes: Dict[str, List[Tuple[int, int]]]) -> int:
        """
        Remove episode records that no longer have video files on disk.
        
        Args:
            existing_episodes: Dict mapping imdb_id -> [(season, episode), ...]
        """
        removed_count = 0
        
        with self.get_connection() as conn:
            # Get all series from database
            series_rows = conn.execute("SELECT imdb_id FROM series").fetchall()
            
            for (imdb_id,) in series_rows:
                if imdb_id not in existing_episodes:
                    # Series no longer exists on disk
                    removed = conn.execute("DELETE FROM episode_dates WHERE imdb_id = ?", (imdb_id,)).rowcount
                    removed_count += removed
                    _log("INFO", f"Removed {removed} episodes for missing series {imdb_id}")
                    continue
                
                # Get episodes in database for this series
                db_episodes = conn.execute(
                    "SELECT season, episode FROM episode_dates WHERE imdb_id = ?", 
                    (imdb_id,)
                ).fetchall()
                
                disk_episodes = set(existing_episodes[imdb_id])
                
                for season, episode in db_episodes:
                    if (season, episode) not in disk_episodes:
                        conn.execute(
                            "DELETE FROM episode_dates WHERE imdb_id = ? AND season = ? AND episode = ?",
                            (imdb_id, season, episode)
                        )
                        removed_count += 1
                        _log("DEBUG", f"Removed orphaned episode {imdb_id} S{season:02d}E{episode:02d}")
            
            conn.commit()
        
        return removed_count

    # Movie Operations  
    def upsert_movie(self, imdb_id: str, movie_path: str) -> None:
        """Insert or update movie record"""
        with self.get_connection() as conn:
            now = datetime.now(timezone.utc).isoformat(timespec="seconds")
            conn.execute(
                "INSERT INTO movies(imdb_id, movie_path, last_applied) VALUES(?,?,?) "
                "ON CONFLICT(imdb_id) DO UPDATE SET movie_path=excluded.movie_path, last_applied=excluded.last_applied",
                (imdb_id, movie_path, now)
            )
            conn.commit()

    def get_movie_dates(self, imdb_id: str) -> Optional[Dict[str, Any]]:
        """Get stored movie dates"""
        with self.get_connection() as conn:
            row = conn.execute(
                "SELECT released, dateadded, source, has_video_file FROM movie_dates WHERE imdb_id = ?", 
                (imdb_id,)
            ).fetchone()
            
            if row:
                return {
                    "released": row[0],
                    "dateadded": row[1],
                    "source": row[2],
                    "has_video_file": bool(row[3])
                }
        return None

    def upsert_movie_dates(
        self,
        imdb_id: str,
        released: Optional[str] = None,
        dateadded: Optional[str] = None,
        source: Optional[str] = None,
        has_video_file: bool = False
    ) -> None:
        """Insert or update movie dates"""
        with self.get_connection() as conn:
            conn.execute(
                "INSERT INTO movie_dates(imdb_id, released, dateadded, source, has_video_file) VALUES(?,?,?,?,?) "
                "ON CONFLICT(imdb_id) DO UPDATE SET "
                "released=excluded.released, dateadded=excluded.dateadded, source=excluded.source, has_video_file=excluded.has_video_file",
                (imdb_id, released, dateadded, source, 1 if has_video_file else 0)
            )
            conn.commit()

    def cleanup_orphaned_movies(self, existing_movies: List[str]) -> int:
        """Remove movie records that no longer exist on disk"""
        if not existing_movies:
            return 0
            
        removed_count = 0
        existing_set = set(existing_movies)
        
        with self.get_connection() as conn:
            # Get all movies from database
            db_movies = conn.execute("SELECT imdb_id FROM movies").fetchall()
            
            for (imdb_id,) in db_movies:
                if imdb_id not in existing_set:
                    # Movie no longer exists on disk
                    conn.execute("DELETE FROM movies WHERE imdb_id = ?", (imdb_id,))
                    conn.execute("DELETE FROM movie_dates WHERE imdb_id = ?", (imdb_id,))
                    removed_count += 1
                    _log("DEBUG", f"Removed orphaned movie {imdb_id}")
            
            conn.commit()
        
        return removed_count

    # Statistics and Reporting
    def get_stats(self) -> Dict[str, Any]:
        """Get database statistics"""
        with self.get_connection() as conn:
            stats = {}
            
            # TV stats
            stats["tv_series_count"] = conn.execute("SELECT COUNT(*) FROM series").fetchone()[0]
            stats["tv_episode_count"] = conn.execute("SELECT COUNT(*) FROM episode_dates").fetchone()[0] 
            stats["tv_episodes_with_files"] = conn.execute("SELECT COUNT(*) FROM episode_dates WHERE has_video_file = 1").fetchone()[0]
            
            # Movie stats
            stats["movie_count"] = conn.execute("SELECT COUNT(*) FROM movies").fetchone()[0]
            stats["movies_with_files"] = conn.execute("SELECT COUNT(*) FROM movie_dates WHERE has_video_file = 1").fetchone()[0]
            
            # Source breakdowns
            tv_sources = conn.execute("""
                SELECT source, COUNT(*) as count
                FROM episode_dates  
                WHERE has_video_file = 1
                GROUP BY source
                ORDER BY count DESC
            """).fetchall()
            stats["tv_sources"] = dict(tv_sources)
            
            movie_sources = conn.execute("""
                SELECT source, COUNT(*) as count
                FROM movie_dates
                WHERE has_video_file = 1  
                GROUP BY source
                ORDER BY count DESC
            """).fetchall()
            stats["movie_sources"] = dict(movie_sources)
            
            return stats

    def vacuum_database(self) -> None:
        """Vacuum database to reclaim space"""
        with self.get_connection() as conn:
            conn.execute("VACUUM;")
            _log("INFO", "Database vacuumed")

    # Utility Methods
    def clear_all_data(self) -> None:
        """Clear all data (keep schema) - useful for testing"""
        with self.get_connection() as conn:
            conn.execute("DELETE FROM episode_dates;")
            conn.execute("DELETE FROM series;")
            conn.execute("DELETE FROM movie_dates;")
            conn.execute("DELETE FROM movies;")
            conn.execute("DELETE FROM meta;")
            conn.commit()
            _log("WARNING", "All database data cleared")

    def export_series_episodes(self, imdb_id: str) -> List[Dict[str, Any]]:
        """Export all episode data for a series (for debugging)"""
        episodes = self.get_series_episodes(imdb_id)
        _log("INFO", f"Series {imdb_id} has {len(episodes)} episodes in database")
        return episodes

    def check_database_health(self) -> Dict[str, Any]:
        """Check database health and integrity"""
        health = {"status": "healthy", "issues": []}
        
        try:
            with self.get_connection() as conn:
                # Check integrity
                integrity = conn.execute("PRAGMA integrity_check;").fetchone()[0]
                if integrity != "ok":
                    health["status"] = "corrupted"
                    health["issues"].append(f"Integrity check failed: {integrity}")
                
                # Check for orphaned records
                orphaned_episodes = conn.execute("""
                    SELECT COUNT(*) FROM episode_dates 
                    WHERE imdb_id NOT IN (SELECT imdb_id FROM series)
                """).fetchone()[0]
                
                if orphaned_episodes > 0:
                    health["issues"].append(f"{orphaned_episodes} orphaned episodes")
                
                orphaned_movie_dates = conn.execute("""
                    SELECT COUNT(*) FROM movie_dates
                    WHERE imdb_id NOT IN (SELECT imdb_id FROM movies)  
                """).fetchone()[0]
                
                if orphaned_movie_dates > 0:
                    health["issues"].append(f"{orphaned_movie_dates} orphaned movie dates")
                
                health["orphaned_episodes"] = orphaned_episodes
                health["orphaned_movie_dates"] = orphaned_movie_dates
                
        except Exception as e:
            health["status"] = "error"
            health["error"] = str(e)
        
        return health


if __name__ == "__main__":
    # Test the database
    test_db = NFOGuardDatabase(Path("/tmp/test_nfoguard.db"))
    
    # Test TV operations
    test_db.upsert_series("tt1234567", "/media/tv/Test Series [imdb-tt1234567]")
    test_db.upsert_episode_date("tt1234567", 1, 1, "2020-01-01", "2025-07-07T12:00:00+00:00", "test:import", True)
    
    episodes = test_db.get_series_episodes("tt1234567")
    print(f"Test series episodes: {episodes}")
    
    # Test movie operations
    test_db.upsert_movie("tt7654321", "/media/movies/Test Movie [imdb-tt7654321]")
    test_db.upsert_movie_dates("tt7654321", "2020-01-01", "2025-07-07T12:00:00+00:00", "test:import", True)
    
    movie_dates = test_db.get_movie_dates("tt7654321")
    print(f"Test movie dates: {movie_dates}")
    
    # Get stats
    stats = test_db.get_stats()
    print(f"Database stats: {stats}")
    
    # Check health
    health = test_db.check_database_health()
    print(f"Database health: {health}")
    
    print("Database test completed successfully!")