import csv
import json
import os
import pathlib
import platform
import tempfile
import zlib
from unittest.mock import mock_open, patch

import pandas as pd
import pytest

from test_unstructured.unit_utils import assign_hash_ids, input_path
from unstructured.documents.elements import (
    Address,
    CheckBox,
    CoordinatesMetadata,
    CoordinateSystem,
    DataSourceMetadata,
    Element,
    ElementMetadata,
    ElementType,
    FigureCaption,
    Form,
    Formula,
    Image,
    Link,
    ListItem,
    NarrativeText,
    PageBreak,
    Table,
    TableChunk,
    Text,
    Title,
)
from unstructured.errors import DecompressedSizeExceededError
from unstructured.partition.email import partition_email
from unstructured.partition.text import partition_text
from unstructured.staging import base


def test_base64_gzipped_json_to_elements_can_deserialize_compressed_elements_from_a_JSON_string():
    base64_elements_str = (
        "eJyFzcsKwjAQheFXKVm7yDS3xjcQXNaViKTJjBR6o46glr67zVI3Lmf4Dv95EdhhjwNf2yT2hYDGUaWtJVm5WDoq"
        "NUL0UoJrqtLHJHaF6JFDChw2v6zbzfjkvD2OM/YZ8GvC/Khb7lBs5LcilUwRyCsblQYTiBQpZRxYZcCA/1spDtP9"
        "8dU6DTEw3sa5fWOqs10vH0cLQn0="
    )

    elements = base.elements_from_base64_gzipped_json(base64_elements_str)

    assert elements == [Title("Lorem"), Text("Lorem Ipsum")]


def test_elements_from_base64_gzipped_json_raises_error_if_decompression_is_incomplete():
    base64_elements_str = (
        "eJyFzcsKwjAQheFXKVm7yDS3xjcQXNaViKTJjBR6o46glr67zVI3Lmf4Dv95EdhhjwNf2yT2hYDGUaWtJVm5WDoq"
        "NUL0UoJrqtLHJHaF6JFDChw2v6zbzfjkvD2OM/YZ8GvC/Khb7lBs5LcilUwRyCsblQYTiBQpZRxYZcCA/1spDtP9"
        "8dU6DTEw3sa5fWOqs10vH0cL="
    )

    with pytest.raises(zlib.error):
        base.elements_from_base64_gzipped_json(base64_elements_str)


def test_elements_from_base64_gzipped_json_raises_error_if_decompression_exceeds_max_size():
    base64_elements_str = (
        "eJyFzcsKwjAQheFXKVm7yDS3xjcQXNaViKTJjBR6o46glr67zVI3Lmf4Dv95EdhhjwNf2yT2hYDGUaWtJVm5WDoq"
        "NUL0UoJrqtLHJHaF6JFDChw2v6zbzfjkvD2OM/YZ8GvC/Khb7lBs5LcilUwRyCsblQYTiBQpZRxYZcCA/1spDtP9"
        "8dU6DTEw3sa5fWOqs10vH0cLQn0="
    )

    with (
        patch("unstructured.staging.base.MAX_DECOMPRESSED_SIZE", 32),
        pytest.raises(DecompressedSizeExceededError),
    ):
        base.elements_from_base64_gzipped_json(base64_elements_str)


def test_elements_to_base64_gzipped_json_can_serialize_elements_to_a_base64_str():
    elements = assign_hash_ids([Title("Lorem"), Text("Lorem Ipsum")])

    b64_str = base.elements_to_base64_gzipped_json(elements)

    # Verify round-trip: decompress and check the JSON content rather than the
    # exact compressed bytes, which vary across zlib implementations (e.g. zlib-ng).
    round_tripped = base.elements_from_base64_gzipped_json(b64_str)
    assert len(round_tripped) == 2
    assert round_tripped[0].text == "Lorem"
    assert round_tripped[1].text == "Lorem Ipsum"


def test_elements_to_dicts():
    elements = [Title(text="Title 1"), NarrativeText(text="Narrative 1")]
    isd = base.elements_to_dicts(elements)

    assert isd[0]["text"] == "Title 1"
    assert isd[0]["type"] == ElementType.TITLE

    assert isd[1]["text"] == "Narrative 1"
    assert isd[1]["type"] == "NarrativeText"


def test_elements_from_dicts():
    element_dicts = [
        {"text": "Blurb1", "type": "NarrativeText"},
        {"text": "Blurb2", "type": "Title"},
        {"text": "Blurb3", "type": "ListItem"},
        {"text": "Blurb4", "type": "BulletedText"},
        {"text": "No Type"},
    ]

    elements = base.elements_from_dicts(element_dicts)
    assert elements == [
        NarrativeText(text="Blurb1"),
        Title(text="Blurb2"),
        ListItem(text="Blurb3"),
        ListItem(text="Blurb4"),
    ]


