"""Tests for the doc_modeler."""
import asyncio
import shutil
import sqlite3
from unittest.mock import patch

import pytest
from PIL import Image
from sqlalchemy.exc import OperationalError

from modules.docs import doc_modeler
from modules.docs.models import Doc, DocSection
from wrolpi.files import lib as files_lib
from wrolpi.files.models import FileGroup
from wrolpi.files.worker import file_worker
from wrolpi.test.common import assert_dict_contains
from wrolpi.vars import PROJECT_DIR


@pytest.mark.asyncio
async def test_doc_modeler_epub(async_client, test_session, test_directory, example_epub, example_mobi, refresh_files):
    """Doc modeler indexes epub files."""
    await refresh_files()

    doc: Doc = test_session.query(Doc).one()
    assert doc.file_group.title == 'WROLPi Test Book'
    assert doc.file_group.mimetype.startswith('application/epub')
    assert doc.file_group.author == 'roland'
    assert doc.file_group.a_text, 'Book title was not indexed'
    assert doc.file_group.b_text, 'Book author was not indexed'
    assert doc.file_group.d_text, 'Book text was not indexed'
    assert doc.file_group.model == 'doc'

    # Both ebooks are assumed to be the same book.
    assert (epubs := doc.file_group.my_files('application/epub+zip')) and len(epubs) == 1
    assert (mobis := doc.file_group.my_files('application/x-mobipocket-ebook')) and len(mobis) == 1
    # Cover was discovered.
    assert len(doc.file_group.my_poster_files()) == 1

    # A second refresh should not re-model the already-indexed doc.
    with patch('modules.docs._model_doc', side_effect=AssertionError('_model_doc should not be called')):
        await refresh_files()


@pytest.mark.asyncio
async def test_doc_modeler_pdf(async_client, test_session, test_directory, example_pdf, refresh_files):
    """Doc modeler indexes PDF files."""
    await refresh_files()

    doc: Doc = test_session.query(Doc).one()
    assert doc.file_group.mimetype == 'application/pdf'
    assert doc.file_group.model == 'doc'
    assert doc.file_group.indexed is True
    assert doc.size  # Has a size.


@pytest.mark.asyncio
async def test_doc_modeler_persists_epub_sections(async_client, test_session, test_directory,
                                                  example_epub, refresh_files):
    """After modeling, per-spine-item DocSection rows exist for an EPUB."""
    await refresh_files()

    doc: Doc = test_session.query(Doc).one()
    sections = test_session.query(DocSection).filter_by(doc_id=doc.id) \
        .order_by(DocSection.ordinal).all()
    assert sections, 'EPUB modeler should produce DocSection rows'
    assert all(s.kind == 'epub_spine' for s in sections)
    assert [s.ordinal for s in sections] == list(range(len(sections)))


@pytest.mark.asyncio
async def test_doc_modeler_persists_pdf_sections(async_client, test_session, test_directory,
                                                 example_pdf, refresh_files):
    """After modeling, per-page DocSection rows exist for a PDF."""
    await refresh_files()

    doc: Doc = test_session.query(Doc).one()
    sections = test_session.query(DocSection).filter_by(doc_id=doc.id) \
        .order_by(DocSection.ordinal).all()
    assert sections, 'PDF modeler should produce DocSection rows'
    assert all(s.kind == 'pdf_page' for s in sections)
    # 1-based ordinals, test PDF has 3 pages.
    assert [s.ordinal for s in sections] == [1, 2, 3]


@pytest.mark.asyncio
async def test_doc_modeler_replaces_sections(async_client, test_session, test_directory,
                                             example_epub, refresh_files):
    """Re-modeling a Doc replaces its DocSection rows (no duplicates)."""
    from modules.docs import _model_doc

    await refresh_files()
    doc: Doc = test_session.query(Doc).one()
    initial_count = test_session.query(DocSection).filter_by(doc_id=doc.id).count()
    assert initial_count > 0

    # Re-run modeling on the same file_group and confirm we don't accumulate rows.
    _model_doc(doc.file_group, test_session)
    test_session.commit()
    final_count = test_session.query(DocSection).filter_by(doc_id=doc.id).count()
    assert final_count == initial_count


