"""Integrity checks for document/source filesystem-to-database reconciliation."""

from __future__ import annotations

from pathlib import Path
from uuid import UUID
from uuid import uuid4

import pytest
from sqlmodel import func
from sqlmodel import select

from transcription.config import Settings
from transcription.db import session_scope
from transcription.db.models import Document
from transcription.db.models import Source


def _document_folder_ids(root: Path) -> set[str]:
    documents_root = root / "documents"
    if not documents_root.exists():
        return set()
    return {entry.name for entry in documents_root.iterdir() if entry.is_dir()}


async def _document_ids(settings: Settings) -> set[str]:
    async with session_scope(settings=settings) as session:
        rows = await session.exec(select(Document.id))
        return {str(item) for item in rows.all()}


async def _source_counts_by_document(settings: Settings) -> dict[str, int]:
    async with session_scope(settings=settings) as session:
        rows = await session.exec(
            select(Source.document_id, func.count(Source.id)).group_by(Source.document_id)
        )
        return {str(document_id): int(count) for document_id, count in rows}


def _source_file_count_for_document(root: Path, document_id: str) -> int:
    directory = root / "documents" / document_id
    if not directory.exists():
        return 0
    return sum(1 for entry in directory.iterdir() if entry.is_file())


async def assert_storage_reconciliation(*, upload_dir: Path, settings: Settings) -> None:
    folder_ids = _document_folder_ids(upload_dir)
    doc_ids = await _document_ids(settings)

    missing_in_table = sorted(folder_ids - doc_ids)
    missing_in_folders = sorted(doc_ids - folder_ids)
    if missing_in_table:
        folder = missing_in_table[0]
        raise AssertionError(
            "Document directory count and document.doc_id count do not agree. "
            f"./data/documents/{folder} does not appear in document table"
        )
    if missing_in_folders:
        doc_id = missing_in_folders[0]
        raise AssertionError(
            "Document directory count and document.doc_id count do not agree. "
            f"document.doc_id {doc_id} has no corresponding folder in ./data/documents"
        )

    source_counts = await _source_counts_by_document(settings)
    for document_id in sorted(doc_ids):
        db_count = source_counts.get(document_id, 0)
        file_count = _source_file_count_for_document(upload_dir, document_id)
        if file_count > db_count:
            raise AssertionError(
                "Source file count and source.source_id count do not agree. "
                f"[UPLOAD_DIR]/documents/{document_id}/ contains file(s) with no source row"
            )
        if db_count > file_count:
            raise AssertionError(
                "Source file count and source.source_id count do not agree. "
                f"source.source_id rows exist without files in [UPLOAD_DIR]/documents/{document_id}"
            )


@pytest.mark.asyncio
async def test_storage_reconciliation_passes_for_matching_counts(tmp_path, default_settings: Settings):
    upload_dir = tmp_path / "uploads"
    document_id = str(uuid4())
    document_uuid = UUID(document_id)
    file_name = f"{uuid4()}.png"
    (upload_dir / "documents" / document_id).mkdir(parents=True, exist_ok=True)
    (upload_dir / "documents" / document_id / file_name).write_bytes(b"ok")

    async with session_scope(settings=default_settings) as session:
        session.add(Document(id=document_uuid, name="Doc"))
        session.add(
            Source(
                document_id=document_uuid,
                page_number=1,
                upload_name=file_name,
                filename=file_name,
                file_path=f"documents/{document_id}/{file_name}",
                file_hash="a" * 64,
                file_size_bytes=2,
            )
        )
        await session.commit()

    await assert_storage_reconciliation(upload_dir=upload_dir, settings=default_settings)


@pytest.mark.asyncio
async def test_storage_reconciliation_reports_actionable_mismatch_message(tmp_path, default_settings: Settings):
    upload_dir = tmp_path / "uploads"
    orphan_dir = upload_dir / "documents" / str(uuid4())
    orphan_dir.mkdir(parents=True, exist_ok=True)

    with pytest.raises(AssertionError, match="document\\.doc_id count do not agree"):
        await assert_storage_reconciliation(upload_dir=upload_dir, settings=default_settings)