bazarr/custom_libs/subzero/subtitle_storage.py

# coding=utf-8
from __future__ import absolute_import
import datetime
import gzip
import hashlib
import os
import logging
import traceback
import types
import zlib

import sys

from json_tricks.nonp import loads
from subzero.lib.json import dumps
try:
    from os import scandir
    _scandir_generic = scandir
except ImportError:
    from scandir import scandir, scandir_generic as _scandir_generic
from .constants import mode_map
import six

logger = logging.getLogger(__name__)


class JSONStoredSubtitle(object):
    score = None
    storage_type = None
    hash = None
    provider_name = None
    id = None
    date_added = None
    mode = "a"  # auto/manual/auto-better (a/m/b)
    content = None
    mods = None
    encoding = None
    last_mod = None  # file modification

    def initialize(self, score, storage_type, hash, provider_name, id, date_added=None, mode="a", content=None,
                 mods=None, encoding=None, last_mod=None):
        self.score = int(score)
        self.storage_type = storage_type
        self.hash = hash
        self.provider_name = provider_name
        self.id = id
        self.date_added = date_added or datetime.datetime.now()
        self.mode = mode
        self.content = content
        self.mods = mods or []
        self.encoding = encoding
        self.last_mod = last_mod

    def add_mod(self, identifier):
        self.mods = self.mods or []
        if identifier is None:
            self.mods = []
            return

        self.mods.append(identifier)

    @classmethod
    def get_mode_verbose(cls, mode):
        return mode_map.get(mode, "Unknown")

    @property
    def mode_verbose(self):
        return self.get_mode_verbose(self.mode)

    def serialize(self):
        return self.__dict__

    def deserialize(self, data):
        if data["content"]:
            # legacy: storage was unicode; content is always present in encoded form
            if isinstance(data["content"], str):
                data["content"] = data["content"].encode(data["encoding"])

        self.initialize(**data)

    @property
    def key(self):
        return self.provider_name, self.id