def test_elements_from_dicts_form():
    element_dicts = [
        {"text": "Applicant Name: Jane Doe", "type": "Form"},
    ]

    elements = base.elements_from_dicts(element_dicts)

    assert elements == [Form(text="Applicant Name: Jane Doe")]


def test_convert_to_csv(tmp_path: str):
    output_csv_path = os.path.join(tmp_path, "isd_data.csv")
    elements = [Title(text="Title 1"), NarrativeText(text="Narrative 1")]
    with open(output_csv_path, "w+") as csv_file:
        isd_csv_string = base.convert_to_csv(elements)
        csv_file.write(isd_csv_string)

    with open(output_csv_path) as csv_file:
        csv_rows = csv.DictReader(csv_file)
        assert all(set(row.keys()) == set(base.TABLE_FIELDNAMES) for row in csv_rows)


def test_convert_to_dataframe():
    elements = [Title(text="Title 1"), NarrativeText(text="Narrative 1")]
    df = base.convert_to_dataframe(elements)
    expected_df = pd.DataFrame(
        {
            "type": ["Title", "NarrativeText"],
            "text": ["Title 1", "Narrative 1"],
        },
    )
    assert df.type.equals(expected_df.type) is True  # type: ignore
    assert df.text.equals(expected_df.text) is True  # type: ignore


def test_convert_to_dataframe_maintains_fields():
    elements = partition_email(
        "example-docs/eml/fake-email-attachment.eml",
        process_attachements=True,
    )
    df = base.convert_to_dataframe(elements)
    for element in elements:
        metadata = element.metadata.to_dict()
        for key in metadata:
            assert key in df.columns


def test_default_pandas_dtypes():
    """Ensure all element fields have a dtype in dict returned by get_default_pandas_dtypes()."""
    full_element = Text(
        text="some text",
        element_id="123",
        coordinates=((1, 2), (3, 4)),
        coordinate_system=CoordinateSystem(width=12.3, height=99.4),
        detection_origin="some origin",
        embeddings=[1.1, 2.2, 3.3, 4.4],
        metadata=ElementMetadata(
            coordinates=CoordinatesMetadata(
                points=((1, 2), (3, 4)),
                system=CoordinateSystem(width=12.3, height=99.4),
            ),
            data_source=DataSourceMetadata(
                url="http://mysite.com",
                version="123",
                record_locator={"some": "data", "value": 3},
                date_created="then",
                date_processed="now",
                date_modified="before",
                permissions_data=[{"data": 1}, {"data": 2}],
            ),
            filename="filename",
            file_directory="file_directory",
            last_modified="last_modified",
            filetype="filetype",
            attached_to_filename="attached_to_filename",
            parent_id="parent_id",
            category_depth=1,
            image_path="image_path",
            languages=["eng", "spa"],
            page_number=1,
            page_name="page_name",
            url="url",
            link_urls=["links", "url"],
            link_texts=["links", "texts"],
            links=[Link(text="text", url="url", start_index=1)],
            sent_from=["sent", "from"],
            sent_to=["sent", "to"],
            subject="subject",
            header_footer_type="header_footer_type",
            emphasized_text_contents=["emphasized", "text", "contents"],
            emphasized_text_tags=["emphasized", "text", "tags"],
            text_as_html="text_as_html",
            is_continuation=True,
            detection_class_prob=0.5,
        ),
    )
    element_as_dict = full_element.to_dict()
    element_as_dict.update(
        base.flatten_dict(
            element_as_dict.pop("metadata"), keys_to_omit=["data_source_record_locator"]
        ),
    )
    flattened_element_keys = element_as_dict.keys()
    default_dtypes = base.get_default_pandas_dtypes()
    dtype_keys = default_dtypes.keys()
    for key in flattened_element_keys:
        assert key in dtype_keys


@pytest.mark.skipif(
    platform.system() == "Windows",
    reason="Posix Paths are not available on Windows",
)
def test_elements_to_dicts_serializes_with_posix_paths():
    metadata = ElementMetadata(filename=pathlib.PosixPath("../../fake-file.txt"))
    elements = [
        Title(text="Title 1", metadata=metadata),
        NarrativeText(text="Narrative 1", metadata=metadata),
    ]
    output = base.elements_to_dicts(elements)
    # NOTE(robinson) - json.dumps should run without raising an exception
    json.dumps(output)


