#!/usr/bin/env python3
"""Repair ChatGPT citation tokens in Markdown exports."""

from __future__ import annotations

import argparse
import json
import pathlib
import re
import sys
from dataclasses import dataclass
from typing import Any


CITATION_RE = re.compile(r"cite(?:turn\d+(?:view|search)\d+)+")
REF_ID_RE = re.compile(r"turn\d+(?:view|search)\d+")
FENCE_RE = re.compile(r"^( {0,3})(`{3,}|~{3,})")
LINK_RE = re.compile(r"\[[^\]\n]+\]\([^) \n]+(?: [^)]+)?\)")
REFERENCE_HEADING_RE = re.compile(r"^## +参考资料\s*$", re.MULTILINE)
HTML_METADATA_RE = re.compile(
    r"<!--\s*deep-research-metadata\s*(?P<json>\{.*?\})\s*-->",
    re.DOTALL,
)


@dataclass(frozen=True)
class CitationRecord:
    ref_id: str
    url: str
    title: str = ""
    attribution: str = ""

    def footnote_text(self) -> str:
        label = self.title or self.url
        if self.attribution and self.title:
            return f"{self.attribution}：{self.title}。{self.url}"
        return f"{label}。{self.url}"


def extract_ref_ids(token: str) -> list[str]:
    return REF_ID_RE.findall(token)


def ref_id_from_ref(ref: dict[str, Any]) -> str | None:
    turn_index = ref.get("turn_index")
    ref_type = ref.get("ref_type")
    ref_index = ref.get("ref_index")
    if turn_index is None or ref_type not in {"view", "search"} or ref_index is None:
        return None
    return f"turn{turn_index}{ref_type}{ref_index}"


def _first_string(*values: Any) -> str:
    for value in values:
        if isinstance(value, str) and value.strip():
            return value.strip()
    return ""


def record_from_ref(ref: dict[str, Any]) -> CitationRecord | None:
    ref_id = _first_string(ref.get("ref_id"), ref.get("id")) or ref_id_from_ref(ref)
    url = _first_string(ref.get("url"), ref.get("safe_url"), ref.get("source_url"))
    title = _first_string(ref.get("title"), ref.get("name"), ref.get("text"))
    attribution = _first_string(ref.get("attribution"), ref.get("source"), ref.get("site_name"))
    if not ref_id or not url:
        return None
    return CitationRecord(ref_id=ref_id, url=url, title=title, attribution=attribution)


def _walk_dicts(value: Any) -> list[dict[str, Any]]:
    found: list[dict[str, Any]] = []
    if isinstance(value, dict):
        found.append(value)
        for child in value.values():
            found.extend(_walk_dicts(child))
    elif isinstance(value, list):
        for child in value:
            found.extend(_walk_dicts(child))
    return found


def build_citation_map(metadata: dict[str, Any] | None) -> dict[str, CitationRecord]:
    if not metadata:
        return {}

    citation_map: dict[str, CitationRecord] = {}
    for node in _walk_dicts(metadata):
        record = record_from_ref(node)
        if record and record.ref_id not in citation_map:
            citation_map[record.ref_id] = record

    return citation_map


def split_existing_references(markdown: str) -> tuple[str, str]:
    match = REFERENCE_HEADING_RE.search(markdown)
    if not match:
        return markdown.rstrip(), ""
    return markdown[: match.start()].rstrip(), markdown[match.start() :].strip()


def collect_markdown_links(markdown: str) -> list[str]:
    seen: set[str] = set()
    links: list[str] = []
    for link in LINK_RE.findall(markdown):
        if link not in seen:
            seen.add(link)
            links.append(link)
    return links


def replace_citations(
    markdown: str, citation_map: dict[str, CitationRecord]
) -> tuple[str, list[CitationRecord], list[str]]:
    footnote_numbers: dict[str, int] = {}
    ordered_records: list[CitationRecord] = []
    missing: list[str] = []

    def marker_for_ref(ref_id: str) -> str:
        if ref_id not in citation_map:
            if ref_id not in missing:
                missing.append(ref_id)
            return ""
        if ref_id not in footnote_numbers:
            footnote_numbers[ref_id] = len(ordered_records) + 1
            ordered_records.append(citation_map[ref_id])
        return f"[^{footnote_numbers[ref_id]}]"

    output: list[str] = []
    in_fence = False
    fence_marker = ""

    for line in markdown.splitlines(keepends=True):
        body = line[:-1] if line.endswith("\n") else line
        newline = "\n" if line.endswith("\n") else ""
        if body.endswith("\r"):
            body = body[:-1]
            newline = "\r" + newline

        fence_match = FENCE_RE.match(body)
        if fence_match:
            marker = fence_match.group(2)
            marker_char = marker[0]
            if not in_fence:
                in_fence = True
                fence_marker = marker_char * len(marker)
            elif marker_char == fence_marker[0] and len(marker) >= len(fence_marker):
                in_fence = False
                fence_marker = ""
            output.append(line)
            continue

        if in_fence:
            output.append(line)
            continue

        def replace_token(match: re.Match[str]) -> str:
            return "".join(marker_for_ref(ref_id) for ref_id in extract_ref_ids(match.group(0)))

        output.append(CITATION_RE.sub(replace_token, body) + newline)

    return "".join(output), ordered_records, missing


