#!/usr/bin/env python3
"""Convert Deep Research source IDs into numeric citations for final output."""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any


SRC_CITE_RE = re.compile(r"\[((?:src_[A-Za-z0-9_-]+)(?:\s*,\s*src_[A-Za-z0-9_-]+)*)\]")


def load_sources(path: Path) -> dict[str, dict[str, Any]]:
    sources: dict[str, dict[str, Any]] = {}
    if not path.exists():
        return sources
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        sid = obj.get("id") or obj.get("source_id")
        if sid:
            sources[str(sid)] = obj
    return sources


def extract_ordered_source_ids(text: str) -> list[str]:
    ordered: list[str] = []
    seen: set[str] = set()
    for match in SRC_CITE_RE.finditer(text):
        for sid in [item.strip() for item in match.group(1).split(",")]:
            if sid and sid not in seen:
                seen.add(sid)
                ordered.append(sid)
    return ordered


def _canonical_source_key(sid: str, source: dict[str, Any] | None) -> str:
    """Return a stable de-duplication key for a source record.

    Phase 2 often creates chapter-local source IDs for the same local PDF or
    official guideline. Final references should cite the underlying source
    once, while citation_map.json keeps the full src_id traceability.
    """
    if not source:
        return f"missing:{sid}"
    title = re.sub(r"\s+", " ", str(source.get("title") or source.get("name") or sid)).strip().lower()
    title = title.removesuffix(" ocr").removesuffix(".ocr").strip()
    doi = str(source.get("doi") or "").strip().lower()
    if doi:
        return f"doi:{doi}"
    path = str(source.get("path") or "").strip()
    url = str(source.get("url") or "").strip()
    if title and ("phase0/extracted/" in path or "phase0/extracted/" in url):
        return f"local-material:{title}"
    for field in ("url", "path"):
        value = str(source.get(field) or "").strip()
        if value:
            return f"{field}:{value.rstrip('/').lower()}"
    return f"title:{title or sid}"


def build_numeric_mapping(
    ordered_ids: list[str],
    sources: dict[str, dict[str, Any]],
) -> tuple[dict[str, int], list[dict[str, Any]]]:
    mapping: dict[str, int] = {}
    records: list[dict[str, Any]] = []
    seen_keys: dict[str, int] = {}
    record_by_number: dict[int, dict[str, Any]] = {}
    for sid in ordered_ids:
        source = sources.get(sid)
        key = _canonical_source_key(sid, source)
        if key in seen_keys:
            number = seen_keys[key]
            mapping[sid] = number
            record_by_number[number].setdefault("source_ids", []).append(sid)
            continue
        number = len(records) + 1
        seen_keys[key] = number
        mapping[sid] = number
        record = {
            "number": number,
            "source_id": sid,
            "source_ids": [sid],
            "source": source or {},
            "dedupe_key": key,
        }
        records.append(record)
        record_by_number[number] = record
    return mapping, records


def format_reference(number: int, sid: str, source: dict[str, Any] | None) -> str:
    if not source:
        return f"{number}. {sid}. （sources.jsonl 未找到该来源）"
    authors = ", ".join(source.get("authors", [])) if source.get("authors") else ""
    year = source.get("year") or source.get("date") or ""
    title = source.get("title") or source.get("name") or sid
    title = re.sub(r"(?i)(?:\s+OCR|\.ocr)$", "", str(title)).strip()
    publisher = source.get("publisher") or source.get("venue") or source.get("source") or ""
    url = source.get("url") or source.get("path") or ""
    parts = [f"{number}. "]
    if authors:
        parts.append(f"{authors}. ")
    if year:
        parts.append(f"({year}). ")
    parts.append(str(title))
    if publisher:
        parts.append(f". {publisher}")
    if url:
        parts.append(f". {url}")
    return "".join(parts)


def convert_citations(text: str, mapping: dict[str, int]) -> str:
    def repl(match: re.Match[str]) -> str:
        ids = [item.strip() for item in match.group(1).split(",") if item.strip()]
        nums: list[str] = []
        seen: set[int] = set()
        for sid in ids:
            if sid not in mapping:
                continue
            number = mapping[sid]
            if number in seen:
                continue
            seen.add(number)
            nums.append(str(number))
        return "<sup>[" + ", ".join(nums) + "]</sup>" if nums else match.group(0)

    return SRC_CITE_RE.sub(repl, text)


def strip_existing_reference_section(text: str) -> str:
    pattern = re.compile(r"\n##\s*(?:参考文献|参考来源清单|References)\s*\n.*\Z", re.S)
    return pattern.sub("", text).rstrip() + "\n"


def number_citations(
    *,
    text: str,
    sources: dict[str, dict[str, Any]],
) -> tuple[str, list[dict[str, Any]]]:
    ordered_ids = extract_ordered_source_ids(text)
    mapping, records = build_numeric_mapping(ordered_ids, sources)
    body = convert_citations(strip_existing_reference_section(text), mapping).rstrip()
    ref_lines = ["", "## 参考来源清单", ""]
    for record in records:
        ref_lines.append(format_reference(record["number"], record["source_id"], record["source"]))
    return body + "\n" + "\n".join(ref_lines).rstrip() + "\n", records


def main() -> int:
    parser = argparse.ArgumentParser(description="Convert [src_xxx] citations to numeric citations")
    parser.add_argument("project", help="Project directory")
    parser.add_argument("--input", default="phase4/final_zh.md")
    parser.add_argument("--output", default="phase4/final_zh_numbered.md")
    parser.add_argument("--sources", default="phase2/sources.jsonl")
    parser.add_argument("--map", default="phase4/citation_map.json")
    args = parser.parse_args()

    project = Path(args.project)
    src_path = project / args.input
    out_path = project / args.output
    sources_path = project / args.sources
    map_path = project / args.map
    if not src_path.exists():
        raise SystemExit(f"input not found: {src_path}")
    sources = load_sources(sources_path)
    numbered, records = number_citations(text=src_path.read_text(encoding="utf-8"), sources=sources)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(numbered, encoding="utf-8")
    map_path.parent.mkdir(parents=True, exist_ok=True)
    map_path.write_text(json.dumps(records, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote: {out_path.relative_to(project)}")
    print(f"Wrote: {map_path.relative_to(project)}")
    print(f"Citations: {len(records)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