def test_all_elements_preserved_when_serialized():
    metadata = ElementMetadata(filename="fake-file.txt")
    elements = [
        Address(text="address", metadata=metadata, element_id="1"),
        CheckBox(checked=True, metadata=metadata, element_id="2"),
        FigureCaption(text="caption", metadata=metadata, element_id="3"),
        Title(text="title", metadata=metadata, element_id="4"),
        NarrativeText(text="narrative", metadata=metadata, element_id="5"),
        ListItem(text="list", metadata=metadata, element_id="6"),
        Image(text="image", metadata=metadata, element_id="7"),
        Text(text="text", metadata=metadata, element_id="8"),
        TableChunk(text="table chunk", metadata=metadata, element_id="9"),
        PageBreak(text=""),
    ]

    element_dicts = base.elements_to_dicts(elements)
    assert base.elements_to_dicts(base.elements_from_dicts(element_dicts)) == element_dicts


def test_serialized_deserialize_elements_to_json(tmpdir: str):
    filename = os.path.join(tmpdir, "fake-elements.json")
    metadata = ElementMetadata(filename="fake-file.txt")
    elements = [
        Address(text="address", metadata=metadata, element_id="1"),
        CheckBox(checked=True, metadata=metadata, element_id="2"),
        FigureCaption(text="caption", metadata=metadata, element_id="3"),
        Title(text="title", metadata=metadata, element_id="4"),
        NarrativeText(text="narrative", metadata=metadata, element_id="5"),
        ListItem(text="list", metadata=metadata, element_id="6"),
        Image(text="image", metadata=metadata, element_id="7"),
        Text(text="text", metadata=metadata, element_id="8"),
        TableChunk(text="table chunk", metadata=metadata, element_id="9"),
        PageBreak(text=""),
    ]

    base.elements_to_json(elements, filename=filename)
    new_elements_filename = base.elements_from_json(filename=filename)
    assert elements == new_elements_filename

    elements_str = base.elements_to_json(elements)
    assert elements_str is not None
    new_elements_text = base.elements_from_json(text=elements_str)
    assert elements == new_elements_text


def test_read_and_write_json_with_encoding():
    elements = partition_text("example-docs/fake-text-utf-16-be.txt")

    with tempfile.TemporaryDirectory() as tmp_dir_path:
        tmp_file_path = os.path.join(tmp_dir_path, "tmp_file")
        base.elements_to_json(elements, filename=tmp_file_path, encoding="utf-16")
        new_elements_filename = base.elements_from_json(filename=tmp_file_path, encoding="utf-16")

    assert elements == new_elements_filename


def test_filter_element_types_with_include_element_type():
    element_types = [Title]
    elements = partition_text("example-docs/fake-text.txt")
    elements = base.filter_element_types(elements=elements, include_element_types=element_types)
    for element in elements:
        assert type(element) in element_types


def test_filter_element_types_with_exclude_element_type():
    element_types = [Title]
    elements = partition_text("example-docs/fake-text.txt")
    elements = base.filter_element_types(elements=elements, exclude_element_types=element_types)
    for element in elements:
        assert type(element) not in element_types


def test_filter_element_types_with_exclude_and_include_element_type():
    element_types = [Title]
    elements = partition_text("example-docs/fake-text.txt")
    with pytest.raises(ValueError):
        elements = base.filter_element_types(
            elements=elements,
            exclude_element_types=element_types,
            include_element_types=element_types,
        )


def test_convert_to_coco():
    elements = [
        Text(
            text="some text",
            element_id="123",
            detection_origin="some origin",
            embeddings=[1.1, 2.2, 3.3, 4.4],
            metadata=ElementMetadata(
                coordinates=CoordinatesMetadata(
                    points=((1, 2), (1, 4), (3, 4), (3, 2)),
                    system=CoordinateSystem(width=12.3, height=99.4),
                ),
                data_source=DataSourceMetadata(
                    url="http://mysite.com",
                    version="123",
                    record_locator={"some": "data", "value": 3},
                    date_created="then",
                    date_processed="now",
                    date_modified="before",
                    permissions_data=[{"data": 1}, {"data": 2}],
                ),
                filename="filename",
                file_directory="file_directory",
                last_modified="last_modified",
                filetype="filetype",
                attached_to_filename="attached_to_filename",
                parent_id="parent_id",
                category_depth=1,
                image_path="image_path",
                languages=["eng", "spa"],
                page_number=1,
                page_name="page_name",
                url="url",
                link_urls=["links", "url"],
                link_texts=["links", "texts"],
                links=[Link(text="text", url="url", start_index=1)],
                sent_from=["sent", "from"],
                sent_to=["sent", "to"],
                subject="subject",
                header_footer_type="header_footer_type",
                emphasized_text_contents=["emphasized", "text", "contents"],
                emphasized_text_tags=["emphasized", "text", "tags"],
                text_as_html="text_as_html",
                is_continuation=True,
                detection_class_prob=0.5,
            ),
        )
    ]
    missing_elements = [
        Text(
            text="some text",
            element_id="123",
            detection_origin="some origin",
            embeddings=[1.1, 2.2, 3.3, 4.4],
            metadata=ElementMetadata(
                data_source=DataSourceMetadata(
                    url="http://mysite.com",
                    version="123",
                    record_locator={"some": "data", "value": 3},
                    date_created="then",
                    date_processed="now",
                    date_modified="before",
                    permissions_data=[{"data": 1}, {"data": 2}],
                ),
                filename="filename",
                file_directory="file_directory",
                last_modified="last_modified",
                filetype="filetype",
                attached_to_filename="attached_to_filename",
                parent_id="parent_id",
                category_depth=1,
                image_path="image_path",
                languages=["eng", "spa"],
                page_number=1,
                page_name="page_name",
                url="url",
                link_urls=["links", "url"],
                link_texts=["links", "texts"],
                links=[Link(text="text", url="url", start_index=1)],
                sent_from=["sent", "from"],
                sent_to=["sent", "to"],
                subject="subject",
                header_footer_type="header_footer_type",
                emphasized_text_contents=["emphasized", "text", "contents"],
                emphasized_text_tags=["emphasized", "text", "tags"],
                text_as_html="text_as_html",
                is_continuation=True,
                detection_class_prob=0.5,
            ),
        )
    ]
    full_coco = base.convert_to_coco(elements)
    limited_coco = base.convert_to_coco(missing_elements)
    assert full_coco["annotations"][0]["area"]
    assert limited_coco["annotations"][0]["area"] is None


