from __future__ import annotations

import hashlib
import io
from uuid import UUID

import pytest
from PIL import Image

from app.contracts import Attachment
from app.file_pipeline import ObjectChanged, collect_and_hash, detect_format


def image_bytes(format_name: str) -> bytes:
    output = io.BytesIO()
    Image.new("RGB", (2, 2), "white").save(output, format=format_name)
    return output.getvalue()


@pytest.mark.parametrize(
    "format_name,mime",
    [("JPEG", "image/jpeg"), ("PNG", "image/png"), ("WEBP", "image/webp")],
)
def test_bounded_image_detector(format_name: str, mime: str) -> None:
    assert detect_format(image_bytes(format_name), mime) is None
    assert detect_format(image_bytes(format_name), "application/pdf") == "file.format_mismatch"


def test_pdf_active_encrypted_and_malformed() -> None:
    clean = b"%PDF-1.7\n1 0 obj <<>> endobj\nstartxref\n0\n%%EOF"
    assert detect_format(clean, "application/pdf") is None
    assert (
        detect_format(clean.replace(b"<<>>", b"<</Encrypt 2 0 R>>"), "application/pdf")
        == "file.encrypted_content"
    )
    assert (
        detect_format(clean.replace(b"<<>>", b"<</JavaScript 2 0 R>>"), "application/pdf")
        == "file.active_content"
    )
    assert detect_format(b"%PDF-1.7 no eof", "application/pdf") == "file.polyglot_or_ambiguous"


class Reader:
    def __init__(self, data: bytes) -> None:
        self.data = data

    async def stream(self, attachment):
        yield self.data[:2]
        yield self.data[2:]


def attachment(data: bytes, *, size: int | None = None) -> Attachment:
    return Attachment(
        attachment_id=UUID("00000000-0000-4000-8000-000000000003"),
        quarantine_object_key=(
            "quarantine/users/00000000-0000-4000-8000-000000000001/"
            "dialogs/00000000-0000-4000-8000-000000000002/"
            "00000000-0000-4000-8000-000000000003"
        ),
        quarantine_version_id="v1",
        quarantine_etag='"etag"',
        mime_type="application/pdf",
        size_bytes=size or len(data),
        checksum="sha256:" + hashlib.sha256(data).hexdigest(),
    )


async def test_authoritative_stream_hash_and_size() -> None:
    data = b"content"
    body, digest = await collect_and_hash(Reader(data), attachment(data), max_size=100)
    assert body == data and digest == hashlib.sha256(data).digest()
    with pytest.raises(ObjectChanged):
        await collect_and_hash(Reader(data), attachment(data, size=len(data) + 1), max_size=100)