from __future__ import annotations

import importlib.util
import pathlib
import sys
import textwrap
import unittest


SCRIPT_PATH = (
    pathlib.Path(__file__).resolve().parents[1]
    / "scripts"
    / "repair_markdown_citations.py"
)

spec = importlib.util.spec_from_file_location("repair_markdown_citations", SCRIPT_PATH)
repair = importlib.util.module_from_spec(spec)
assert spec is not None
assert spec.loader is not None
sys.modules["repair_markdown_citations"] = repair
spec.loader.exec_module(repair)


class RepairMarkdownCitationsTest(unittest.TestCase):
    def test_extracts_turn_ids_from_chatgpt_citation_token(self) -> None:
        token = "citeturn34view0turn20search2"

        self.assertEqual(
            repair.extract_ref_ids(token),
            ["turn34view0", "turn20search2"],
        )

    def test_replaces_tokens_with_first_use_footnotes_from_metadata(self) -> None:
        markdown = textwrap.dedent(
            """\
            武汉早期疫情通报强调未发现明显人传人。citeturn34view0

            郑州暴雨和唐山事件报道也体现了不同报道框架。citeturn20search2turn34view6
            """
        )
        metadata = {
            "content_references": [
                {
                    "items": [
                        {
                            "refs": [
                                {
                                    "turn_index": 34,
                                    "ref_type": "view",
                                    "ref_index": 0,
                                    "url": "https://wjw.hubei.gov.cn/example.shtml",
                                    "title": "湖北省卫健委通报",
                                    "attribution": "湖北省卫健委",
                                },
                                {
                                    "turn_index": 20,
                                    "ref_type": "search",
                                    "ref_index": 2,
                                    "url": "https://www.xinhuanet.com/zhengzhou.htm",
                                    "title": "新华社郑州暴雨报道",
                                },
                                {
                                    "turn_index": 34,
                                    "ref_type": "view",
                                    "ref_index": 6,
                                    "url": "https://www.xinhuanet.com/tangshan.htm",
                                    "title": "新华社唐山打人事件报道",
                                },
                            ]
                        }
                    ]
                }
            ]
        }

        fixed = repair.repair_markdown(markdown, metadata=metadata)

        self.assertIn("武汉早期疫情通报强调未发现明显人传人。[^1]", fixed)
        self.assertIn("报道框架。[^2][^3]", fixed)
        self.assertIn("## 参考资料", fixed)
        self.assertIn(
            "[^1]: 湖北省卫健委：湖北省卫健委通报。https://wjw.hubei.gov.cn/example.shtml",
            fixed,
        )
        self.assertIn(
            "[^2]: 新华社郑州暴雨报道。https://www.xinhuanet.com/zhengzhou.htm",
            fixed,
        )
        self.assertIn(
            "[^3]: 新华社唐山打人事件报道。https://www.xinhuanet.com/tangshan.htm",
            fixed,
        )

    def test_deduplicates_repeated_references(self) -> None:
        markdown = "Aciteturn1view0 Bciteturn1view0"
        metadata = {
            "content_references": [
                {
                    "items": [
                        {
                            "refs": [
                                {
                                    "turn_index": 1,
                                    "ref_type": "view",
                                    "ref_index": 0,
                                    "url": "https://example.com/a",
                                    "title": "Example A",
                                }
                            ]
                        }
                    ]
                }
            ]
        }

        fixed = repair.repair_markdown(markdown, metadata=metadata)

        self.assertIn("A[^1] B[^1]", fixed)
        self.assertEqual(fixed.count("[^1]:"), 1)

    def test_without_metadata_removes_tokens_and_adds_warning(self) -> None:
        markdown = "正文citeturn1view0\n\n[来源](https://example.com)"

        fixed = repair.repair_markdown(markdown, metadata=None)

        self.assertNotIn("cite", fixed)
        self.assertIn("## 参考资料", fixed)
        self.assertIn("[来源](https://example.com)", fixed)
        self.assertIn("无法恢复逐句脚注映射", fixed)

    def test_loads_embedded_metadata_from_html_comment(self) -> None:
        markdown = textwrap.dedent(
            """\
            正文citeturn1view0

            <!-- deep-research-metadata
            {"content_references":[{"items":[{"refs":[{"turn_index":1,"ref_type":"view","ref_index":0,"url":"https://example.com/a","title":"Example A"}]}]}]}
            -->
            """
        )

        metadata = repair.extract_embedded_metadata(markdown)
        fixed = repair.repair_markdown(markdown, metadata=metadata)

        self.assertIn("正文[^1]", fixed)
        self.assertIn("[^1]: Example A。https://example.com/a", fixed)

    def test_does_not_rewrite_tokens_inside_fenced_code_blocks(self) -> None:
        markdown = textwrap.dedent(
            """\
            正文citeturn1view0

            ```md
            示例citeturn1view0
            ```
            """
        )
        metadata = {
            "content_references": [
                {
                    "items": [
                        {
                            "refs": [
                                {
                                    "turn_index": 1,
                                    "ref_type": "view",
                                    "ref_index": 0,
                                    "url": "https://example.com/a",
                                    "title": "Example A",
                                }
                            ]
                        }
                    ]
                }
            ]
        }

        fixed = repair.repair_markdown(markdown, metadata=metadata)

        self.assertIn("正文[^1]", fixed)
        self.assertIn("示例citeturn1view0", fixed)


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