@pytest.mark.asyncio
async def test_doc_do_model_sets_model(test_session, example_pdf):
    """Doc.do_model sets file_group.model so the FileGroup is recognized as a doc."""
    files_lib._upsert_files([example_pdf])
    file_group = test_session.query(FileGroup).one()
    # Reset to simulate a newly discovered file.
    file_group.model = None
    file_group.indexed = False
    test_session.commit()

    file_group.do_model(test_session)
    test_session.commit()

    assert file_group.model == 'doc', 'do_model must set file_group.model'
    assert file_group.indexed is True


@pytest.mark.asyncio
async def test_doc_modeler_models_indexed_filegroup_without_doc(async_client, test_session, test_directory,
                                                               example_pdf):
    """A doc-mimetype FileGroup that was already indexed (indexed=True) but never modeled (no Doc
    row) must still be picked up by the doc_modeler.

    Regression: apply_indexers marks every FileGroup indexed=True.  If a doc file gets indexed
    before it is modeled (e.g. the modeler errored, or the file was imported already-indexed), the
    old `indexed == False` gate excluded it forever -- leaving thousands of docs invisible in
    /api/docs even though they exist on disk.
    """
    # Simulate the stuck production state: the FileGroup has been swept by apply_indexers
    # (indexed=True) but has no Doc row and no model.
    file_group = FileGroup.from_paths(test_session, example_pdf)
    file_group.indexed = True
    file_group.model = None
    test_session.commit()
    assert test_session.query(Doc).count() == 0, 'Precondition: no Doc row exists yet'

    await doc_modeler()

    doc: Doc = test_session.query(Doc).one()
    assert doc.file_group_id == file_group.id
    assert doc.file_group.model == 'doc'


@pytest.mark.asyncio
async def test_doc_modeler_terminates_on_persistent_failure(async_client, test_session, test_directory):
    """doc_modeler must terminate even when modeling consistently fails and never creates a Doc row.

    Regression: the `Doc.id IS NULL` gate keeps matching a doc that fails to model.  Without a
    per-run guard, a batch of >=DOC_PROCESSING_LIMIT permanently-failing docs is re-selected every
    iteration, so the loop never breaks and later docs are starved.
    """
    ebook_dir = test_directory / 'ebooks'
    ebook_dir.mkdir(parents=True)
    num_docs = 15  # More than one batch (DOC_PROCESSING_LIMIT == 10).
    for i in range(num_docs):
        path = ebook_dir / f'test_ebook_{i:03d}.epub'
        shutil.copy(PROJECT_DIR / 'test/ebook example.epub', path)
    for path in ebook_dir.iterdir():
        fg = FileGroup.from_paths(test_session, path)
        fg.indexed = True  # Simulate apply_indexers having already run.
    test_session.commit()

    # _model_doc always raises without persisting a Doc row (patch replaces the whole function).
    # PYTEST re-raises inside the loop's per-item handler; disable that so we exercise the failure
    # path and the loop's termination guard, not the re-raise.
    with patch('modules.docs._model_doc', side_effect=RuntimeError('boom')), \
            patch('wrolpi.files.modeler.PYTEST', False):
        await asyncio.wait_for(doc_modeler(), timeout=10)

    # No Doc rows were created (every attempt failed), but the loop still terminated.
    assert test_session.query(Doc).count() == 0


@pytest.mark.asyncio
async def test_doc_modeler_survives_database_locked(async_client, test_session, test_directory, refresh_files):
    """A transient SQLite "database is locked" on one doc must not abort the whole modeler run.

    Regression: doc_modeler held one write transaction across a whole batch of slow extractions.
    When another connection committed mid-batch, the deferred transaction failed to upgrade with
    SQLITE_BUSY_SNAPSHOT ("database is locked"), the error handler crashed on the rolled-back
    session, and the entire run aborted -- so no further docs were modeled.  Each doc now runs in
    its own short transaction and a lock is retried.
    """
    ebook_dir = test_directory / 'ebooks'
    ebook_dir.mkdir(parents=True)
    for i in range(3):
        shutil.copy(PROJECT_DIR / 'test/ebook example.epub', ebook_dir / f'book_{i:03d}.epub')

    from modules.docs import _model_doc as real_model_doc

    calls = {'n': 0}

    def flaky_model_doc(file_group, session):
        # Fail the very first modeling attempt with a lock error, then behave normally.  Proves the
        # run recovers (retries the doc, keeps modeling the rest) instead of aborting.
        calls['n'] += 1
        if calls['n'] == 1:
            raise OperationalError('stmt', {}, sqlite3.OperationalError('database is locked'))
        return real_model_doc(file_group, session)

    with patch('modules.docs._model_doc', side_effect=flaky_model_doc):
        await refresh_files()

    # All three ebooks were modeled despite the initial lock error.
    assert test_session.query(Doc).count() == 3
    assert calls['n'] >= 4  # 1 failed attempt + 1 retry + 2 more docs.


