from pathlib import Path
from unittest.mock import patch

import pytest

from kiln_ai.adapters.extractors.base_extractor import (
    BaseExtractor,
    ExtractionInput,
    ExtractionOutput,
)
from kiln_ai.datamodel.extraction import (
    ExtractorConfig,
    ExtractorType,
    LitellmExtractorConfigProperties,
    OutputFormat,
)


class MockBaseExtractor(BaseExtractor):
    async def _extract(self, input: ExtractionInput) -> ExtractionOutput:
        return ExtractionOutput(
            is_passthrough=False,
            content="mock concrete extractor output",
            content_format=OutputFormat.MARKDOWN,
        )

    async def clear_cache_for_file_path(self, file_path: Path) -> None:
        pass


@pytest.fixture
def mock_litellm_properties():
    return {
        "extractor_type": ExtractorType.LITELLM,
        "prompt_document": "mock prompt for document",
        "prompt_image": "mock prompt for image",
        "prompt_video": "mock prompt for video",
        "prompt_audio": "mock prompt for audio",
    }


@pytest.fixture
def mock_extractor(mock_litellm_properties):
    return MockBaseExtractor(
        ExtractorConfig(
            name="mock",
            model_provider_name="gemini_api",
            model_name="gemini-2.0-flash",
            extractor_type=ExtractorType.LITELLM,
            output_format=OutputFormat.MARKDOWN,
            properties=mock_litellm_properties,
        )
    )


def mock_extractor_with_passthroughs(
    properties: LitellmExtractorConfigProperties,
    mimetypes: list[OutputFormat],
    output_format: OutputFormat,
):
    return MockBaseExtractor(
        ExtractorConfig(
            name="mock",
            model_provider_name="gemini_api",
            model_name="gemini-2.0-flash",
            extractor_type=ExtractorType.LITELLM,
            passthrough_mimetypes=mimetypes,
            output_format=output_format,
            properties=properties,
        )
    )


def test_should_passthrough(mock_litellm_properties):
    extractor = mock_extractor_with_passthroughs(
        mock_litellm_properties,
        [OutputFormat.TEXT, OutputFormat.MARKDOWN],
        OutputFormat.TEXT,
    )

    # should passthrough
    assert extractor._should_passthrough("text/plain")
    assert extractor._should_passthrough("text/markdown")

    # should not passthrough
    assert not extractor._should_passthrough("image/png")
    assert not extractor._should_passthrough("application/pdf")
    assert not extractor._should_passthrough("text/html")
    assert not extractor._should_passthrough("image/jpeg")


async def test_extract_passthrough(mock_litellm_properties):
    """
    Tests that when a file's MIME type is configured for passthrough, the extractor skips
    the concrete extraction method and returns the file's contents directly with the
    correct passthrough output format.
    """
    extractor = mock_extractor_with_passthroughs(
        mock_litellm_properties,
        [OutputFormat.TEXT, OutputFormat.MARKDOWN],
        OutputFormat.TEXT,
    )
    with (
        patch.object(
            extractor,
            "_extract",
            return_value=ExtractionOutput(
                is_passthrough=False,
                content="mock concrete extractor output",
                content_format=OutputFormat.TEXT,
            ),
        ) as mock_extract,
        patch(
            "pathlib.Path.read_text",
            return_value=b"test content",
        ),
    ):
        result = await extractor.extract(
            ExtractionInput(
                path="test.txt",
                mime_type="text/plain",
            )
        )

        # Verify _extract was not called
        mock_extract.assert_not_called()

        # Verify correct passthrough result
        assert result.is_passthrough
        assert result.content == "test content"
        assert result.content_format == OutputFormat.TEXT


@pytest.mark.parametrize(
    "output_format",
    [
        "text/plain",
        "text/markdown",
    ],
)
async def test_extract_passthrough_output_format(
    mock_litellm_properties, output_format
):
    extractor = mock_extractor_with_passthroughs(
        mock_litellm_properties,
        [OutputFormat.TEXT, OutputFormat.MARKDOWN],
        output_format,
    )
    with (
        patch.object(
            extractor,
            "_extract",
            return_value=ExtractionOutput(
                is_passthrough=False,
                content="mock concrete extractor output",
                content_format=output_format,
            ),
        ) as mock_extract,
        patch(
            "pathlib.Path.read_text",
            return_value="test content",
        ),
    ):
        result = await extractor.extract(
            ExtractionInput(
                path="test.txt",
                mime_type="text/plain",
            )
        )

        # Verify _extract was not called
        mock_extract.assert_not_called()

        # Verify correct passthrough result
        assert result.is_passthrough
        assert result.content == "test content"
        assert result.content_format == output_format


@pytest.mark.parametrize(
    "path, mime_type, output_format",
    [
        ("test.mp3", "audio/mpeg", OutputFormat.TEXT),
        ("test.png", "image/png", OutputFormat.TEXT),
        ("test.pdf", "application/pdf", OutputFormat.TEXT),
        ("test.txt", "text/plain", OutputFormat.MARKDOWN),
        ("test.txt", "text/markdown", OutputFormat.MARKDOWN),
        ("test.html", "text/html", OutputFormat.MARKDOWN),
    ],
)
async def test_extract_non_passthrough(
    mock_extractor, path: str, mime_type: str, output_format: OutputFormat
):
    with (
        patch.object(
            mock_extractor,
            "_extract",
            return_value=ExtractionOutput(
                is_passthrough=False,
                content="mock concrete extractor output",
                content_format=output_format,
            ),
        ) as mock_extract,
    ):
        # first we call the base class extract method
        result = await mock_extractor.extract(
            ExtractionInput(
                path=path,
                mime_type=mime_type,
            )
        )

        # then we call the subclass _extract method and add validated mime_type
        mock_extract.assert_called_once_with(
            ExtractionInput(
                path=path,
                mime_type=mime_type,
            )
        )

        assert not result.is_passthrough
        assert result.content == "mock concrete extractor output"
        assert result.content_format == output_format


async def test_default_output_format(mock_litellm_properties):
    config = ExtractorConfig(
        name="mock",
        model_provider_name="gemini_api",
        model_name="gemini-2.0-flash",
        extractor_type=ExtractorType.LITELLM,
        properties=mock_litellm_properties,
    )
    assert config.output_format == OutputFormat.MARKDOWN


async def test_extract_failure_from_concrete_extractor(mock_extractor):
    with patch.object(
        mock_extractor,
        "_extract",
        side_effect=Exception("error from concrete extractor"),
    ):
        with pytest.raises(ValueError, match="error from concrete extractor"):
            await mock_extractor.extract(
                ExtractionInput(
                    path="test.txt",
                    mime_type="text/plain",
                )
            )


async def test_output_format(mock_extractor):
    assert mock_extractor.output_format() == OutputFormat.MARKDOWN