def test_flatten_dict():
    """Flattening a simple dictionary"""
    dictionary = {"a": 1, "b": 2, "c": 3}
    expected_result = {"a": 1, "b": 2, "c": 3}
    assert base.flatten_dict(dictionary) == expected_result


def test_flatten_nested_dict():
    """Flattening a nested dictionary"""
    dictionary = {"a": 1, "b": {"c": 2, "d": 3}, "e": 4}
    expected_result = {"a": 1, "b_c": 2, "b_d": 3, "e": 4}
    assert base.flatten_dict(dictionary) == expected_result


def test_flatten_dict_with_tuples():
    """Flattening a dictionary with tuples"""
    dictionary = {"a": 1, "b": (2, 3, 4), "c": {"d": 5, "e": (6, 7)}}
    expected_result = {"a": 1, "b": (2, 3, 4), "c_d": 5, "c_e": (6, 7)}
    assert base.flatten_dict(dictionary) == expected_result


def test_flatten_dict_with_lists():
    """Flattening a dictionary with lists"""
    dictionary = {"a": 1, "b": [2, 3, 4], "c": {"d": 5, "e": [6, 7]}}
    expected_result = {"a": 1, "b": [2, 3, 4], "c_d": 5, "c_e": [6, 7]}
    assert base.flatten_dict(dictionary) == expected_result


def test_flatten_dict_with_omit_keys():
    """Flattening a dictionary with keys to omit"""
    dictionary = {"a": 1, "b": {"c": 2, "d": 3}, "e": 3}
    keys_to_omit = ["b"]
    expected_result = {"a": 1, "b": {"c": 2, "d": 3}, "e": 3}
    assert base.flatten_dict(dictionary, keys_to_omit=keys_to_omit) == expected_result


def test_flatten_dict_alt_separator():
    """Flattening a dictionary with separator other than "_" """
    dictionary = {"a": 1, "b": {"c": 2, "d": 3}, "e": 4}
    separator = "-"
    expected_result = {"a": 1, "b-c": 2, "b-d": 3, "e": 4}
    assert base.flatten_dict(dictionary, separator=separator) == expected_result


def test_flatten_dict_flatten_tuple():
    """Flattening a dictionary with flatten_lists set to True, to flatten tuples"""
    dictionary = {"a": 1, "b": (2, 3, 4), "c": {"d": 5, "e": (6, 7)}}
    expected_result = {"a": 1, "b_0": 2, "b_1": 3, "b_2": 4, "c_d": 5, "c_e_0": 6, "c_e_1": 7}
    assert base.flatten_dict(dictionary, flatten_lists=True) == expected_result


def test_flatten_dict_flatten_list():
    """Flattening a dictionary with flatten_lists set to True"""
    dictionary = {"a": 1, "b": [2, 3, 4], "c": {"d": 5, "e": [6, 7]}}
    expected_result = {"a": 1, "b_0": 2, "b_1": 3, "b_2": 4, "c_d": 5, "c_e_0": 6, "c_e_1": 7}
    assert base.flatten_dict(dictionary, flatten_lists=True) == expected_result


def test_flatten_dict_flatten_list_omit_keys():
    """Flattening a dictionary with flatten_lists set to True and also omitting keys"""
    dictionary = {"a": 1, "b": [2, 3, 4], "c": {"d": 5, "e": [6, 7]}}
    keys_to_omit = ["c"]
    expected_result = {"a": 1, "b_0": 2, "b_1": 3, "b_2": 4, "c": {"d": 5, "e": [6, 7]}}
    assert (
        base.flatten_dict(dictionary, keys_to_omit=keys_to_omit, flatten_lists=True)
        == expected_result
    )


