#
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
#

import asyncio
from datetime import datetime
from unittest import mock
from unittest.mock import MagicMock, call, mock_open, patch

import pytest
import requests
from unstructured.documents.elements import ElementMetadata, Formula, ListItem, Text, Title
from unstructured.file_utils.filetype import FileType

from airbyte_cdk.models import FailureType
from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
from airbyte_cdk.sources.file_based.config.unstructured_format import (
    APIParameterConfigModel,
    APIProcessingConfigModel,
    UnstructuredFormat,
)
from airbyte_cdk.sources.file_based.exceptions import RecordParseError
from airbyte_cdk.sources.file_based.file_types import UnstructuredParser
from airbyte_cdk.sources.file_based.remote_file import RemoteFile
from airbyte_cdk.utils.traced_exception import AirbyteTracedException

FILE_URI = "path/to/file.xyz"


@pytest.mark.parametrize(
    "filetype, format_config, raises",
    [
        pytest.param(
            FileType.MD,
            UnstructuredFormat(skip_unprocessable_files=False),
            False,
            id="markdown_file",
        ),
        pytest.param(
            FileType.CSV,
            UnstructuredFormat(skip_unprocessable_files=False),
            True,
            id="wrong_file_format",
        ),
        pytest.param(
            FileType.CSV,
            UnstructuredFormat(skip_unprocessable_files=True),
            False,
            id="wrong_file_format_skipping",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(skip_unprocessable_files=False),
            False,
            id="pdf_file",
        ),
        pytest.param(
            FileType.DOCX,
            UnstructuredFormat(skip_unprocessable_files=False),
            False,
            id="docx_file",
        ),
        pytest.param(
            FileType.PPTX,
            UnstructuredFormat(skip_unprocessable_files=False),
            False,
            id="pptx_file",
        ),
    ],
)
@patch("airbyte_cdk.sources.file_based.file_types.unstructured_parser.detect_filetype")
def test_infer_schema(mock_detect_filetype, filetype, format_config, raises):
    # use a fresh event loop to avoid leaking into other tests
    main_loop = asyncio.get_event_loop()
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    stream_reader = MagicMock()
    mock_open(stream_reader.open_file)
    fake_file = MagicMock()
    fake_file.uri = FILE_URI
    logger = MagicMock()
    mock_detect_filetype.return_value = filetype
    config = MagicMock()
    config.format = format_config
    if raises:
        with pytest.raises(RecordParseError):
            loop.run_until_complete(
                UnstructuredParser().infer_schema(config, fake_file, stream_reader, logger)
            )
    else:
        schema = loop.run_until_complete(
            UnstructuredParser().infer_schema(config, MagicMock(), MagicMock(), MagicMock())
        )
        assert schema == {
            "content": {
                "type": "string",
                "description": "Content of the file as markdown. Might be null if the file could not be parsed",
            },
            "document_key": {
                "type": "string",
                "description": "Unique identifier of the document, e.g. the file path",
            },
            "_ab_source_file_parse_error": {
                "type": "string",
                "description": "Error message if the file could not be parsed even though the file is supported",
            },
        }
    loop.close()
    asyncio.set_event_loop(main_loop)