class JSONStoredVideoSubtitles(object):
    """
    manages stored subtitles for video_id per media_part/language combination
    """
    video_id = None  # rating_key
    title = None
    parts = None
    version = None
    item_type = None  # movie / episode
    added_at = None

    def initialize(self, plex_item, version=None):
        self.video_id = str(plex_item.rating_key)

        self.title = plex_item.title
        self.parts = {}
        self.version = version
        self.item_type = plex_item.type
        self.added_at = datetime.datetime.fromtimestamp(plex_item.added_at)

    def deserialize(self, data):
        parts = data.pop("parts")
        self.parts = {}
        self.__dict__.update(data)

        if parts:
            for part_id, part in six.iteritems(parts):
                self.parts[part_id] = {}
                for language, sub_data in six.iteritems(part):
                    self.parts[part_id][language] = {}

                    for sub_key, subtitle_data in six.iteritems(sub_data):
                        if sub_key == "current":
                            if not isinstance(subtitle_data, tuple):
                                subtitle_data = tuple(subtitle_data.split("__"))
                            self.parts[part_id][language]["current"] = subtitle_data
                        elif sub_key == "blacklist":
                            bl = dict((tuple([str(a) for a in k.split("__")]), v) for k, v in six.iteritems(subtitle_data))
                            self.parts[part_id][language]["blacklist"] = bl
                        else:
                            sub = JSONStoredSubtitle()

                            sub.initialize(**subtitle_data)
                            if not isinstance(sub_key, tuple):
                                sub_key = tuple(sub_key.split("__"))

                            self.parts[part_id][language][sub_key] = sub

    def serialize(self):
        data = {"parts": {}}
        for key, value in six.iteritems(self.__dict__):
            if key != "parts":
                data[key] = value

        for part_id, part in six.iteritems(self.parts):
            data["parts"][part_id] = {}
            for language, sub_data in six.iteritems(part):
                data["parts"][part_id][language] = {}

                for sub_key, stored_subtitle in six.iteritems(sub_data):
                    if sub_key == "current":
                        data["parts"][part_id][language]["current"] = "__".join(stored_subtitle)
                    elif sub_key == "blacklist":
                        data["parts"][part_id][language]["blacklist"] = dict(("__".join(k), v) for k, v in
                                                                             six.iteritems(stored_subtitle))
                    else:
                        if stored_subtitle.content and not stored_subtitle.encoding:
                            continue

                        serialized_sub = stored_subtitle.serialize()
                        if serialized_sub:
                            data["parts"][part_id][language]["__".join(sub_key)] = serialized_sub

        return data

    def add(self, part_id, lang, subtitle, storage_type, date_added=None, mode="a", last_mod=None, set_current=True):
        part_id = str(part_id)
        part = self.parts.get(part_id)
        if not part:
            self.parts[part_id] = {}
            part = self.parts[part_id]

        subs = part.get(lang)
        if not subs:
            part[lang] = {}
            subs = part[lang]

        sub_key = self.get_sub_key(subtitle.provider_name, subtitle.id)
        subs[sub_key] = JSONStoredSubtitle()
        subs[sub_key].initialize(subtitle.score, storage_type, hashlib.md5(subtitle.content).hexdigest(),
                                 subtitle.provider_name, subtitle.id, date_added=date_added, mode=mode,
                                 content=subtitle.content, mods=subtitle.mods, encoding="utf-8",
                                 last_mod=last_mod)

        if set_current:
            logger.debug(u"Setting subtitle as current: %r", subtitle)
            subs["current"] = sub_key

        return True

    def get_any(self, part_id, lang):
        part_id = str(part_id)
        part = self.parts.get(part_id)
        if not part:
            return

        subs = part.get(str(lang))
        if not subs:
            return

        if "current" in subs and subs["current"]:
            return subs.get(subs["current"])

    def get(self, part_id, lang, sub_key):
        subs = self.get_all(part_id, lang)
        if not subs:
            return

        return subs.get(sub_key)

    def get_all(self, part_id, lang):
        part_id = str(part_id)
        part = self.parts.get(part_id)
        if not part:
            return

        return part.get(str(lang))

    def set_current(self, part_id, lang, sub_key):
        subs = self.get_all(part_id, lang)
        if not subs:
            return

        if sub_key not in subs:
            logger.info("Tried setting unknown subtitle %s as current" % sub_key)
            return

        subs["current"] = sub_key
        logger.debug("Set subtitle %s as current for %s, %s" % (sub_key, part_id, lang))

    def get_by_provider(self, part_id, lang, provider_name):
        out = []
        all_subs = self.get_all(part_id, lang)
        if not all_subs:
            return out

        for key, subtitle in six.iteritems(all_subs):
            if key in ("current", "blacklist"):
                continue

            if subtitle.provider_name == provider_name:
                out.append(subtitle)
        return out

    def count(self, part_id, lang):
        part_id = str(part_id)
        part = self.parts.get(part_id)
        if not part:
            return 0

        subs = part.get(str(lang))
        return len([key for key in list(subs.keys()) if key not in ("current", "blacklist")])

    def get_sub_key(self, provider_name, id):
        return provider_name, str(id)

    def get_blacklist(self, part_id, lang):
        part_id = str(part_id)
        part = self.parts.get(part_id)
        if not part:
            return {}, {}

        subs = part.get(str(lang))
        if not subs:
            return {}, {}

        current_bl = subs.get("blacklist", {})
        return current_bl, subs

    def blacklist(self, part_id, lang, sub_key, add=True):
        current_bl, subs = self.get_blacklist(part_id, lang)
        sub = subs.get(subs["current"])
        if not sub:
            return

        if sub_key in current_bl:
            if add:
                return
            else:
                del current_bl[sub_key]
                subs["blacklist"] = current_bl
                return

        current_bl[sub_key] = {"date_added": sub.date_added, "score": sub.score, "mode": sub.mode, "storage_type":
            sub.storage_type}
        subs["blacklist"] = current_bl

    def __repr__(self):
        return six.text_type(self)

    def __unicode__(self):
        return u"%s (%s)" % (self.title, self.video_id)

    def __str__(self):
        return str(self.video_id)