def test_flatten_dict_flatten_list_omit_keys_remove_none():
    """Flattening a dictionary with flatten_lists set to True and also omitting keys
    and setting remove_none to True"""
    dictionary = {"a": None, "b": [2, 3, 4], "c": {"d": None, "e": [6, 7]}}
    keys_to_omit = ["c"]
    expected_result = {"b_0": 2, "b_1": 3, "b_2": 4, "c": {"d": None, "e": [6, 7]}}
    assert (
        base.flatten_dict(
            dictionary, keys_to_omit=keys_to_omit, flatten_lists=True, remove_none=True
        )
        == expected_result
    )


def test_flatten_dict_flatten_list_remove_none():
    """Flattening a dictionary with flatten_lists set to True and setting remove_none to True"""
    dictionary = {"a": None, "b": [2, 3, 4], "c": {"d": None, "e": [6, 7]}}
    expected_result = {"b_0": 2, "b_1": 3, "b_2": 4, "c_e_0": 6, "c_e_1": 7}
    assert base.flatten_dict(dictionary, flatten_lists=True, remove_none=True) == expected_result


def test_flatten_dict_flatten_list_none_in_list_remove_none():
    """Flattening a dictionary with flatten_lists and remove_none set to True and None in list"""
    dictionary = {"a": 1, "b": [2, 3, 4], "c": {"d": None, "e": [6, None]}}
    expected_result = {"a": 1, "b_0": 2, "b_1": 3, "b_2": 4, "c_e_0": 6}
    assert base.flatten_dict(dictionary, flatten_lists=True, remove_none=True) == expected_result


def test_flatten_dict_flatten_list_omit_keys2():
    """Flattening a dictionary with flatten_lists set to True and also omitting keys"""
    dictionary = {"a": 1, "b": [2, 3, 4], "c": {"d": 5, "e": [6, 7]}}
    keys_to_omit = ["b"]
    expected_result = {"a": 1, "b": [2, 3, 4], "c_d": 5, "c_e_0": 6, "c_e_1": 7}
    assert (
        base.flatten_dict(dictionary, keys_to_omit=keys_to_omit, flatten_lists=True)
        == expected_result
    )


def test_flatten_dict_flatten_list_omit_keys3():
    """Flattening a dictionary with flatten_lists set to True and also omitting nested keys"""
    dictionary = {"a": 1, "b": [2, 3, 4], "c": {"d": 5, "e": [6, 7]}}
    keys_to_omit = ["c_e"]
    expected_result = {"a": 1, "b_0": 2, "b_1": 3, "b_2": 4, "c_d": 5, "c_e": [6, 7]}
    assert (
        base.flatten_dict(dictionary, keys_to_omit=keys_to_omit, flatten_lists=True)
        == expected_result
    )


def test_flatten_dict_flatten_list_omit_keys4():
    """Flattening a dictionary with flatten_lists set to True and also omitting nested keys"""
    dictionary = {"a": 1, "b": [2, 3, 4], "c": {"d": 5, "e": {"f": 6, "g": 7}}}
    keys_to_omit = ["c_e"]
    expected_result = {"a": 1, "b_0": 2, "b_1": 3, "b_2": 4, "c_d": 5, "c_e": {"f": 6, "g": 7}}
    assert (
        base.flatten_dict(dictionary, keys_to_omit=keys_to_omit, flatten_lists=True)
        == expected_result
    )


def test_flatten_empty_dict():
    """Flattening an empty dictionary"""
    assert base.flatten_dict({}) == {}


def test_flatten_dict_empty_lists():
    """Flattening a dictionary with empty lists"""
    assert base.flatten_dict({"a": [], "b": {"c": []}}) == {"a": [], "b_c": []}


@pytest.mark.parametrize(
    ("json_filename", "expected_md_filename"),
    [
        (
            "test_unstructured/testfiles/staging/UDHR_first_article_all.txt.json",
            "test_unstructured/testfiles/staging/UDHR_first_article_all.txt.md",
        ),
        (
            "test_unstructured/testfiles/staging/embedded-images.pdf.json",
            "test_unstructured/testfiles/staging/embedded-images.pdf.md",
        ),
    ],
)
def test_elements_to_md_conversion(json_filename: str, expected_md_filename: str):
    """Test that elements_from_json followed by elements_to_md produces expected markdown output."""
    # Rehydrate elements from JSON
    elements = base.elements_from_json(json_filename)

    # Convert to markdown
    markdown_output = base.elements_to_md(elements)

    # Read expected markdown fixture
    with open(expected_md_filename) as f:
        expected_markdown = f.read()

    # Compare outputs
    assert markdown_output == expected_markdown