@pytest.mark.parametrize(
    "filetype, format_config, parse_result, raises, expected_records, parsing_error",
    [
        pytest.param(
            FileType.MD,
            UnstructuredFormat(skip_unprocessable_files=False),
            "test",
            False,
            [
                {
                    "content": "test",
                    "document_key": FILE_URI,
                    "_ab_source_file_parse_error": None,
                }
            ],
            False,
            id="markdown_file",
        ),
        pytest.param(
            FileType.CSV,
            UnstructuredFormat(skip_unprocessable_files=False),
            None,
            True,
            None,
            False,
            id="wrong_file_format",
        ),
        pytest.param(
            FileType.CSV,
            UnstructuredFormat(skip_unprocessable_files=True),
            None,
            False,
            [
                {
                    "content": None,
                    "document_key": FILE_URI,
                    "_ab_source_file_parse_error": "Error parsing record. This could be due to a mismatch between the config's file type and the actual file type, or because the file or record is not parseable. Contact Support if you need assistance.\nfilename=path/to/file.xyz message=File type FileType.CSV is not supported. Supported file types are FileType.MD, FileType.PDF, FileType.DOCX, FileType.PPTX, FileType.TXT",
                }
            ],
            False,
            id="skip_unprocessable_files",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(skip_unprocessable_files=False),
            [
                Title("heading"),
                Text("This is the text"),
                ListItem("This is a list item"),
                Formula("This is a formula"),
            ],
            False,
            [
                {
                    "content": "# heading\n\nThis is the text\n\n- This is a list item\n\n```\nThis is a formula\n```",
                    "document_key": FILE_URI,
                    "_ab_source_file_parse_error": None,
                }
            ],
            False,
            id="pdf_file",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(skip_unprocessable_files=False),
            [
                Title("first level heading", metadata=ElementMetadata(category_depth=1)),
                Title("second level heading", metadata=ElementMetadata(category_depth=2)),
            ],
            False,
            [
                {
                    "content": "# first level heading\n\n## second level heading",
                    "document_key": FILE_URI,
                    "_ab_source_file_parse_error": None,
                }
            ],
            False,
            id="multi_level_headings",
        ),
        pytest.param(
            FileType.DOCX,
            UnstructuredFormat(skip_unprocessable_files=False),
            [
                Title("heading"),
                Text("This is the text"),
                ListItem("This is a list item"),
                Formula("This is a formula"),
            ],
            False,
            [
                {
                    "content": "# heading\n\nThis is the text\n\n- This is a list item\n\n```\nThis is a formula\n```",
                    "document_key": FILE_URI,
                    "_ab_source_file_parse_error": None,
                }
            ],
            False,
            id="docx_file",
        ),
        pytest.param(
            FileType.DOCX,
            UnstructuredFormat(skip_unprocessable_files=True),
            "",
            False,
            [
                {
                    "content": None,
                    "document_key": FILE_URI,
                    "_ab_source_file_parse_error": "Error parsing record. This could be due to a mismatch between the config's file type and the actual file type, or because the file or record is not parseable. Contact Support if you need assistance.\nfilename=path/to/file.xyz message=weird parsing error",
                }
            ],
            True,
            id="exception_during_parsing",
        ),
    ],
)
@patch("unstructured.partition.pdf.partition_pdf")
@patch("unstructured.partition.pptx.partition_pptx")
@patch("unstructured.partition.docx.partition_docx")
@patch("airbyte_cdk.sources.file_based.file_types.unstructured_parser.detect_filetype")
def test_parse_records(
    mock_detect_filetype,
    mock_partition_docx,
    mock_partition_pptx,
    mock_partition_pdf,
    filetype,
    format_config,
    parse_result,
    raises,
    expected_records,
    parsing_error,
):
    stream_reader = MagicMock()
    mock_open(stream_reader.open_file, read_data=bytes(str(parse_result), "utf-8"))
    fake_file = RemoteFile(uri=FILE_URI, last_modified=datetime.now())
    fake_file.uri = FILE_URI
    logger = MagicMock()
    config = MagicMock()
    config.format = format_config
    mock_detect_filetype.return_value = filetype
    if parsing_error:
        mock_partition_docx.side_effect = Exception("weird parsing error")
        mock_partition_pptx.side_effect = Exception("weird parsing error")
        mock_partition_pdf.side_effect = Exception("weird parsing error")
    else:
        mock_partition_docx.return_value = parse_result
        mock_partition_pptx.return_value = parse_result
        mock_partition_pdf.return_value = parse_result
    if raises:
        with pytest.raises(RecordParseError):
            list(
                UnstructuredParser().parse_records(
                    config, fake_file, stream_reader, logger, MagicMock()
                )
            )
    else:
        assert (
            list(
                UnstructuredParser().parse_records(
                    config, fake_file, stream_reader, logger, MagicMock()
                )
            )
            == expected_records
        )


@pytest.mark.parametrize(
    "format_config, raises_for_status, json_response, is_ok, expected_error",
    [
        pytest.param(
            UnstructuredFormat(skip_unprocessable_file_types=False),
            False,
            {"status": "ok"},
            True,
            None,
            id="local",
        ),
        pytest.param(
            UnstructuredFormat(skip_unprocessable_file_types=False, strategy="fast"),
            False,
            {"status": "ok"},
            True,
            None,
            id="local_ok_strategy",
        ),
        pytest.param(
            UnstructuredFormat(skip_unprocessable_file_types=False, strategy="hi_res"),
            False,
            {"status": "ok"},
            False,
            "Hi-res strategy is not supported for local processing",
            id="local_unsupported_strategy",
        ),
        pytest.param(
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            False,
            [{"type": "Title", "text": "Airbyte source connection test"}],
            True,
            None,
            id="api_ok",
        ),
        pytest.param(
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            True,
            None,
            False,
            "API error",
            id="api_error",
        ),
        pytest.param(
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            False,
            {"unexpected": "response"},
            False,
            "Error",
            id="unexpected_handling_error",
        ),
    ],
)
@patch("airbyte_cdk.sources.file_based.file_types.unstructured_parser.requests")
def test_check_config(
    requests_mock, format_config, raises_for_status, json_response, is_ok, expected_error
):
    mock_response = MagicMock()
    mock_response.json.return_value = json_response
    if raises_for_status:
        mock_response.raise_for_status.side_effect = Exception("API error")
    requests_mock.post.return_value = mock_response
    result, error = UnstructuredParser().check_config(
        FileBasedStreamConfig(name="test", format=format_config)
    )
    assert result == is_ok
    if expected_error:
        assert expected_error in error


