import json
import pytest

from marker.converters.extraction import ExtractionConverter
from marker.extractors.page import PageExtractionSchema
from marker.extractors.document import DocumentExtractionSchema
from marker.services import BaseService


class MockLLMService(BaseService):
    def __call__(self, prompt, image=None, page=None, response_schema=None, **kwargs):
        if response_schema == PageExtractionSchema:
            return {
                "description": "Mock extraction description",
                "detailed_notes": "Mock detailed notes for page extraction",
            }
        elif response_schema == DocumentExtractionSchema:
            return {
                "analysis": "Mock document analysis",
                "document_json": json.dumps({"test_key": "test_value"}),
            }
        return {}


@pytest.fixture
def mock_llm_service():
    return MockLLMService


@pytest.fixture
def extraction_converter(config, model_dict, mock_llm_service):
    test_schema = {
        "title": "TestSchema",
        "type": "object",
        "properties": {"test_key": {"title": "Test Key", "type": "string"}},
        "required": ["test_key"],
    }

    config["page_schema"] = json.dumps(test_schema)
    config["output_format"] = "markdown"
    model_dict["llm_service"] = mock_llm_service

    converter = ExtractionConverter(
        artifact_dict=model_dict, processor_list=None, config=config
    )
    converter.llm_service = mock_llm_service
    converter.default_llm_service = MockLLMService
    return converter


@pytest.mark.config({"page_range": [0]})
def test_extraction_converter(config, model_dict, mock_llm_service, temp_doc):
    config["page_schema"] = "invalid json"

    model_dict["llm_service"] = mock_llm_service
    converter = ExtractionConverter(
        artifact_dict=model_dict, processor_list=None, config=config
    )
    converter.artifact_dict["llm_service"] = mock_llm_service()

    results = converter(temp_doc.name)
    assert results.document_json == '{"test_key": "test_value"}'


@pytest.mark.config({"page_range": [0, 1]})
def test_extraction_converter_multiple_pages(extraction_converter, temp_doc):
    result = extraction_converter(temp_doc.name)

    assert result is not None
    assert result.document_json is not None
    assert json.loads(result.document_json) == {"test_key": "test_value"}
    assert result.analysis == "Mock document analysis"