"""
Textractor module tests
"""

import platform
import unittest

from txtai.pipeline import Textractor

# pylint: disable=C0411
from utils import Utils


class TestTextractor(unittest.TestCase):
    """
    Textractor tests.
    """

    def testClean(self):
        """
        Test text cleaning method
        """

        # Default text cleaning
        textractor = Textractor()
        self.assertEqual(textractor(" a  b  c "), "a b c")

        # Require text to be minlength
        textractor = Textractor(minlength=10)
        self.assertEqual(textractor(" a  b  c "), None)

        # Disable text cleaning
        textractor = Textractor(cleantext=False, minlength=10)
        self.assertEqual(textractor(" a  b  c "), " a  b  c ")

    def testChonkie(self):
        """
        Test a chonkie chunker
        """

        # Test chonkie chunking
        textractor = Textractor(chunker="sentence", chunk_size=5, chunk_overlap=0)
        self.assertEqual(textractor("This is a test. And another test."), ["This is a test.", "And another test."])

        # Test bad chunker throws an exception
        with self.assertRaises(AttributeError):
            textractor = Textractor(chunker="badchunker")

    def testDefault(self):
        """
        Test default text extraction
        """

        # Text input
        textractor = Textractor(backend=None)
        text = textractor(Utils.PATH + "/tabular.csv")
        self.assertEqual(len(text), 125)

        # Markdown input
        textractor = Textractor(sections=True)
        sections = textractor("# Heading 1\nText1\n\n# Heading 2\nText2\n")

        # Check number of sections is as expected
        self.assertEqual(len(sections), 2)

    @unittest.skipIf(platform.system() == "Darwin", "Docling skipped on macOS to avoid MPS issues")
    def testDocling(self):
        """
        Test docling backend
        """

        textractor = Textractor(backend="docling")

        # Extract text and check for Markdown formatting
        text = textractor(Utils.PATH + "/article.pdf")
        self.assertTrue("## Introducing txtai" in text)

    def testLines(self):
        """
        Test extraction to lines
        """

        textractor = Textractor(lines=True)

        # Extract text as lines
        lines = textractor(Utils.PATH + "/article.pdf")

        # Check number of lines is as expected
        self.assertEqual(len(lines), 35)

    def testLiteParse(self):
        """
        Test liteparse backend
        """

        textractor = Textractor(backend="liteparse")

        # Extract text and check for Markdown formatting
        text = textractor(Utils.PATH + "/article.pdf")
        self.assertTrue("# Introducing txtai" in text)

    def testHTML(self):
        """
        Test HTML to Markdown
        """

        # Headings
        self.assertMarkdown("<h1>This is a test</h1>", "# This is a test")
        self.assertMarkdown("<h6>This is a test</h6>", "###### This is a test")

        # Blockquotes
        self.assertMarkdown("<blockquote>This is a test</blockquote>", "> This is a test")

        # Lists
        self.assertMarkdown("<ul><li>Test1</li><li>Test2</li></ul>", "- Test1\n- Test2")
        self.assertMarkdown("<ol><li>Test1</li><li>Test2</li></ol>", "1. Test1\n2. Test2")

        # Code
        self.assertMarkdown("<code>This is a test</code>", "```\nThis is a test\n```")
        self.assertMarkdown("<pre>This is a test</pre>", "```\nThis is a test\n```")

        # Tables
        self.assertMarkdown(
            "<table><tr><th>Header1</th><th>Header2</th></tr><tr><td>Test1</td><td>Test2</td></tr></table>",
            "|Header1|Header2|\n|---|---|\n|Test1|Test2|",
        )

        # Ignore list
        self.assertMarkdown("<aside>This is a test</aside>", "")

        # Text formatting
        self.assertMarkdown("<p>This is a test</p>", "This is a test")
        self.assertMarkdown("<p>This is a <b>test</b</p>", "This is a **test**")
        self.assertMarkdown("<p>This is a <strong>test</strong></p>", "This is a **test**")
        self.assertMarkdown("<p>This is a <i>test</i></p>", "This is a *test*")
        self.assertMarkdown("<p>This is a <em>test</em></p>", "This is a *test*")
        self.assertMarkdown("<p>This is a <a href='link'>test</a>", "This is a [test](link)")

        # Collapse to outer tag
        self.assertMarkdown("<p>This is a <strong><em>test</em></strong></p>", "This is a **test**")
        self.assertMarkdown("<p>This is a <em><strong>test</strong></em></p>", "This is a *test*")

    def testParagraphs(self):
        """
        Test extraction to paragraphs
        """

        textractor = Textractor(paragraphs=True)

        # Extract text as paragraphs
        paragraphs = textractor(Utils.PATH + "/article.pdf")

        # Check number of paragraphs is as expected
        self.assertEqual(len(paragraphs), 11)

    def testSections(self):
        """
        Test extraction to sections
        """

        textractor = Textractor(sections=True)

        # Extract as sections
        sections = textractor(Utils.PATH + "/document.pdf")

        # Check number of sections is as expected
        self.assertEqual(len(sections), 3)

    def testSentences(self):
        """
        Test extraction to sentences
        """

        textractor = Textractor(sentences=True)

        # Extract text as sentences
        sentences = textractor(Utils.PATH + "/article.pdf")

        # Check number of sentences is as expected
        self.assertEqual(len(sentences), 17)

    def testSingle(self):
        """
        Test a single extraction with no tokenization of the results
        """

        textractor = Textractor()

        # Extract text as a single block
        text = textractor(Utils.PATH + "/article.pdf")

        # Check length of text is as expected
        self.assertEqual(len(text), 2471)

    def testTable(self):
        """
        Test table extraction
        """

        textractor = Textractor()

        # Extract text as a single block
        for name in ["document.docx", "spreadsheet.xlsx"]:
            text = textractor(f"{Utils.PATH}/{name}")

            # Check for table header
            self.assertTrue("|---|" in text)

    def testTikaFlag(self):
        """
        Test legacy tika flag
        """

        textractor = Textractor(tika=True)
        self.assertIsNotNone(textractor.html)

        textractor = Textractor(tika=False)
        self.assertIsNone(textractor.html)

    def testTuples(self):
        """
        Test output tuples
        """

        # Default text cleaning
        textractor = Textractor(tuples=True)

        path, text = textractor(Utils.PATH + "/article.pdf")
        self.assertEqual(path, Utils.PATH + "/article.pdf")
        self.assertEqual(len(text), 2471)

    def testURL(self):
        """
        Test parsing a remote URL
        """

        # Test parsing URLs for each backend
        for backend in ["docling", "liteparse", "tika"]:
            textractor = Textractor(backend=backend)
            text = textractor("https://github.com/neuml/txtai")
            self.assertTrue("txtai is an all-in-one AI framework" in text)

    def assertMarkdown(self, html, expected):
        """
        Helper method to assert generated markdown is as expected.

        Args:
            html: input html snippet
            expected: expected markdown text
        """

        textractor = Textractor()
        self.assertEqual(textractor(f"<html><body>{html}</body></html>"), expected)