@pytest.mark.asyncio
async def test_doc_modeler_processes_more_than_batch_limit(async_client, test_session, test_directory):
    """Doc modeler processes more than one batch of files (limit=10)."""
    ebook_dir = test_directory / 'ebooks'
    ebook_dir.mkdir(parents=True)

    num_docs = 15
    for i in range(num_docs):
        path = ebook_dir / f'test_ebook_{i:03d}.epub'
        shutil.copy(PROJECT_DIR / 'test/ebook example.epub', path)

    # Create FileGroups.
    for path in ebook_dir.iterdir():
        fg = FileGroup.from_paths(test_session, path)
        assert fg.mimetype.startswith('application/epub')
    test_session.commit()

    await doc_modeler()

    doc_count = test_session.query(Doc).count()
    assert doc_count == num_docs, \
        f'doc_modeler should process ALL {num_docs} files, but only processed {doc_count}'

    still_unindexed = test_session.query(FileGroup).filter(
        FileGroup.indexed == False,
        FileGroup.mimetype == 'application/epub+zip',
    ).count()
    assert still_unindexed == 0


@pytest.mark.asyncio
async def test_doc_modeler_reindex_creates_collections(async_client, test_session, test_directory, example_epub,
                                                       refresh_files):
    """When a doc is re-indexed (indexed=false), author/subject collections should be re-created."""
    from wrolpi.collections.models import Collection

    from wrolpi.db import get_db_session

    await refresh_files()

    # Author collection was created during first modeling.
    with get_db_session() as session:
        doc = session.query(Doc).one()
        assert doc.file_group.author == 'roland'
        author_col = session.query(Collection).filter_by(kind='author', name='Roland').one_or_none()
        assert author_col, 'Author collection should exist after first indexing'

    # Simulate post-migration state: delete the collection and mark as unindexed.
    with get_db_session(commit=True) as session:
        author_col = session.query(Collection).filter_by(kind='author', name='Roland').one()
        session.delete(author_col)
        doc = session.query(Doc).one()
        doc.file_group.indexed = False

    # Verify collection is gone.
    with get_db_session() as session:
        assert session.query(Collection).filter_by(kind='author', name='Roland').one_or_none() is None

    # Re-run the modeler.
    await doc_modeler()

    # Author collection should be re-created.
    with get_db_session() as session:
        author_col = session.query(Collection).filter_by(kind='author', name='Roland').one_or_none()
        assert author_col, 'Author collection should be re-created when doc is re-indexed'


@pytest.mark.asyncio
async def test_doc_modeler_calibre_cover(test_session, async_client, test_directory, example_epub, example_mobi,
                                         image_file, refresh_files):
    """Calibre cover near an ebook file is discovered."""
    metadata = test_directory / 'metadata.opf'
    metadata.touch()
    cover_image = test_directory / 'cover.jpg'
    shutil.move(image_file, cover_image)

    await refresh_files()
    doc: Doc = test_session.query(Doc).one()
    assert doc.file_group.title == 'WROLPi Test Book'
    assert doc.cover_path
    assert {i['path'].name for i in doc.file_group.my_files()} == {'ebook example.epub', 'ebook example.mobi',
                                                                   'cover.jpg'}
    # Metadata and cover FileGroups were deleted.
    assert test_session.query(FileGroup).count() == 1


@pytest.mark.asyncio
async def test_doc_modeler_normalizes_author(async_client, test_session, test_directory, example_epub, refresh_files):
    """Author names are normalized before creating collections (e.g. quotes stripped)."""
    from wrolpi.collections.models import Collection
    from wrolpi.db import get_db_session

    # Patch extract_metadata to return a quoted author.
    from modules.docs.extractors import DocMetadata
    fake_metadata = DocMetadata(author='"Tony R. Kuphaldt"', title='Test Book')
    with patch('modules.docs.extract_metadata', return_value=fake_metadata):
        await refresh_files()

    with get_db_session() as session:
        # Collection should have normalized name (no quotes).
        col = session.query(Collection).filter_by(kind='author', name='Tony R. Kuphaldt').one_or_none()
        assert col, 'Author collection should have normalized name without quotes'
        # No collection with the quoted name.
        quoted = session.query(Collection).filter_by(kind='author', name='"Tony R. Kuphaldt"').one_or_none()
        assert not quoted, 'Should not create collection with quoted author name'


