#!/usr/bin/env python
# License: GPLv3 Copyright: 2014, Kovid Goyal <kovid at kovidgoyal.net>

import errno
import os
import shutil
import time
import traceback
from collections import defaultdict, namedtuple
from io import BytesIO
from queue import Empty
from typing import cast

from qt.core import QObject, Qt, pyqtSignal

from calibre import force_unicode, prints
from calibre.constants import DEBUG
from calibre.customize.ui import can_set_metadata
from calibre.db.constants import DATA_FILE_PATTERN
from calibre.db.errors import NoSuchFormat
from calibre.ebooks.metadata import authors_to_string
from calibre.ebooks.metadata.opf2 import metadata_to_opf
from calibre.gui2 import error_dialog, gprefs, open_local_file, warning_dialog
from calibre.gui2.dialogs.progress import ProgressDialog
from calibre.library.save_to_disk import find_plugboard, get_path_components, plugboard_save_to_disk_value, sanitize_args
from calibre.ptempfile import PersistentTemporaryDirectory, SpooledTemporaryFile
from calibre.utils.filenames import make_long_path_useable
from calibre.utils.ipc.pool import Failure, Pool
from calibre.utils.localization import _

BookId = namedtuple('BookId', 'title authors')


def ensure_unique_components(data):  # {{{
    cmap = defaultdict(set)
    bid_map = {}
    for book_id, (mi, components, fmts) in data.items():
        cmap[tuple(components)].add(book_id)
        bid_map[book_id] = components

    for book_ids in cmap.values():
        if len(book_ids) > 1:
            for i, book_id in enumerate(sorted(book_ids)[1:]):
                suffix = f' ({i + 1})'
                components = bid_map[book_id]
                components[-1] = components[-1] + suffix


# }}}


class SpooledFile(SpooledTemporaryFile):  # {{{
    def __init__(self, file_obj, max_size=50 * 1024 * 1024):
        self._file_obj = file_obj
        SpooledTemporaryFile.__init__(self, max_size)

    def rollover(self):
        if self._rolled:
            return
        orig = self._file
        newfile = self._file = self._file_obj
        del self._TemporaryFileArgs  # type: ignore

        newfile.write(cast(BytesIO, orig).getvalue())
        newfile.seek(orig.tell(), 0)

        self._rolled = True


# }}}