@pytest.mark.parametrize(
    ("element", "expected_markdown", "exclude_binary"),
    [
        (Title("Test Title"), "# Test Title", False),
        (NarrativeText("This is some narrative text."), "This is some narrative text.", False),
        (Formula(r"\int_a^b x^2 dx"), "$$\n\\int_a^b x^2 dx\n$$", False),
        (
            Image(
                "Test Image",
                metadata=ElementMetadata(
                    image_base64=(
                        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhf"
                        "DwAChwGA60e6kgAAAABJRU5ErkJggg=="
                    ),
                    image_mime_type="image/png",
                ),
            ),
            (
                "![Test Image](data:image/png;base64,"
                "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgA"
                "AAABJRU5ErkJggg==)"
            ),
            False,
        ),
        (
            Image(
                "Test Image",
                metadata=ElementMetadata(
                    image_base64=(
                        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60"
                        "e6kgAAAABJRU5ErkJggg=="
                    ),
                    image_mime_type="image/png",
                ),
            ),
            "Test Image",
            True,
        ),
        (
            Image(
                "Test Image", metadata=ElementMetadata(image_url="https://example.com/image.jpg")
            ),
            "![Test Image](https://example.com/image.jpg)",
            False,
        ),
        (
            Table(
                "Table Text",
                metadata=ElementMetadata(text_as_html="<table><tr><td>Test</td></tr></table>"),
            ),
            "<table><tr><td>Test</td></tr></table>",
            False,
        ),
        (Table("Table Text"), "Table Text", False),
    ],
)
def test_element_to_md_conversion(element: "Element", expected_markdown: str, exclude_binary: bool):
    """Test individual element to markdown conversion for different element types."""
    assert (
        base.element_to_md(element, exclude_binary_image_data=exclude_binary) == expected_markdown
    )


def test_element_to_md_formula_normalizes_common_math_symbols():
    element = Formula("x ∈ A and y ≤ z and a × b = c")
    assert base.element_to_md(element) == (
        "$$\nx \\in{} A and y \\leq{} z and a \\times{} b = c\n$$"
    )


def test_element_to_md_formula_can_disable_normalization():
    element = Formula("x ∈ A and y ≤ z and a × b = c")
    assert (
        base.element_to_md(element, normalize_formula=False)
        == "$$\nx ∈ A and y ≤ z and a × b = c\n$$"
    )


def test_element_to_md_formula_preserves_unicode_square_root():
    """√ must not become \\sqrt{} (would break √2, √(x+1), etc.)."""
    assert base.element_to_md(Formula("√2")) == "$$\n√2\n$$"
    assert base.element_to_md(Formula("√(x+1)")) == "$$\n√(x+1)\n$$"
    assert base.element_to_md(Formula("√2 ≤ x")) == "$$\n√2 \\leq{} x\n$$"


def test_elements_to_md_positional_encoding_backward_compat():
    """Legacy positional (... filename, exclude_binary_image_data, encoding) must still work."""
    m = mock_open()
    with patch("builtins.open", m):
        base.elements_to_md([Title("x")], "out.md", False, "latin-1")
    assert m.call_count == 1
    _args, kwargs = m.call_args
    assert kwargs.get("encoding") == "latin-1"


def test_create_file_from_elements_positional_no_group_by_page_backward_compat():
    """Legacy HTML positional args through no_group_by_page still bind correctly."""
    with patch("unstructured.partition.html.convert.elements_to_html") as mock_html:
        mock_html.return_value = "<html></html>"
        base.create_file_from_elements([Title("P")], "html", None, "utf-8", False, False)
    mock_html.assert_called_once()
    assert mock_html.call_args.kwargs["exclude_binary_image_data"] is False
    assert mock_html.call_args.kwargs["no_group_by_page"] is False


def test_elements_to_md_formula_can_disable_normalization():
    elements = [Formula("x ∈ A")]
    assert base.elements_to_md(elements, normalize_formula=False) == "$$\nx ∈ A\n$$"


def test_create_file_from_elements_markdown_passes_formula_normalization_flag():
    elements = [Formula("x ∈ A")]
    content = base.create_file_from_elements(
        elements,
        output_format="markdown",
        normalize_formula=False,
    )
    assert content == "$$\nx ∈ A\n$$"


def test_element_to_md_formula_auto_plain_for_noisy_ocr():
    text = "_ CRo-CR O= OR"
    assert base.element_to_md(Formula(text)) == text


def test_element_to_md_formula_auto_plain_when_embedded_dollar_delimiters():
    assert base.element_to_md(Formula("a $$ b")) == "a $$ b"
    assert base.element_to_md(Formula("inline $x$ math")) == "inline $x$ math"


def test_element_to_md_formula_display_math_fallback_when_unsafe_delimiters():
    raw = "a $$ b"
    assert (
        base.element_to_md(
            Formula(raw),
            formula_markdown_style=base.FORMULA_MARKDOWN_DISPLAY_MATH,
        )
        == raw
    )