def strip_citation_tokens(markdown: str) -> str:
    lines: list[str] = []
    in_fence = False
    fence_marker = ""

    for line in markdown.splitlines(keepends=True):
        body = line[:-1] if line.endswith("\n") else line
        newline = "\n" if line.endswith("\n") else ""
        if body.endswith("\r"):
            body = body[:-1]
            newline = "\r" + newline

        fence_match = FENCE_RE.match(body)
        if fence_match:
            marker = fence_match.group(2)
            marker_char = marker[0]
            if not in_fence:
                in_fence = True
                fence_marker = marker_char * len(marker)
            elif marker_char == fence_marker[0] and len(marker) >= len(fence_marker):
                in_fence = False
                fence_marker = ""
            lines.append(line)
            continue

        if in_fence:
            lines.append(line)
        else:
            lines.append(CITATION_RE.sub("", body) + newline)

    return "".join(lines)


def render_reference_section(
    records: list[CitationRecord], existing_links: list[str], missing: list[str], degraded: bool
) -> str:
    lines = ["## 参考资料"]

    if degraded:
        lines.append("")
        lines.append("> 无法恢复逐句脚注映射：未找到可用元数据或部分 turn ID 缺少 URL。")

    if records:
        lines.append("")
        for index, record in enumerate(records, start=1):
            lines.append(f"[^{index}]: {record.footnote_text()}")

    if existing_links:
        lines.append("")
        lines.append("### 原始来源链接")
        for link in existing_links:
            lines.append(f"- {link}")

    if missing:
        lines.append("")
        lines.append("### 未解析引用 ID")
        for ref_id in missing:
            lines.append(f"- `{ref_id}`")

    return "\n".join(lines)


def repair_markdown(markdown: str, metadata: dict[str, Any] | None = None) -> str:
    body, existing_reference_section = split_existing_references(markdown)
    existing_links = collect_markdown_links(existing_reference_section)
    citation_map = build_citation_map(metadata)

    if citation_map:
        replaced, records, missing = replace_citations(body, citation_map)
        degraded = bool(missing)
        cleaned = strip_citation_tokens(replaced) if missing else replaced
        reference_section = render_reference_section(records, existing_links, missing, degraded)
        return f"{cleaned.rstrip()}\n\n{reference_section}\n"

    cleaned = strip_citation_tokens(body)
    reference_section = render_reference_section([], existing_links, [], degraded=True)
    return f"{cleaned.rstrip()}\n\n{reference_section}\n"


def load_json(path: pathlib.Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        value = json.load(handle)
    if not isinstance(value, dict):
        raise ValueError("metadata JSON must contain an object at the top level")
    return value


def extract_embedded_metadata(markdown: str) -> dict[str, Any] | None:
    match = HTML_METADATA_RE.search(markdown)
    if not match:
        return None
    value = json.loads(match.group("json"))
    if not isinstance(value, dict):
        raise ValueError("embedded metadata JSON must contain an object at the top level")
    return value


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Repair ChatGPT/Deep Research citation tokens in Markdown."
    )
    parser.add_argument("path", type=pathlib.Path, help="Markdown file to repair")
    parser.add_argument("--metadata", type=pathlib.Path, help="Deep Research metadata JSON")
    parser.add_argument("-o", "--output", type=pathlib.Path, help="Write repaired Markdown here")
    parser.add_argument("--in-place", action="store_true", help="Rewrite the input Markdown file")
    parser.add_argument("--dry-run", action="store_true", help="Print repaired Markdown to stdout")
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    if sum(bool(value) for value in (args.output, args.in_place, args.dry_run)) != 1:
        print("error: choose exactly one of --output, --in-place, or --dry-run", file=sys.stderr)
        return 2

    try:
        markdown = args.path.read_text(encoding="utf-8")
        metadata = load_json(args.metadata) if args.metadata else extract_embedded_metadata(markdown)
        fixed = repair_markdown(markdown, metadata=metadata)
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    if args.dry_run:
        sys.stdout.write(fixed)
        return 0

    destination = args.path if args.in_place else args.output
    assert destination is not None
    try:
        destination.write_text(fixed, encoding="utf-8")
    except OSError as exc:
        print(f"error: cannot write {destination}: {exc}", file=sys.stderr)
        return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))