@pytest.mark.asyncio
async def test_doc_modeler_validates_subject(async_client, test_session, test_directory, example_epub, refresh_files):
    """Junk subjects are rejected during modeling."""
    from wrolpi.collections.models import Collection
    from wrolpi.db import get_db_session

    from modules.docs.extractors import DocMetadata
    fake_metadata = DocMetadata(subject='01/2004', title='Test Book')
    with patch('modules.docs.extract_metadata', return_value=fake_metadata):
        await refresh_files()

    with get_db_session() as session:
        col = session.query(Collection).filter_by(kind='subject', name='01/2004').one_or_none()
        assert not col, 'Should not create collection for date-like subject'


@pytest.mark.asyncio
async def test_doc_modeler_discover_local_cover(test_session, test_directory, example_epub, image_bytes_factory,
                                                refresh_files):
    """A cover file near the eBook is discovered."""
    cover_path = example_epub.with_suffix('.jpg')
    cover_path.write_bytes(image_bytes_factory())
    await refresh_files()

    doc: Doc = test_session.query(Doc).one()
    assert doc.cover_path.read_bytes() == cover_path.read_bytes()


@pytest.mark.asyncio
async def test_doc_modeler_extract_cover(test_session, test_directory, example_epub, refresh_files):
    """First image is extracted from the Ebook and used as the cover."""
    await refresh_files()

    doc: Doc = test_session.query(Doc).one()
    # The WROLPi logo is the first image in the example epub.
    assert doc.cover_path and doc.cover_path.stat().st_size == 297099


@pytest.mark.asyncio
async def test_doc_modeler_move_ebook(async_client, test_session, test_directory, example_epub, image_file, tag_factory,
                                      refresh_files, await_background_tasks):
    """An ebook is re-indexed when moved."""
    tag = await tag_factory()
    shutil.move(image_file, example_epub.with_suffix('.jpg'))
    await refresh_files()

    doc: Doc = test_session.query(Doc).one()
    assert doc.cover_path
    assert_dict_contains(doc.file_group.data, dict(author='roland'))
    assert sorted([i['path'] for i in doc.file_group.my_files()]) == [
        test_directory / 'ebook example.epub',
        test_directory / 'ebook example.jpg',
    ]
    doc.file_group.add_tag(test_session, tag.id)
    test_session.commit()

    new_directory = test_directory / 'new'
    new_directory.mkdir()
    file_worker.queue_move(new_directory, [doc.file_group.primary_path])
    await await_background_tasks()
    test_session.expire_all()

    doc: Doc = test_session.query(Doc).one()
    assert doc.file_group.primary_path == test_directory / 'new/ebook example.epub'
    assert doc.cover_path
    assert_dict_contains(doc.file_group.data, dict(author='roland'))
    assert sorted([i['path'] for i in doc.file_group.my_files()]) == [
        test_directory / 'new/ebook example.epub',
        test_directory / 'new/ebook example.jpg',
    ]


@pytest.mark.asyncio
async def test_doc_modeler_pdf_max_size(test_session, example_pdf):
    """The contents of a large PDF are not indexed."""
    example_pdf.write_bytes(example_pdf.read_bytes() * 5000)

    files_lib._upsert_files([example_pdf])

    await doc_modeler()
    file_group = test_session.query(FileGroup).one()
    assert file_group.d_text is None


@pytest.mark.asyncio
async def test_doc_modeler_pdf_poster(test_session, example_pdf):
    """PDF modeler finds an image with the same stem and uses it as the poster."""
    poster_path = example_pdf.with_suffix('.jpg')
    Image.new('RGB', (25, 25), color='grey').save(poster_path)

    files_lib._upsert_files([example_pdf, poster_path])

    await doc_modeler()
    file_group = test_session.query(FileGroup).one()
    assert file_group.indexed is True
    assert file_group.data is not None
    assert file_group.data.get('cover_path') == poster_path.name