from __future__ import annotations

import argparse
import json
import re
from dataclasses import dataclass
from pathlib import Path
from zipfile import ZipFile
from xml.etree import ElementTree as ET


W_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
W = f"{{{W_NS}}}"

TEXT_TAGS = {
    f"{W}t",
    f"{W}delText",
    f"{W}instrText",
}

CONTENT_TAGS = TEXT_TAGS | {
    f"{W}drawing",
    f"{W}object",
    f"{W}pict",
    f"{W}oMath",
    f"{W}oMathPara",
    f"{W}noBreakHyphen",
    f"{W}softHyphen",
    f"{W}tab",
    f"{W}br",
    f"{W}cr",
}

IMAGE_RE = re.compile(r"!\[[^\]]*\]\([^)]+\)")
IMAGE_GLUE_RE = re.compile(
    r"!\[[^\]]*\]\([^)]+\)(?=[^\s<>)\].,;:!?])|(?<=[^\s<(\[.,;:!?])!\[[^\]]*\]\([^)]+\)"
)
FOUR_PLUS_STARS_RE = re.compile(r"(?<!^)\*{4,}(?!$)", re.M)
ODD_BACKTICK_LINE_RE = re.compile(r"`")


@dataclass
class SourceStats:
    total_paragraphs: int = 0
    nonempty_paragraphs: int = 0


def local_name(tag: str) -> str:
    return tag.rsplit("}", 1)[-1] if "}" in tag else tag


def paragraph_has_content(p: ET.Element) -> bool:
    for node in p.iter():
        if node.tag in TEXT_TAGS and (node.text or "").strip():
            return True
        if node.tag in CONTENT_TAGS - TEXT_TAGS:
            return True
    return False


def source_stats(docx_path: Path) -> SourceStats:
    stats = SourceStats()
    with ZipFile(docx_path) as zf:
        root = ET.fromstring(zf.read("word/document.xml"))
    for p in root.iter(f"{W}p"):
        stats.total_paragraphs += 1
        if paragraph_has_content(p):
            stats.nonempty_paragraphs += 1
    return stats


def load_json(path: Path) -> dict | list:
    return json.loads(path.read_text(encoding="utf-8"))


def count_odd_backtick_lines(text: str) -> int:
    count = 0
    for line in text.splitlines():
        if line.count("`") % 2:
            count += 1
    return count


def scan_markdown(md_path: Path) -> dict[str, int]:
    text = md_path.read_text(encoding="utf-8", errors="replace")
    four_plus = 0
    for match in FOUR_PLUS_STARS_RE.finditer(text):
        if match.group(0).strip("*"):
            four_plus += 1
    return {
        "image_count": len(IMAGE_RE.findall(text)),
        "image_glue_count": len(IMAGE_GLUE_RE.findall(text)),
        "four_plus_stars_count": four_plus,
        "odd_backtick_line_count": count_odd_backtick_lines(text),
    }


def audit_entry(entry: dict) -> dict:
    input_path = Path(entry["input"])
    output_path = Path(entry["output"])
    export_report_path = output_path.parent / "export-report.json"

    problems: list[str] = []
    source = source_stats(input_path)

    if not output_path.exists():
        problems.append("missing_markdown_output")
        return {
            "input": str(input_path),
            "output": str(output_path),
            "problems": problems,
        }

    if not export_report_path.exists():
        problems.append("missing_export_report")
        return {
            "input": str(input_path),
            "output": str(output_path),
            "problems": problems,
        }

    export_report = load_json(export_report_path)
    report_stats = export_report.get("stats", {})
    batch_stats = entry.get("stats", {})
    markdown = scan_markdown(output_path)

    rendered_block_count = (
        int(report_stats.get("paragraphs", 0))
        + int(report_stats.get("headings", 0))
        + int(report_stats.get("lists", 0))
        + int(report_stats.get("code_blocks", 0))
    )

    if report_stats != batch_stats:
        problems.append("batch_report_mismatch")
    if rendered_block_count != source.nonempty_paragraphs:
        problems.append("paragraph_count_mismatch")
    if markdown["image_count"] != int(report_stats.get("images", 0)):
        problems.append("image_count_mismatch")
    if markdown["image_glue_count"]:
        problems.append("image_glue")
    if markdown["four_plus_stars_count"]:
        problems.append("four_plus_stars")
    if markdown["odd_backtick_line_count"]:
        problems.append("odd_backtick_lines")

    return {
        "input": str(input_path),
        "output": str(output_path),
        "source_total_paragraphs": source.total_paragraphs,
        "source_nonempty_paragraphs": source.nonempty_paragraphs,
        "report_rendered_blocks": rendered_block_count,
        "report_images": int(report_stats.get("images", 0)),
        "markdown_images": markdown["image_count"],
        "image_glue_count": markdown["image_glue_count"],
        "four_plus_stars_count": markdown["four_plus_stars_count"],
        "odd_backtick_line_count": markdown["odd_backtick_line_count"],
        "warnings": list(report_stats.get("warnings", [])),
        "problems": problems,
    }


def summarize(results: list[dict]) -> dict:
    problem_counts: dict[str, int] = {}
    for result in results:
        for problem in result.get("problems", []):
            problem_counts[problem] = problem_counts.get(problem, 0) + 1
    return {
        "files_audited": len(results),
        "files_with_problems": sum(1 for result in results if result.get("problems")),
        "problem_counts": problem_counts,
        "problem_examples": [result for result in results if result.get("problems")][:25],
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Audit build-corpus batch output against source .docx files.")
    parser.add_argument("--batch-report", required=True, help="Path to build-corpus-batch-report.json")
    parser.add_argument("--out", default="", help="Optional path for the audit JSON report")
    args = parser.parse_args()

    batch_report_path = Path(args.batch_report).resolve()
    entries = load_json(batch_report_path)
    if not isinstance(entries, list):
        raise SystemExit("Batch report must be a JSON array.")

    results = [audit_entry(entry) for entry in entries]
    summary = summarize(results)
    payload = {
        "batch_report": str(batch_report_path),
        "summary": summary,
        "results": results,
    }

    out_path = Path(args.out).resolve() if args.out else batch_report_path.with_name("corpus-audit-report.json")
    out_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2))
    print(f"WROTE {out_path}")
    return 0


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