def test_element_to_md_formula_display_math_wraps_when_auto_would_plain():
    assert base.element_to_md(Formula("x = 1")) == "x = 1"
    assert (
        base.element_to_md(
            Formula("x = 1"),
            formula_markdown_style=base.FORMULA_MARKDOWN_DISPLAY_MATH,
        )
        == "$$\nx = 1\n$$"
    )


def test_element_to_md_formula_auto_plain_for_prose_style_caption():
    text = (
        "The corrosion rate (CR) was calculated using Eq. (1) "
        "and we reference [1–5] for detail in this manuscript."
    )
    assert base.element_to_md(Formula(text)) == text


def test_element_to_md_formula_invalid_style_raises():
    with pytest.raises(ValueError, match="formula_markdown_style"):
        base.element_to_md(Formula("x=1"), formula_markdown_style="nope")


def test_elements_to_md_formula_markdown_style_keyword_only():
    els = [Formula("x ∈ A")]
    out = base.elements_to_md(els, formula_markdown_style=base.FORMULA_MARKDOWN_PLAIN)
    assert out == "x ∈ A"
    out_plain_unicode = base.elements_to_md(
        els,
        normalize_formula=False,
        formula_markdown_style=base.FORMULA_MARKDOWN_PLAIN,
    )
    assert out_plain_unicode == "x ∈ A"


def test_element_to_md_formula_plain_never_normalizes_unicode_minus():
    assert (
        base.element_to_md(
            Formula("a − b"),
            formula_markdown_style=base.FORMULA_MARKDOWN_PLAIN,
        )
        == "a − b"
    )


def test_element_to_md_formula_in_brace_boundary_after_symbol():
    out = base.element_to_md(
        Formula("x∈S"),
        formula_markdown_style=base.FORMULA_MARKDOWN_DISPLAY_MATH,
    )
    assert out == "$$\nx\\in{}S\n$$"


def test_formula_auto_scores_raw_text_prose_with_one_symbol_stays_plain():
    text = (
        "E ≤ threshold where E is the energy and threshold was determined experimentally "
        "in the laboratory setup described above herein."
    )
    assert base.element_to_md(Formula(text)) == text


def test_elements_from_json_to_md_with_formula_fixture():
    """JSON fixture → elements_from_json → elements_to_md (real Formula types)."""
    path = input_path("staging/formula-elements.json")
    elements = base.elements_from_json(filename=path)
    assert len(elements) == 2
    assert elements[0].category == "NarrativeText"
    assert elements[1].category == "Formula"
    md = base.elements_to_md(elements)
    assert md == "See equation below.\n$$\nE = mc^2\n$$"


def test_elements_to_md_five_positional_args_order():
    """Lock (elements, filename, exclude_binary_image_data, encoding, normalize_formula)."""
    m = mock_open()
    elements = [Formula("x ∈ A")]
    with patch("builtins.open", m):
        out = base.elements_to_md(elements, "out.md", False, "latin-1", False)
    assert out == "$$\nx ∈ A\n$$"
    assert m.call_args.kwargs.get("encoding") == "latin-1"


def test_create_file_from_elements_propagates_formula_markdown_style():
    elements = [Formula("x = 1")]
    assert base.create_file_from_elements(elements, output_format="markdown") == "x = 1"
    dm = base.create_file_from_elements(
        elements,
        output_format="markdown",
        formula_markdown_style=base.FORMULA_MARKDOWN_DISPLAY_MATH,
    )
    assert dm == "$$\nx = 1\n$$"
    plain = base.create_file_from_elements(
        elements,
        output_format="markdown",
        formula_markdown_style=base.FORMULA_MARKDOWN_PLAIN,
    )
    assert plain == "x = 1"


def test_elements_to_md_file_output():
    """Test elements_to_md function with file output."""

    elements = [Title("Test Title"), NarrativeText("Test content.")]

    # Test file output
    with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as tmp_file:
        tmp_filename = tmp_file.name

    try:
        markdown_content = base.elements_to_md(elements, filename=tmp_filename)

        # Check that the function returns the content
        assert markdown_content == "# Test Title\nTest content."

        # Check that the file was created with correct content
        with open(tmp_filename) as f:
            file_content = f.read()
        assert file_content == "# Test Title\nTest content."

    finally:
        # Clean up
        if os.path.exists(tmp_filename):
            os.unlink(tmp_filename)


def test_create_file_from_elements_markdown():
    """Test create_file_from_elements with format=markdown returns and optionally writes file."""
    elements = [Title("Heading"), NarrativeText("Some body text.")]
    content = base.create_file_from_elements(elements, output_format="markdown")
    assert content == "# Heading\nSome body text."

    with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as tmp_file:
        tmp_filename = tmp_file.name
    try:
        out = base.create_file_from_elements(
            elements, output_format="markdown", filename=tmp_filename
        )
        assert out == content
        with open(tmp_filename) as f:
            assert f.read() == content
    finally:
        if os.path.exists(tmp_filename):
            os.unlink(tmp_filename)