class Saver(QObject):
    do_one_signal = pyqtSignal()
    pd_title = ''

    def __init__(self, book_ids, db, opts, root, parent=None, pool=None):
        QObject.__init__(self, parent)
        self.db = db.new_api
        self.plugboards = self.db.pref('plugboards', {})
        self.template_functions = self.db.pref('user_template_functions', [])
        self.library_id = self.db.library_id
        # This call to load_user_template_functions isn't needed because
        # __init__ is running on the GUI thread. It must be done in the separate
        # process by the worker
        # from calibre.gui2 import is_gui_thread
        # print(f'Saver __init__ is_gui_thread: {is_gui_thread()}')
        # load_user_template_functions('', self.template_functions)
        self.collected_data = {}
        self.errors = defaultdict(list)
        self._book_id_data = {}
        self.all_book_ids = frozenset(book_ids)
        self.pd = ProgressDialog(
            self.pd_title or (_('Saving %d books...') % len(self.all_book_ids)),
            _('Collecting metadata...'),
            min=0,
            max=0,
            parent=parent,
            icon='save.png',
        )
        self.do_one_signal.connect(self.tick, type=Qt.ConnectionType.QueuedConnection)
        self.do_one = self.do_one_collect
        self.ids_to_collect = iter(self.all_book_ids)
        self.tdir = PersistentTemporaryDirectory('_save_to_disk')
        self.pool = pool

        self.pd.show()
        self.root, self.opts, self.path_length = sanitize_args(root, opts)
        self.do_one_signal.emit()
        if DEBUG:
            self.start_time = time.time()

    def tick(self):
        pd = self.pd
        assert pd is not None
        if pd.canceled:
            pd.close()
            pd.deleteLater()
            self.break_cycles()
            return
        self.do_one()

    def break_cycles(self):
        shutil.rmtree(self.tdir, ignore_errors=True)
        if self.pool is not None:
            self.pool.shutdown()
        self.setParent(None)
        self.jobs = self.pool = self.plugboards = self.template_functions = self.library_id = self.collected_data = self.all_book_ids = self.pd = self.db = None
        self.deleteLater()

    def book_id_data(self, book_id):
        ans = self._book_id_data.get(book_id)
        if ans is None:
            try:
                db = self.db
                assert db is not None
                ans = BookId(db.field_for('title', book_id), db.field_for('authors', book_id))
            except Exception:
                ans = BookId((_('Unknown') + f' ({book_id})'), (_('Unknown'),))
            self._book_id_data[book_id] = ans
        return ans

    def do_one_collect(self):
        try:
            book_id = next(self.ids_to_collect)
        except StopIteration:
            self.collection_finished()
            return
        try:
            self.collect_data(book_id)
        except Exception:
            self.errors[book_id].append(('critical', traceback.format_exc()))
        self.do_one_signal.emit()

    def collect_data(self, book_id):
        db = self.db
        assert db is not None
        mi = db.get_metadata(book_id)
        self._book_id_data[book_id] = BookId(mi.title, mi.authors)
        components = get_path_components(self.opts, mi, book_id, self.path_length)
        collected_data = self.collected_data
        assert collected_data is not None
        collected_data[book_id] = (mi, components, {fmt.lower() for fmt in db.formats(book_id)})

    def collection_finished(self):
        self.do_one = self.do_one_write
        collected_data = self.collected_data
        assert collected_data is not None
        ensure_unique_components(collected_data)
        self.ids_to_write = iter(collected_data)
        pd = self.pd
        assert pd is not None
        pd.title = _('Copying files and writing metadata...') if self.opts.update_metadata else _('Copying files...')
        pd.max = len(collected_data)
        pd.value = 0
        if self.opts.update_metadata:
            all_fmts = {fmt for data in collected_data.values() for fmt in data[2]}
            plugboards_cache = {fmt: find_plugboard(plugboard_save_to_disk_value, fmt, self.plugboards) for fmt in all_fmts}
            self.pool = Pool(name='SaveToDisk') if self.pool is None else self.pool
            try:
                self.pool.set_common_data({
                    'plugboard_cache': plugboards_cache,
                    'template_functions': self.template_functions,
                    'library_id': self.library_id,
                })
            except Failure as err:
                error_dialog(
                    pd,
                    _('Critical failure'),
                    _('Could not save books to disk, click "Show details" for more information'),
                    det_msg=force_unicode(err.failure_message) + '\n' + force_unicode(err.details),
                    show=True,
                )
                pd.canceled = True
        self.do_one_signal.emit()

    def do_one_write(self):
        try:
            book_id = next(self.ids_to_write)
        except StopIteration:
            self.writing_finished()
            return
        if not self.opts.update_metadata:
            pd = self.pd
            assert pd is not None
            pd.msg = self.book_id_data(book_id).title
            pd.value += 1
        collected_data = self.collected_data
        assert collected_data is not None
        try:
            self.write_book(book_id, *collected_data[book_id])
        except Exception:
            self.errors[book_id].append(('critical', traceback.format_exc()))
        self.consume_results()
        self.do_one_signal.emit()

    def consume_results(self):
        if self.pool is not None:
            while True:
                try:
                    worker_result = self.pool.results.get_nowait()
                except Empty:
                    break
                book_id = worker_result.id
                pd = self.pd
                assert pd is not None
                if worker_result.is_terminal_failure:
                    error_dialog(
                        pd,
                        _('Critical failure'),
                        _('The update metadata worker process crashed while processing the book %s. Saving is aborted.') % self.book_id_data(book_id).title,
                        show=True,
                    )
                    pd.canceled = True
                    return
                result = worker_result.result
                pd.value += 1
                pd.msg = self.book_id_data(book_id).title
                if result.err is not None:
                    self.errors[book_id].append(('metadata', (None, result.err + '\n' + result.traceback)))
                if result.value:
                    for fmt, tb in result.value:
                        self.errors[book_id].append(('metadata', (fmt, tb)))

    def write_book(self, book_id, mi, components, fmts):
        db = self.db
        assert db is not None
        base_path = os.path.join(self.root, *components)
        base_dir = os.path.dirname(base_path)
        if self.opts.formats and self.opts.formats != 'all':
            if self.opts.formats == '..cover..':
                fmts = set()
            else:
                asked_formats = {x.lower().strip() for x in self.opts.formats.split(',')}
                fmts = asked_formats.intersection(fmts)
                if not fmts:
                    self.errors[book_id].append(('critical', _('Requested formats not available')))
                    return

        extra_files = {}
        if self.opts.save_extra_files:
            extra_files = {}
            for efx in db.new_api.list_extra_files(int(book_id), pattern=DATA_FILE_PATTERN):
                extra_files[efx.relpath] = efx.file_path
        if not fmts and not self.opts.write_opf and not self.opts.save_cover and not extra_files:
            return

        # On windows python incorrectly raises an access denied exception
        # when trying to create the root of a drive, like C:\
        if os.path.dirname(base_dir) != base_dir:
            try:
                os.makedirs(base_dir)
            except OSError as err:
                if err.errno != errno.EEXIST:
                    raise

        if self.opts.update_metadata:
            d = {}
            d['last_modified'] = mi.last_modified.isoformat()

        cdata = db.cover(book_id)
        mi.cover, mi.cover_data = None, (None, None)

        if cdata:
            fname = None
            if self.opts.save_cover:
                fname = base_path + os.extsep + 'jpg'
                mi.cover = os.path.basename(fname)
            elif self.opts.update_metadata:
                fname = os.path.join(self.tdir, f'{book_id}.jpg')

            if fname:
                with open(fname, 'wb') as f:
                    f.write(cdata)
                if self.opts.update_metadata:
                    d['cover'] = fname

        fname = None
        if self.opts.write_opf:
            fname = base_path + os.extsep + 'opf'
        elif self.opts.update_metadata:
            fname = os.path.join(self.tdir, f'{book_id}.opf')
        if fname:
            opf = metadata_to_opf(mi)
            with open(fname, 'wb') as f:
                f.write(opf)
            if self.opts.update_metadata:
                d['opf'] = fname
        mi.cover, mi.cover_data = None, (None, None)
        fmts_list: list[str] = []
        if self.opts.update_metadata:
            d['fmts'] = fmts_list
        if extra_files:
            for relpath, src_path in extra_files.items():
                src_path = make_long_path_useable(src_path)
                if os.access(src_path, os.R_OK):
                    dest = make_long_path_useable(os.path.abspath(os.path.join(base_dir, relpath)))
                    try:
                        shutil.copy2(src_path, dest)
                    except FileNotFoundError:
                        os.makedirs(os.path.dirname(dest), exist_ok=True)
                        shutil.copy2(src_path, dest)
        for fmt in fmts:
            try:
                fmtpath = self.write_fmt(book_id, fmt, base_path)
                if fmtpath and self.opts.update_metadata and can_set_metadata(fmt):
                    fmts_list.append(fmtpath)
            except Exception:
                self.errors[book_id].append(('fmt', (fmt, traceback.format_exc())))
        if self.opts.update_metadata:
            if fmts_list:
                pool = self.pool
                assert pool is not None
                pd = self.pd
                assert pd is not None
                try:
                    pool(book_id, 'calibre.library.save_to_disk', 'update_serialized_metadata', d)
                except Failure as err:
                    error_dialog(
                        pd,
                        _('Critical failure'),
                        _('Could not save books to disk, click "Show details" for more information'),
                        det_msg=str(err.failure_message) + '\n' + str(err.details),
                        show=True,
                    )
                    pd.canceled = True
            else:
                pd = self.pd
                assert pd is not None
                pd.value += 1
                pd.msg = self.book_id_data(book_id).title

    def write_fmt(self, book_id, fmt, base_path):
        db = self.db
        assert db is not None
        fmtpath = base_path + os.extsep + fmt
        written = False
        with open(fmtpath, 'w+b') as f:
            try:
                db.copy_format_to(book_id, fmt, f)
                written = True
            except NoSuchFormat:
                self.errors[book_id].append(('fmt', (fmt, _('No %s format file present') % fmt.upper())))
        if not written:
            os.remove(fmtpath)
        if written:
            return fmtpath

    def writing_finished(self):
        if not self.opts.update_metadata:
            self.updating_metadata_finished()
        else:
            self.do_one = self.do_one_update
            self.do_one_signal.emit()

    def do_one_update(self):
        self.consume_results()
        pool = self.pool
        assert pool is not None
        pd = self.pd
        assert pd is not None
        try:
            pool.wait_for_tasks(0.1)
        except Failure as err:
            error_dialog(
                pd,
                _('Critical failure'),
                _('Could not save books to disk, click "Show details" for more information'),
                det_msg=str(err.failure_message) + '\n' + str(err.details),
                show=True,
            )
            pd.canceled = True
        except RuntimeError:
            pass  # tasks not completed
        else:
            self.consume_results()
            return self.updating_metadata_finished()
        self.do_one_signal.emit()

    def updating_metadata_finished(self):
        if DEBUG:
            all_book_ids = self.all_book_ids
            assert all_book_ids is not None
            prints(f'Saved {len(all_book_ids)} books in {time.time() - self.start_time:.1f} seconds')
        pd = self.pd
        assert pd is not None
        pd.close()
        pd.deleteLater()
        self.report()
        self.break_cycles()
        self.on_complete()

    def on_complete(self):
        if gprefs['show_files_after_save']:
            open_local_file(self.root)

    def format_report(self):
        report = []
        a = report.append

        def indent(text):
            text = force_unicode(text)
            return '\xa0\xa0\xa0\xa0' + '\n\xa0\xa0\xa0\xa0'.join(text.splitlines())

        for book_id, errors in self.errors.items():
            types = {t for t, data in errors}
            title, authors = self.book_id_data(book_id).title, authors_to_string(self.book_id_data(book_id).authors[:1])
            if report:
                a('\n' + ('_' * 70) + '\n')
            if 'critical' in types:
                a(_('Failed to save: {0} by {1} to disk, with error:').format(title, authors))
                for t, tb in errors:
                    if t == 'critical':
                        a(indent(tb))
            else:
                errs = defaultdict(list)
                for t, data in errors:
                    errs[t].append(data)
                for fmt, tb in errs['fmt']:
                    a(_('Failed to save the {2} format of: {0} by {1} to disk, with error:').format(title, authors, fmt.upper()))
                    a(indent(tb)), a('')
                for fmt, tb in errs['metadata']:
                    if fmt:
                        a(_('Failed to update the metadata in the {2} format of: {0} by {1}, with error:').format(title, authors, fmt.upper()))
                    else:
                        a(_('Failed to update the metadata in all formats of: {0} by {1}, with error:').format(title, authors))
                    a(indent(tb)), a('')
        return '\n'.join(report)

    def report(self):
        if not self.errors:
            return
        all_book_ids = self.all_book_ids
        assert all_book_ids is not None
        err_types = {e[0] for errors in self.errors.values() for e in errors}
        if err_types == {'metadata'}:
            msg = _('Failed to update metadata in some books, click "Show details" for more information')
            d = warning_dialog
        elif len(self.errors) == len(all_book_ids):
            msg = _('Failed to save any books to disk, click "Show details" for more information')
            d = error_dialog
        else:
            msg = _('Failed to save some books to disk, click "Show details" for more information')
            d = warning_dialog
        d(self.parent(), _('Error while saving'), msg, det_msg=self.format_report(), show=True)
