from __future__ import annotations

import argparse
import re
from pathlib import Path
from typing import Iterable

import pandas as pd
from spire.doc import Document

from compare_annotation import compare_with_log

# Map raw comment authors to unified review item names.
COMMENT_AUTHOR_MAPPING: dict[str, str] = {
    "三方货款审查": "第三方审查",
    "履行义务审查": "第三方审查",
    "债务转移审查": "第三方审查",
    "违约条款审查": "违约与延期审查",
    "延期审查": "违约与延期审查",
}


def clean_illegal(value: object) -> object:
    if isinstance(value, str):
        return re.compile(r"[\x00-\x08\x0B-\x0C\x0E-\x1F]").sub("", value)
    return value


def normalize_comment_author(author: str) -> str:
    author = author.strip()
    if not author:
        return author
    return COMMENT_AUTHOR_MAPPING.get(author, author)


def extract_annotaion(
    datasets_dir: Path,
    output_dir: Path,
    strip_suffixes: list[str] | None = None,
) -> None:
    """Extract review comments from Word files to xlsx files."""
    datasets_dir = datasets_dir.resolve()
    output_dir = output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    strip_suffixes = strip_suffixes or []

    for item in sorted(datasets_dir.iterdir()):
        if item.suffix.lower() == ".xlsx" or not item.is_file():
            continue

        document = Document()
        document.LoadFromFile(str(item))

        comments: list[dict[str, str]] = []
        for i in range(document.Comments.Count):
            comment = document.Comments[i]
            comment_text = ""
            for j in range(comment.Body.Paragraphs.Count):
                paragraph = comment.Body.Paragraphs[j]
                comment_text += paragraph.Text + "\n"

            comment_author = comment.Format.Author
            # 通过｜作为分隔符，只拿以后的进行比对
            author_split_idx = comment_author.find("|")
            comment_author = (
                comment_author[author_split_idx + 1 :]
                if author_split_idx != -1
                else comment_author
            )
            comment_author = normalize_comment_author(comment_author)
            comments.append(
                {
                    "审查项": clean_illegal(comment_author),
                    "合同原文": clean_illegal(comment.OwnerParagraph.Text),
                    "建议": clean_illegal(comment_text),
                }
            )

        df = pd.DataFrame(comments)
        clean_stem = _strip_suffix_once(item.stem, strip_suffixes)
        output_stem = clean_stem or item.stem
        output_file = output_dir / f"{output_stem}.xlsx"
        df.to_excel(output_file, index=False)
        document.Close()


def compare_annotaion(val_dir: Path, answer_dir: Path) -> None:
    """Run benchmark comparison on extracted annotations."""
    log_path = compare_with_log(val_dir=val_dir, answer_dir=answer_dir)
    print(f"Compare log written to: {log_path}")


def _strip_suffix_once(stem: str, suffixes: Iterable[str]) -> str:
    for suffix in suffixes:
        if suffix and stem.endswith(suffix):
            return stem[: -len(suffix)]
    return stem


def eval(
    datasets_dir: Path,
    answer_dir: Path,
    val_dir: Path,
    strip_suffixes: list[str] | None = None,
) -> None:
    """Pipeline: extract annotations first, then compare against ground truth."""
    strip_suffixes = strip_suffixes or []
    extract_annotaion(
        datasets_dir=datasets_dir,
        output_dir=val_dir,
        strip_suffixes=strip_suffixes,
    )
    compare_annotaion(val_dir=val_dir, answer_dir=answer_dir)


def _parse_args() -> argparse.Namespace:
    base = Path(__file__).parent
    parser = argparse.ArgumentParser(
        description="Extract review comments from docs and evaluate against answers."
    )
    parser.add_argument(
        "--datasets-dir",
        type=Path,
        default=base / "results" / "jp-output-lufa-20260511-101828",
        help="Directory containing Word files with annotations.",
    )
    parser.add_argument(
        "--strip-suffixes",
        nargs="*",
        default=["_麓发改进", "_人机交互", "_麓发迁移"],
        help=(
            "Optional filename suffixes to strip from generated val xlsx stems before "
            "comparison, e.g. --strip-suffixes _v1 _审阅版"
        ),
    )
    return parser.parse_args()


if __name__ == "__main__":
    base = Path(__file__).parent
    args = _parse_args()
    eval(
        datasets_dir=args.datasets_dir,
        answer_dir=base / "审查答案",
        val_dir=args.datasets_dir.with_name(
            f"{args.datasets_dir.name}-extract-comment"
        ),
        strip_suffixes=args.strip_suffixes,
    )
