import tempfile
import unittest
from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile

from pypdf import PdfWriter
from pypdf.generic import DecodedStreamObject, DictionaryObject, NameObject

from rag.document_parser import parse_document


class DocumentParserTests(unittest.TestCase):
    def test_parses_txt_and_md_as_utf8_text(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            txt = root / "notice.txt"
            md = root / "notice.md"
            txt.write_text("基金风险提示", encoding="utf-8")
            md.write_text("# 产品说明\n\n开放式基金", encoding="utf-8")

            self.assertEqual(parse_document(txt), "基金风险提示")
            self.assertEqual(parse_document(md), "# 产品说明\n\n开放式基金")

    def test_extracts_text_from_docx(self):
        with tempfile.TemporaryDirectory() as tmp:
            docx = Path(tmp) / "notice.docx"
            document_xml = (
                '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
                '<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
                '<w:body><w:p><w:r><w:t>基金产品说明</w:t></w:r></w:p>'
                '<w:p><w:r><w:t>风险揭示</w:t></w:r></w:p></w:body></w:document>'
            )
            with ZipFile(docx, "w", ZIP_DEFLATED) as archive:
                archive.writestr("word/document.xml", document_xml)

            self.assertEqual(parse_document(docx), "基金产品说明\n风险揭示")

    def test_extracts_text_from_pdf(self):
        with tempfile.TemporaryDirectory() as tmp:
            pdf = Path(tmp) / "notice.pdf"
            writer = PdfWriter()
            page = writer.add_blank_page(width=612, height=792)
            font = writer._add_object(
                DictionaryObject(
                    {
                        NameObject("/Type"): NameObject("/Font"),
                        NameObject("/Subtype"): NameObject("/Type1"),
                        NameObject("/BaseFont"): NameObject("/Helvetica"),
                    }
                )
            )
            page[NameObject("/Resources")] = DictionaryObject(
                {NameObject("/Font"): DictionaryObject({NameObject("/F1"): font})}
            )
            page[NameObject("/Contents")] = DecodedStreamObject()
            page[NameObject("/Contents")].set_data(b"BT /F1 12 Tf 72 720 Td (Fund FAQ) Tj ET")
            with pdf.open("wb") as output:
                writer.write(output)

            self.assertEqual(parse_document(pdf), "Fund FAQ")

    def test_rejects_unsupported_extension(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "notice.xlsx"
            path.write_bytes(b"not supported")

            with self.assertRaises(ValueError):
                parse_document(path)


if __name__ == "__main__":
    unittest.main()