transcription/tests/services/test_normalization.py

"""Tests for V4.5 metadata-directed orientation normalization."""

import hashlib
import io
from pathlib import Path

import pytest
from PIL import Image

from transcription.config import Settings
from transcription.db.models import Document
from transcription.db.models import Job
from transcription.db.models import JobSource
from transcription.db.models import Source
from transcription.providers import RequestManifest
from transcription.providers import TranscriptionResult
from transcription.providers.evidence import SourceEvidenceReference
from transcription.providers.evidence import build_software_context
from transcription.services import ServiceBundle
from transcription.services.documents import DocumentService
from transcription.services.jobs import JobService
from transcription.services.normalization import normalize_orientation
from transcription.services.normalization import normalize_orientation_async
from transcription.services.sources import SourceService
from transcription.services.workflows import process_queued_job


def _write_oriented_jpeg(path: Path, *, orientation: int) -> bytes:
    image = Image.new("RGB", (2, 3))
    image.putdata(
        [
            (255, 0, 0),
            (255, 0, 0),
            (0, 255, 0),
            (0, 255, 0),
            (0, 0, 255),
            (0, 0, 255),
        ]
    )
    exif = Image.Exif()
    exif[274] = orientation
    image.save(path, format="JPEG", quality=100, subsampling=0, exif=exif)
    return path.read_bytes()


def _write_oriented_image(path: Path, *, orientation: int, image_format: str) -> None:
    image = Image.new("RGB", (2, 3), color="white")
    exif = Image.Exif()
    exif[274] = orientation
    image.save(path, format=image_format, exif=exif)


@pytest.mark.unit
def test_orientation_three_is_physically_rotated_and_metadata_removed(tmp_path):
    path = tmp_path / "upside-down.jpg"
    original = _write_oriented_jpeg(path, orientation=3)

    result = normalize_orientation(path, media_type="image/jpeg")

    assert result is not None
    assert result.applied_rotation_degrees == 180
    assert path.read_bytes() == original
    with Image.open(path) as source_image, Image.open(io.BytesIO(result.content)) as derivative:
        assert source_image.getexif()[274] == 3
        assert derivative.getexif().get(274, 1) == 1
        pixel = derivative.getpixel((0, 0))
        assert isinstance(pixel, tuple)
        assert pixel[2] > pixel[0]


@pytest.mark.unit
def test_orientation_one_is_noop(tmp_path):
    path = tmp_path / "upright.jpg"
    _write_oriented_jpeg(path, orientation=1)

    assert normalize_orientation(path, media_type="image/jpeg") is None


@pytest.mark.parametrize(
    ("filename", "image_format", "media_type"),
    [
        ("oriented.png", "PNG", "image/png"),
        ("oriented.tiff", "TIFF", "image/tiff"),
    ],
)
def test_supported_non_jpeg_orientation_is_normalized(
    tmp_path,
    filename,
    image_format,
    media_type,
):
    path = tmp_path / filename
    _write_oriented_image(path, orientation=6, image_format=image_format)

    result = normalize_orientation(path, media_type=media_type)

    assert result is not None
    assert result.applied_rotation_degrees == 90
    assert (result.derivative_width, result.derivative_height) == (3, 2)
    with Image.open(io.BytesIO(result.content)) as derivative:
        assert derivative.getexif().get(274, 1) == 1


@pytest.mark.integration
@pytest.mark.asyncio
async def test_resolve_provider_input_persists_exact_derivative(default_session_factory, tmp_path):
    source_path = tmp_path / "source.jpg"
    original = _write_oriented_jpeg(source_path, orientation=3)
    settings = Settings(
        openrouter_api_key="test-key",
        artifact_dir=tmp_path / "artifacts",
        provider_models=None,
    )
    documents = DocumentService(session_factory=default_session_factory, settings=settings)
    sources = SourceService(session_factory=default_session_factory, settings=settings)
    document = await documents.create_document(Document(name="Oriented"))
    source = await sources.create_source(
        Source(
            document_id=document.id,
            page_number=1,
            upload_name="source.jpg",
            filename="source.jpg",
            file_path=str(source_path),
            file_hash=hashlib.sha256(original).hexdigest(),
            file_size_bytes=len(original),
        )
    )

    provider_input = await sources.resolve_provider_input(source)
    artifacts = await sources.list_processing_artifacts(source_id=source.id)

    assert source_path.read_bytes() == original
    assert provider_input.derivative_id == artifacts[0].id
    assert provider_input.path.read_bytes() != original
    assert hashlib.sha256(provider_input.path.read_bytes()).hexdigest() == provider_input.digest_sha256
    coordinate_metadata = artifacts[0].coordinate_metadata
    assert coordinate_metadata is not None
    assert coordinate_metadata["original_orientation"] == 3