def test_create_file_from_elements_text():
    """Test create_file_from_elements with format=text."""
    elements = [Title("A"), NarrativeText("B")]
    content = base.create_file_from_elements(elements, output_format="text")
    assert content == "A\nB"


def test_create_file_from_elements_html():
    """Test create_file_from_elements with format=html returns HTML."""
    elements = [Title("Page"), NarrativeText("Content")]
    content = base.create_file_from_elements(elements, output_format="html")
    assert "<!DOCTYPE html" in content
    assert "<body>" in content
    assert "Page" in content
    assert "Content" in content


def test_create_file_from_elements_unsupported_format():
    """Test create_file_from_elements raises for unsupported format."""
    elements = [Title("X")]
    with pytest.raises(ValueError, match="Unsupported format"):
        base.create_file_from_elements(elements, output_format="pdf")


def test_create_file_from_elements_html_group_by_page_drops_elements_without_page_number():
    """With no_group_by_page=False, elements without page_number are skipped (body empty)."""
    elements = [Title("Page"), NarrativeText("Content")]
    content = base.create_file_from_elements(elements, output_format="html", no_group_by_page=False)
    assert "<!DOCTYPE html" in content
    assert "<body>" in content
    # Elements without metadata.page_number are not included when grouping by page
    assert "Page" not in content
    assert "Content" not in content


@pytest.mark.parametrize(
    ("format_name", "expected_in_content"),
    [
        ("markdown", "# Heading\nSome body text."),
        ("text", "Heading\nSome body text."),
        ("html", "<!DOCTYPE html"),
    ],
)
def test_create_file_from_elements_filename_write(format_name: str, expected_in_content: str):
    """Test create_file_from_elements writes correct content to file for all formats."""
    elements = [Title("Heading"), NarrativeText("Some body text.")]
    ext = {"markdown": ".md", "text": ".txt", "html": ".html"}[format_name]
    with tempfile.NamedTemporaryFile(mode="w", suffix=ext, delete=False) as tmp_file:
        tmp_filename = tmp_file.name
    try:
        out = base.create_file_from_elements(
            elements, output_format=format_name, filename=tmp_filename
        )
        assert expected_in_content in out
        with open(tmp_filename) as f:
            written = f.read()
        assert expected_in_content in written
        assert out == written
    finally:
        if os.path.exists(tmp_filename):
            os.unlink(tmp_filename)


def test_create_file_from_elements_exclude_binary_image_data_markdown():
    """exclude_binary_image_data=True passthrough: markdown omits base64 image data."""
    elements = [
        Title("Doc"),
        Image(
            "Alt",
            metadata=ElementMetadata(
                image_base64="abc123",
                image_mime_type="image/png",
            ),
        ),
    ]
    content = base.create_file_from_elements(
        elements, output_format="markdown", exclude_binary_image_data=True
    )
    assert "base64," not in content
    assert "Alt" in content


def test_create_file_from_elements_exclude_binary_image_data_html():
    """exclude_binary_image_data=True passthrough: HTML omits base64 image data."""
    elements = [
        Title("Doc"),
        Image(
            "Alt",
            metadata=ElementMetadata(
                image_base64="abc123",
                image_mime_type="image/png",
            ),
        ),
    ]
    content = base.create_file_from_elements(
        elements, output_format="html", exclude_binary_image_data=True
    )
    assert "abc123" not in content


@pytest.mark.parametrize(
    ("format_arg", "expected_substring"),
    [
        (" Markdown ", "# Heading"),
        ("HTML ", "<!DOCTYPE html"),
        (" TEXT ", "Heading"),
    ],
)
def test_create_file_from_elements_format_normalization(format_arg: str, expected_substring: str):
    """Format string is stripped and lowercased (e.g. ' Markdown ' -> 'markdown')."""
    elements = [Title("Heading"), NarrativeText("Body")]
    content = base.create_file_from_elements(elements, output_format=format_arg)
    assert expected_substring in content


def test_element_to_md_with_none_mime_type():
    """Test element_to_md handles None mime_type gracefully."""
    from unstructured.documents.elements import ElementMetadata, Image

    # Test Image element with None mime_type
    image_metadata = ElementMetadata(
        image_base64=(
            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU"
            "5ErkJggg=="
        ),
        image_mime_type=None,
    )
    image_element = Image("Test Image", metadata=image_metadata)

    # Should handle None mime_type gracefully
    result = base.element_to_md(image_element)
    assert "![Test Image](data:image/*" in result
    assert "base64," in result