@pytest.mark.parametrize(
    "filetype, format_config, raises_for_status, file_content, json_response, expected_requests, raises, expected_records, http_status_code",
    [
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            None,
            "test",
            [{"type": "Text", "text": "test"}],
            [
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                )
            ],
            False,
            [{"content": "test", "document_key": FILE_URI, "_ab_source_file_parse_error": None}],
            200,
            id="basic_request",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                strategy="hi_res",
                processing=APIProcessingConfigModel(
                    mode="api",
                    api_key="test",
                    api_url="http://localhost:8000",
                    parameters=[
                        APIParameterConfigModel(name="include_page_breaks", value="true"),
                        APIParameterConfigModel(name="ocr_languages", value="eng"),
                        APIParameterConfigModel(name="ocr_languages", value="kor"),
                    ],
                ),
            ),
            None,
            "test",
            [{"type": "Text", "text": "test"}],
            [
                call(
                    "http://localhost:8000/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={
                        "strategy": "hi_res",
                        "include_page_breaks": "true",
                        "ocr_languages": ["eng", "kor"],
                    },
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                )
            ],
            False,
            [{"content": "test", "document_key": FILE_URI, "_ab_source_file_parse_error": None}],
            200,
            id="request_with_params",
        ),
        pytest.param(
            FileType.MD,
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            None,
            "# Mymarkdown",
            None,
            None,
            False,
            [
                {
                    "content": "# Mymarkdown",
                    "document_key": FILE_URI,
                    "_ab_source_file_parse_error": None,
                }
            ],
            200,
            id="handle_markdown_locally",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            [
                requests.exceptions.RequestException("API error"),
                requests.exceptions.RequestException("API error"),
                requests.exceptions.RequestException("API error"),
                requests.exceptions.RequestException("API error"),
                requests.exceptions.RequestException("API error"),
            ],
            "test",
            None,
            [
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
            ],
            True,
            None,
            200,
            id="retry_and_raise_on_api_error",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            [
                requests.exceptions.RequestException("API error"),
                requests.exceptions.RequestException("API error"),
                None,
            ],
            "test",
            [{"type": "Text", "text": "test"}],
            [
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
            ],
            False,
            [{"content": "test", "document_key": FILE_URI, "_ab_source_file_parse_error": None}],
            200,
            id="retry_and_recover",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            [
                Exception("Unexpected error"),
            ],
            "test",
            [{"type": "Text", "text": "test"}],
            [
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
            ],
            True,
            None,
            200,
            id="no_retry_on_unexpected_error",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            [
                requests.exceptions.RequestException(
                    "API error", response=MagicMock(status_code=400)
                ),
            ],
            "test",
            [{"type": "Text", "text": "test"}],
            [
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
                call().raise_for_status(),
            ],
            True,
            None,
            400,
            id="no_retry_on_400_error",
        ),
        pytest.param(
            FileType.PDF,
            UnstructuredFormat(
                skip_unprocessable_file_types=False,
                processing=APIProcessingConfigModel(mode="api", api_key="test"),
            ),
            None,
            "test",
            [{"detail": "Something went wrong"}],
            [
                call(
                    "https://api.unstructured.io/general/v0/general",
                    headers={"accept": "application/json", "unstructured-api-key": "test"},
                    data={"strategy": "auto"},
                    files={"files": ("filename", mock.ANY, "application/pdf")},
                ),
            ],
            False,
            [
                {
                    "content": None,
                    "document_key": FILE_URI,
                    "_ab_source_file_parse_error": "Error parsing record. This could be due to a mismatch between the config's file type and the actual file type, or because the file or record is not parseable. Contact Support if you need assistance.\nfilename=path/to/file.xyz message=[{'detail': 'Something went wrong'}]",
                }
            ],
            422,
            id="error_record_on_422_error",
        ),
    ],
)
@patch("airbyte_cdk.sources.file_based.file_types.unstructured_parser.requests")
@patch("airbyte_cdk.sources.file_based.file_types.unstructured_parser.detect_filetype")
@patch("time.sleep", side_effect=lambda _: None)
def test_parse_records_remotely(
    time_mock,
    mock_detect_filetype,
    requests_mock,
    filetype,
    format_config,
    raises_for_status,
    file_content,
    json_response,
    expected_requests,
    raises,
    expected_records,
    http_status_code,
):
    stream_reader = MagicMock()
    mock_open(stream_reader.open_file, read_data=bytes(str(file_content), "utf-8"))
    fake_file = RemoteFile(uri=FILE_URI, last_modified=datetime.now())
    fake_file.uri = FILE_URI
    logger = MagicMock()
    config = MagicMock()
    config.format = format_config
    mock_detect_filetype.return_value = filetype
    mock_response = MagicMock()
    mock_response.json.return_value = json_response
    mock_response.status_code = http_status_code
    if raises_for_status:
        mock_response.raise_for_status.side_effect = raises_for_status
    requests_mock.post.return_value = mock_response
    requests_mock.exceptions.RequestException = requests.exceptions.RequestException

    if raises:
        with pytest.raises(AirbyteTracedException) as exc:
            list(
                UnstructuredParser().parse_records(
                    config, fake_file, stream_reader, logger, MagicMock()
                )
            )
        # Failures from the API are treated as config errors
        assert exc.value.failure_type == FailureType.config_error
    else:
        assert (
            list(
                UnstructuredParser().parse_records(
                    config, fake_file, stream_reader, logger, MagicMock()
                )
            )
            == expected_records
        )

    if expected_requests:
        requests_mock.post.assert_has_calls(expected_requests)
    else:
        requests_mock.post.assert_not_called()