class StoredSubtitlesManager(object):
    """
    manages the storage and retrieval of StoredVideoSubtitles instances for a given video_id
    """
    storage = None
    version = 3
    extension = ".json.gz"

    def __init__(self, storage, threadkit, plexapi_item_getter):
        self.storage = storage
        self.get_item = plexapi_item_getter
        self.threadkit = threadkit

    def destroy(self):
        self.storage = None
        self.get_item = None
        self.threadkit = None

    def get_storage_filename(self, video_id):
        return "subs_%s" % video_id

    @property
    def dataitems_path(self):
        return os.path.join(getattr(self.storage, "_core").storage.data_path, "DataItems")

    def get_json_data_path(self, bare_fn):
        if not bare_fn.endswith(self.extension):
            return os.path.join(self.dataitems_path, "%s%s" % (bare_fn, self.extension))
        return os.path.join(self.dataitems_path, bare_fn)

    def get_all_files(self, scandir_generic=False):
        _scandir = _scandir_generic if scandir_generic else scandir
        for entry in _scandir(self.dataitems_path):
            if entry.is_file(follow_symlinks=False) and \
                    entry.name.startswith("subs_") and \
                    entry.name.endswith(self.extension):
                yield entry.name

    def get_recent_files(self, age_days=30):
        fl = []
        root = self.dataitems_path
        recent_dt = datetime.datetime.now() - datetime.timedelta(days=age_days)

        def run(scandir_generic=False):
            for fn in self.get_all_files(scandir_generic=scandir_generic):
                ctime = os.path.getctime(os.path.join(root, fn))
                created = datetime.datetime.fromtimestamp(ctime)
                if created > recent_dt:
                    fl.append(fn)
        try:
            run()
        except OSError:
            run(scandir_generic=True)
        return fl

    def load_recent_files(self, age_days=30):
        fl = self.get_recent_files(age_days=age_days)
        out = {}
        for fn in fl:
            data = self.load(filename=fn)
            if data:
                out[fn] = data
        return out

    def delete_missing(self, wanted_languages=set(), scandir_generic=False):
        deleted = []

        def delete_fn(filename):
            if filename.endswith(self.extension):
                self.delete(self.get_json_data_path(filename))
            else:
                self.legacy_delete(filename)

        for fn in self.get_all_files(scandir_generic=scandir_generic):
            video_id = os.path.basename(fn).split(".")[0].split("subs_")[1]
            item = self.get_item(video_id)

            # item missing, delete storage
            if not item:
                delete_fn(fn)
                deleted.append(video_id)

            else:
                known_parts = []

                # wrong (legacy) info, delete storage
                if not hasattr(item, "media"):
                    delete_fn(fn)
                    deleted.append(video_id)
                    continue

                for media in item.media:
                    for part in media.parts:
                        known_parts.append(str(part.id))
                stored_subs = self.load(filename=fn)

                if not stored_subs:
                    continue

                missing_parts = set(stored_subs.parts).difference(set(known_parts))

                changed_any = False

                # remove known info about deleted parts
                if missing_parts:
                    logger.debug("Parts removed: %s:%s, removing data", video_id, missing_parts)
                    for missing_part in missing_parts:
                        if missing_part in stored_subs.parts:
                            try:
                                del stored_subs.parts[missing_part]
                                changed_any = True
                            except:
                                pass

                # remove known info about non-existing languages
                for part_id, part in six.iteritems(stored_subs.parts):
                    missing_languages = set(part).difference(wanted_languages)
                    if missing_languages:
                        logger.debug("Languages removed: %s:%s:%s, removing data", video_id, part_id, missing_languages)
                        for missing_language in missing_languages:
                            try:
                                del stored_subs.parts[part_id][missing_language]
                                changed_any = True
                            except:
                                pass

                if changed_any:
                    self.save(stored_subs)
                stored_subs = None
                missing_parts = None
                missing_languages = None

        return deleted

    def migrate_v2(self, subs_for_video):
        plex_item = self.get_item(subs_for_video.video_id)
        if not plex_item:
            return False
        subs_for_video.item_type = plex_item.type
        subs_for_video.added_at = datetime.datetime.fromtimestamp(plex_item.added_at)
        subs_for_video.version = 2
        return True

    def migrate_v3(self, subs_for_video):
        subs_for_video.version = 3
        return True

    def load(self, video_id=None, filename=None):
        subs_for_video = None
        bare_fn = self.get_storage_filename(video_id) if video_id else filename
        json_path = self.get_json_data_path(bare_fn)
        basename = os.path.basename(json_path)

        #logger.debug("Loading subtitle storage data file: %s", basename)

        if os.path.exists(json_path):
            # new style data
            subs_for_video = JSONStoredVideoSubtitles()
            try:
                with self.threadkit.Lock(key="sub_storage_%s" % basename):
                    if sys.platform == "win32":
                        try:
                            with open(json_path, 'rb') as f:
                                s = zlib.decompress(f.read())
                        except zlib.error:
                            # fallback to old gzip win32 implementation
                            with gzip.open(json_path, 'rb', compresslevel=6) as f:
                                s = f.read()

                    else:
                        with gzip.open(json_path, 'rb', compresslevel=6) as f:
                            s = f.read()

                data = loads(s)
            except:
                logger.error("Couldn't load JSON data for %s: %s", bare_fn, traceback.format_exc())
                return

            subs_for_video.deserialize(data)
            data = None

        if not subs_for_video:
            return

        # apply possible migrations
        cur_ver = old_ver = subs_for_video.version

        if cur_ver < self.version:
            success = False
            while cur_ver < self.version:
                cur_ver += 1
                mig_func = "migrate_v%s" % cur_ver
                if hasattr(self, mig_func):
                    logger.info("Migrating subtitle storage for %s %s>%s" % (subs_for_video.video_id, old_ver, cur_ver))
                    success = getattr(self, mig_func)(subs_for_video)
                    if success is False:
                        logger.error("Couldn't migrate %s, removing data", subs_for_video.video_id)
                        self.delete(json_path)
                        break

            if cur_ver > old_ver and success:
                logger.info("Storing migrated subtitle storage for %s" % subs_for_video.video_id)
                self.save(subs_for_video)
            elif not success:
                logger.info("Migration of %s %s>%s failed" % (subs_for_video.video_id, old_ver, cur_ver))

        return subs_for_video

    def new(self, plex_item):
        subs_for_video = JSONStoredVideoSubtitles()
        subs_for_video.initialize(plex_item, version=self.version)
        return subs_for_video

    def load_or_new(self, plex_item, save=False):
        subs_for_video = self.load(plex_item.rating_key)
        if not subs_for_video:
            logger.info("Creating new subtitle storage for: %s", plex_item.rating_key)
            subs_for_video = self.new(plex_item)
            if save:
                self.save(subs_for_video)
        return subs_for_video

    def save(self, subs_for_video):
        data = subs_for_video.serialize()
        temp_fn = self.get_json_data_path(self.get_storage_filename(subs_for_video.video_id) + "_tmp")
        fn = self.get_json_data_path(self.get_storage_filename(subs_for_video.video_id))
        basename = os.path.basename(fn)
        json_data = str(dumps(data, ensure_ascii=False))
        with self.threadkit.Lock(key="sub_storage_%s" % basename):
            if sys.platform == "win32":
                try:
                    f = open(temp_fn, "w+b")

                    try:
                        f.seek(0, os.SEEK_CUR)
                        f.write(zlib.compress(json_data, 6))
                        f.flush()
                    except:
                        logger.error("Something went wrong when writing to: %s: %s", basename, traceback.format_exc())
                    finally:
                        f.close()
                except:
                    logger.error("Something went REALLY wrong when writing to: %s: %s", basename,
                                 traceback.format_exc())
            else:
                with gzip.open(temp_fn, "wb", compresslevel=6) as f:
                    f.write(json_data)

            os.rename(temp_fn, fn)

    def delete(self, filename):
        os.remove(filename)

    def legacy_delete(self, filename):
        try:
            self.storage.Remove(filename)
        except:
            logger.error("Failed to delete item %s: %s" % (filename, traceback.format_exc()))