@pytest.mark.integration
@pytest.mark.asyncio
async def test_worker_sends_exact_derivative_and_links_attempt_evidence(
    default_session_factory,
    tmp_path,
    monkeypatch,
):
    source_path = tmp_path / "source.jpg"
    original = _write_oriented_jpeg(source_path, orientation=3)
    prompt_dir = tmp_path / "prompts"
    prompt_dir.mkdir()
    (prompt_dir / "transcribe_document.md").write_text("Transcribe verbatim.", encoding="utf-8")
    settings = Settings(
        openrouter_api_key="test-key",
        artifact_dir=tmp_path / "artifacts",
        prompt_dir=prompt_dir,
        provider_models=None,
    )
    services = ServiceBundle(
        documents=DocumentService(session_factory=default_session_factory, settings=settings),
        jobs=JobService(session_factory=default_session_factory, settings=settings),
        sources=SourceService(session_factory=default_session_factory, settings=settings),
    )
    document = await services.documents.create_document(Document(name="Pipeline"))
    source = await services.sources.create_source(
        Source(
            document_id=document.id,
            page_number=1,
            upload_name="source.jpg",
            filename="source.jpg",
            file_path=str(source_path),
            file_hash=hashlib.sha256(original).hexdigest(),
            file_size_bytes=len(original),
        )
    )
    job = await services.jobs.create_job(Job(document_id=document.id))
    await services.sources.create_job_source(JobSource(job_id=job.id, source_id=source.id))
    loaded = await services.jobs.read_job(job.id)
    captured: dict[str, object] = {}

    async def fake_transcribe(
        image_path,
        *,
        prompt_name,
        prompt_text,
        temperature,
        top_p,
        settings,
        provider,
        source_reference,
        requested_model,
    ):
        _ = (prompt_name, temperature, top_p, settings, provider, requested_model)
        image_bytes = Path(image_path).read_bytes()
        captured["bytes"] = image_bytes
        captured["source_reference"] = source_reference
        manifest = RequestManifest(
            provider="fixture",
            requested_model="fixture/model",
            request={"model": "fixture/model"},
            source=source_reference,
            optional_parameter_states={"temperature": "omitted", "top_p": "omitted"},
            prompt_content=prompt_text,
            prompt_sha256=hashlib.sha256(prompt_text.encode()).hexdigest(),
            timeout_seconds=20,
            retry_policy="none",
            software=build_software_context(
                adapter_name="fixture",
                adapter_version="1",
                client_library="transcription",
            ),
        )
        return TranscriptionResult(
            text="[document body typewritten]\nDamaged \ufffd text",
            provider="fixture",
            model="fixture/model",
            request_manifest=manifest,
        )

    monkeypatch.setattr("transcription.services.workflows.transcribe_document_image", fake_transcribe)

    await process_queued_job(job=loaded, services=services, settings=settings)

    attempts = await services.sources.list_execution_attempts(source_id=source.id)
    artifacts = await services.sources.list_processing_artifacts(source_id=source.id)
    source_reference = captured["source_reference"]
    assert isinstance(source_reference, SourceEvidenceReference)
    captured_bytes = captured["bytes"]
    assert isinstance(captured_bytes, bytes)
    assert source_path.read_bytes() == original
    assert hashlib.sha256(captured_bytes).hexdigest() == source_reference.digest_sha256
    assert source_reference.derivative_id is not None
    assert {artifact.artifact_type for artifact in artifacts} == {
        "orientation_normalized_model_input",
        "transcription_quality_warnings",
    }
    assert {artifact.execution_attempt_id for artifact in artifacts} == {attempts[0].id}


@pytest.mark.asyncio
async def test_async_wrapper_matches_sync_result_and_precomputes_digest(tmp_path):
    """[MED-01]: Pillow work runs off the event loop and hashes its own output."""
    path = tmp_path / "async-upside-down.jpg"
    _write_oriented_jpeg(path, orientation=3)

    result = await normalize_orientation_async(path, media_type="image/jpeg")
    expected = normalize_orientation(path, media_type="image/jpeg")

    assert result is not None
    assert expected is not None
    assert result.content == expected.content
    assert result.digest_sha256 == hashlib.sha256(result.content).hexdigest()