nfoguard/core/nfo_manager.py

#!/usr/bin/env python3
"""
NFO file management module for movies and TV episodes
"""
import os
import re
from pathlib import Path
from datetime import datetime, timezone
from typing import Optional
import xml.etree.ElementTree as ET


def _log(level: str, msg: str):
    """Placeholder logging function - replace with your actual logger"""
    print(f"[{datetime.now().isoformat()}] {level}: {msg}")


class NFOManager:
    """Handles creation and updating of NFO files for movies and TV episodes"""
    
    # Regex patterns - supports both [imdb-tt123] and [tt123] formats
    IMDB_TAG_PATTERN = re.compile(r"\[(?:imdb-)?(tt\d+)\]", re.IGNORECASE)
    LONG_NFO_PATTERN = re.compile(r".*-S\d{2}E\d{2}-.*\.nfo$", re.IGNORECASE)
    SHORT_NFO_PATTERN = re.compile(r"^S(\d{2})E(\d{2})\.nfo$", re.IGNORECASE)
    
    def __init__(self, manager_brand: str = "NFOGuard"):
        self.manager_brand = manager_brand

    @staticmethod
    def parse_imdb_from_path(path: Path) -> Optional[str]:
        """Extract IMDb ID from directory or filename, with .nfo fallback"""
        # First try path name
        match = NFOManager.IMDB_TAG_PATTERN.search(path.name)
        if match:
            return match.group(1).lower()
        
        # Fallback: scan .nfo files in directory
        if path.is_dir():
            for nfo_file in path.glob("*.nfo"):
                imdb_id = NFOManager.parse_imdb_from_nfo(nfo_file)
                if imdb_id:
                    _log("DEBUG", f"Found IMDb ID {imdb_id} in {nfo_file.name}")
                    return imdb_id
        
        return None

    @staticmethod
    def parse_imdb_from_nfo(nfo_path: Path) -> Optional[str]:
        """Extract IMDb ID from existing NFO file"""
        if not nfo_path.exists():
            return None
            
        try:
            tree = ET.parse(str(nfo_path))
            root = tree.getroot()
            
            # Check uniqueid elements
            for uniqueid in root.findall("uniqueid"):
                if (uniqueid.attrib.get("type") or "").lower() == "imdb":
                    imdb_id = (uniqueid.text or "").strip()
                    if imdb_id.startswith("tt"):
                        return imdb_id.lower()
            
            # Check direct imdbid elements (older format)
            imdbid_elem = root.find("imdbid")
            if imdbid_elem is not None and imdbid_elem.text:
                imdb_id = imdbid_elem.text.strip()
                if imdb_id.startswith("tt"):
                    return imdb_id.lower()
                    
        except Exception as e:
            _log("DEBUG", f"Could not parse IMDb from NFO {nfo_path}: {e}")
        
        return None

    @staticmethod
    def _escape_xml(text: str) -> str:
        """Escape XML special characters"""
        if not text:
            return ""
        return (str(text)
                .replace("&", "&")
                .replace("<", "&lt;")
                .replace(">", "&gt;")
                .replace('"', "&quot;")
                .replace("'", "&apos;"))

    @staticmethod
    def _ensure_child_text(parent: ET.Element, tag: str, text: str, overwrite: bool = True) -> ET.Element:
        """Ensure child element exists with specified text"""
        child = parent.find(tag)
        if child is None:
            child = ET.SubElement(parent, tag)
            child.text = text
        else:
            if overwrite or (child.text is None or str(child.text).strip() == ""):
                child.text = text
        return child

    @staticmethod
    def _ensure_uniqueid_imdb(root: ET.Element, imdb_id: str):
        """Ensure IMDb uniqueid element exists"""
        imdb_elem = None
        for uid in root.findall("uniqueid"):
            if (uid.attrib.get("type") or "").lower() == "imdb":
                imdb_elem = uid
                break
                
        if imdb_elem is None:
            imdb_elem = ET.SubElement(root, "uniqueid", {"type": "imdb", "default": "true"})
        else:
            if "default" not in imdb_elem.attrib:
                imdb_elem.set("default", "true")
                
        imdb_elem.text = imdb_id

    @staticmethod
    def _indent_xml(elem: ET.Element, level: int = 0, space: str = "  "):
        """Pretty-print XML with proper indentation"""
        i = "\n" + level * space
        if len(elem):
            if not elem.text or not elem.text.strip():
                elem.text = i + space
            for idx, e in enumerate(list(elem)):
                NFOManager._indent_xml(e, level + 1, space)
                if not e.tail or not e.tail.strip():
                    e.tail = i + (space if idx < len(elem) - 1 else "")
            if not elem.tail or not elem.tail.strip():
                elem.tail = "\n" + (level - 1) * space if level else "\n"
        else:
            if not elem.text or not elem.text.strip():
                elem.text = ""
            if not elem.tail or not elem.tail.strip():
                elem.tail = "\n" + (level - 1) * space if level else "\n"

    @staticmethod
    def _strip_existing_footer_comments(root: ET.Element):
        """Remove existing manager comments from XML"""
        for node in list(root):
            if hasattr(node, 'text') and isinstance(node.text, str):
                text_lower = node.text.lower()
                if "source:" in text_lower or "managed by" in text_lower:
                    root.remove(node)

    def create_movie_nfo(
        self, 
        movie_dir: Path, 
        imdb_id: str, 
        dateadded: Optional[str] = None,
        released: Optional[str] = None,
        source_detail: str = "unknown",
        lock_metadata: bool = True
    ):
        """Create or update movie.nfo file"""
        canonical_path = movie_dir / "movie.nfo"
        alt_path = movie_dir / f"{movie_dir.name}.nfo"

        # Try to load existing NFO
        root = None
        for nfo_path in (canonical_path, alt_path):
            if nfo_path.exists():
                try:
                    tree = ET.parse(str(nfo_path))
                    existing_root = tree.getroot()
                    if existing_root is not None and existing_root.tag.lower() == "movie":
                        root = existing_root
                        break
                except Exception as e:
                    _log("DEBUG", f"Could not parse existing NFO {nfo_path}: {e}")

        if root is None:
            root = ET.Element("movie")

        # Ensure IMDb ID format
        if imdb_id and not imdb_id.lower().startswith("tt"):
            imdb_id = f"tt{imdb_id}"

        # Add/update elements
        if imdb_id:
            self._ensure_uniqueid_imdb(root, imdb_id)

        if released:
            # Extract date part from ISO string
            release_date = released.split("T")[0]
            self._ensure_child_text(root, "premiered", release_date, overwrite=False)
            self._ensure_child_text(root, "year", release_date.split("-")[0], overwrite=False)

        if dateadded:
            if dateadded == "MANUAL_REVIEW_NEEDED":
                self._ensure_child_text(root, "dateadded", "MANUAL_REVIEW_NEEDED", overwrite=True)
                _log("WARNING", f"Set MANUAL_REVIEW_NEEDED marker in NFO: {movie_dir}/movie.nfo")
            else:
                self._ensure_child_text(root, "dateadded", dateadded, overwrite=True)

        if lock_metadata:
            self._ensure_child_text(root, "lockdata", "true", overwrite=True)

        # Add footer comments with timestamp
        self._strip_existing_footer_comments(root)
        current_time = datetime.now(timezone.utc).isoformat(timespec="seconds")
        root.append(ET.Comment(f" source: Movie; {source_detail} "))
        root.append(ET.Comment(f" managed by {self.manager_brand} at {current_time} "))

        # Pretty-print and save
        self._indent_xml(root)

        # Write to both canonical path and alt path if it exists
        for target in [canonical_path] + ([alt_path] if alt_path.exists() else []):
            try:
                target.parent.mkdir(parents=True, exist_ok=True)
                ET.ElementTree(root).write(str(target), encoding="utf-8", xml_declaration=True)
                _log("DEBUG", f"Updated NFO: {target}")
            except Exception as e:
                _log("ERROR", f"Failed writing {target}: {e}")

    def create_episode_nfo(
        self,
        season_dir: Path,
        season_num: int,
        episode_num: int,
        aired: Optional[str] = None,
        dateadded: Optional[str] = None,
        source_detail: str = "unknown",
        lock_metadata: bool = True,
        enhanced_metadata: Optional[Dict[str, Any]] = None
    ):
        """Create episode NFO file (S01E01.nfo format)"""
        try:
            season_dir.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            _log("ERROR", f"Failed creating season directory {season_dir}: {e}")
            return

        nfo_path = season_dir / f"S{season_num:02d}E{episode_num:02d}.nfo"
        
        # Build XML content
        xml_lines = [
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>',
            "<episodedetails>",
            f"  <season>{season_num}</season>",
            f"  <episode>{episode_num}</episode>",
        ]
        
        # Add enhanced metadata if available
        if enhanced_metadata:
            if enhanced_metadata.get("title"):
                xml_lines.append(f"  <title>{self._escape_xml(enhanced_metadata['title'])}</title>")
            
            if enhanced_metadata.get("overview"):
                xml_lines.append(f"  <plot>{self._escape_xml(enhanced_metadata['overview'])}</plot>")
            
            if enhanced_metadata.get("runtime"):
                xml_lines.append(f"  <runtime>{enhanced_metadata['runtime']}</runtime>")
            
            # Add ratings if available
            if enhanced_metadata.get("ratings"):
                ratings = enhanced_metadata["ratings"]
                if ratings.get("imdb", {}).get("value"):
                    imdb_rating = ratings["imdb"]["value"]
                    imdb_votes = ratings["imdb"].get("votes", 0)
                    xml_lines.append(f"  <rating>{imdb_rating}</rating>")
                    xml_lines.append(f"  <votes>{imdb_votes}</votes>")
        
        if aired:
            # Extract date part from ISO string  
            air_date = aired.split('T')[0]
            xml_lines.append(f"  <aired>{air_date}</aired>")
        
        if dateadded:
            xml_lines.append(f"  <dateadded>{dateadded}</dateadded>")
            # Add premiered tag for broader compatibility
            premiered_date = dateadded.split('T')[0]
            xml_lines.append(f"  <premiered>{premiered_date}</premiered>")
        
        if lock_metadata:
            xml_lines.append("  <lockdata>true</lockdata>")
        
        # Add source comments with timestamp
        if source_detail:
            current_time = datetime.now(timezone.utc).isoformat(timespec="seconds")
            xml_lines.append(f"  <!-- source: TV Episode; {source_detail} -->")
            xml_lines.append(f"  <!-- managed by {self.manager_brand} at {current_time} -->")
        
        xml_lines.append("</episodedetails>")
        xml_lines.append("")  # Final newline
        
        content = "\n".join(xml_lines)
        
        try:
            nfo_path.write_text(content, encoding="utf-8")
            _log("DEBUG", f"Created episode NFO: {nfo_path}")
        except Exception as e:
            _log("ERROR", f"Failed writing {nfo_path}: {e}")

    def create_season_nfo(self, season_dir: Path, season_num: int):
        """Create season.nfo if it doesn't exist"""
        season_nfo = season_dir / "season.nfo"
        if season_nfo.exists():
            return

        xml_content = f"""<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<season>
  <seasonnumber>{season_num}</seasonnumber>
</season>
"""
        try:
            season_nfo.write_text(xml_content, encoding="utf-8")
            _log("DEBUG", f"Created season NFO: {season_nfo}")
        except Exception as e:
            _log("WARNING", f"Could not write season.nfo in {season_dir}: {e}")

    def create_tvshow_nfo(self, series_dir: Path, imdb_id: str):
        """Create tvshow.nfo if it doesn't exist"""
        tvshow_nfo = series_dir / "tvshow.nfo"
        if tvshow_nfo.exists():
            return

        xml_content = f"""<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<tvshow>
  <id>{imdb_id}</id>
  <uniqueid type="imdb" default="true">{imdb_id}</uniqueid>
</tvshow>
"""
        try:
            tvshow_nfo.write_text(xml_content, encoding="utf-8")
            _log("DEBUG", f"Created tvshow NFO: {tvshow_nfo}")
        except Exception as e:
            _log("WARNING", f"Could not create tvshow.nfo in {series_dir}: {e}")

    def remove_long_episode_nfos(self, season_dir: Path) -> int:
        """Remove long-format episode NFO files (e.g., Show.Name.S01E01.nfo)"""
        count = 0
        for nfo_file in season_dir.glob("*.nfo"):
            if self.LONG_NFO_PATTERN.match(nfo_file.name):
                try:
                    nfo_file.unlink()
                    count += 1
                    _log("DEBUG", f"Removed long NFO: {nfo_file}")
                except Exception as e:
                    _log("WARNING", f"Failed to remove long NFO {nfo_file}: {e}")
        return count

    def remove_orphaned_nfos(self, season_dir: Path, existing_episodes: set) -> int:
        """Remove NFO files that don't have corresponding video files"""
        count = 0
        for nfo_file in season_dir.glob("S??E??.nfo"):
            match = self.SHORT_NFO_PATTERN.match(nfo_file.name)
            if match:
                nfo_season = int(match.group(1))
                nfo_episode = int(match.group(2))
                
                # Check if we have a video file for this episode
                if (nfo_season, nfo_episode) not in existing_episodes:
                    try:
                        nfo_file.unlink()
                        count += 1
                        _log("DEBUG", f"Removed orphaned NFO: {nfo_file}")
                    except Exception as e:
                        _log("WARNING", f"Failed to remove orphaned NFO {nfo_file}: {e}")
        return count

    @staticmethod
    def set_file_mtime(file_path: Path, iso_timestamp: str):
        """Set file modification time to match the timestamp"""
        try:
            timestamp = datetime.fromisoformat(iso_timestamp.replace("Z", "+00:00")).timestamp()
            os.utime(file_path, (timestamp, timestamp), follow_symlinks=False)
            _log("DEBUG", f"Updated mtime: {file_path} -> {iso_timestamp}")
        except Exception as e:
            _log("WARNING", f"Failed to set mtime on {file_path}: {e}")

    def set_directory_mtime(self, dir_path: Path, iso_timestamp: str):
        """Set directory modification time"""
        try:
            timestamp = datetime.fromisoformat(iso_timestamp.replace("Z", "+00:00")).timestamp()
            os.utime(dir_path, (timestamp, timestamp), follow_symlinks=False)
            _log("DEBUG", f"Updated directory mtime: {dir_path} -> {iso_timestamp}")
        except Exception as e:
            _log("WARNING", f"Failed to set mtime on directory {dir_path}: {e}")

    def update_movie_files_mtime(self, movie_dir: Path, iso_timestamp: str):
        """Update modification time for all files in movie directory"""
        if not iso_timestamp or iso_timestamp == "MANUAL_REVIEW_NEEDED":
            return

        try:
            timestamp = datetime.fromisoformat(iso_timestamp.replace("Z", "+00:00")).timestamp()
            
            files_updated = 0
            for file_path in movie_dir.iterdir():
                if file_path.is_file():
                    try:
                        os.utime(file_path, (timestamp, timestamp), follow_symlinks=False)
                        files_updated += 1
                    except Exception as e:
                        _log("WARNING", f"Failed to update mtime on {file_path.name}: {e}")
            
            # Update directory timestamp
            os.utime(movie_dir, (timestamp, timestamp), follow_symlinks=False)
            _log("INFO", f"Updated {files_updated} files and directory mtime: {movie_dir.name} -> {iso_timestamp}")
            
        except Exception as e:
            _log("WARNING", f"Failed to update movie directory mtimes: {e}")


if __name__ == "__main__":
    # Test the NFO manager
    nfo_manager = NFOManager("NFOGuard-Test")
    
    # Test movie NFO creation
    test_movie_dir = Path("/tmp/test_movie")
    test_movie_dir.mkdir(exist_ok=True)
    
    nfo_manager.create_movie_nfo(
        test_movie_dir,
        "tt1234567",
        "2025-07-07T12:00:00+00:00",
        "2020-01-01T00:00:00+00:00",
        "test:movie_creation"
    )
    
    # Test episode NFO creation
    test_season_dir = Path("/tmp/test_series/Season 01")
    test_season_dir.mkdir(parents=True, exist_ok=True)
    
    nfo_manager.create_episode_nfo(
        test_season_dir,
        1, 1,
        "2020-01-01T00:00:00+00:00",
        "2025-07-07T12:00:00+00:00", 
        "test:episode_creation"
    )
    
    print("NFO test files created in /tmp/")