#!/usr/bin/env python3
"""
MCRO multi-group cyclic shifts summary
======================================

Scan a directory of MCRO JSON files and output a CSV listing all documents
(filename) that belong to MORE THAN ONE pdf_group (via tracking.groups),
with one row per CYCLIC SHIFT of the groups.

Output columns:
    filename, pdf_group1, pdf_group2, pdf_group3, font_group1, font_group2, font_group3

Rules:
  - Only include rows where the document has >= 2 distinct pdf_group values.
  - For each document:
      * Collect (pdf_group, font_group) from tracking.groups[*].
      * Deduplicate by pdf_group.
      * Sort by pdf_group alphabetically.
      * Keep up to 3 groups (if more exist, take the first 3 alphabetically).
      * Let m = number of groups (2 or 3).
      * Produce m rows, each a cyclic shift of the group list:
            shift 0: [g1, g2, g3]
            shift 1: [g2, g3, g1]
            shift 2: [g3, g1, g2]
        (apply the same rotation to font_group values).
      * Pad up to 3 columns with empty strings if needed.

So a doc with groups A1, C1 yields:
    filename, A1, C1, "", font_for_A1, font_for_C1, ""

And a doc with groups A1, B1, C1 yields:
    filename, A1, B1, C1, fA, fB, fC
    filename, B1, C1, A1, fB, fC, fA
    filename, C1, A1, B1, fC, fA, fB

Usage:
  python3 mcro_multigroup_shifts.py \
      --input /path/to/json_dir \
      --out multigroup_shifts.csv \
      [--pattern "*.json"] \
      [--recurse]
"""

import argparse
import csv
import json
from pathlib import Path
from typing import Any, Dict, List, Tuple
import sys


def as_list(x):
    if x is None:
        return []
    if isinstance(x, list):
        return x
    return [x]


def main():
    ap = argparse.ArgumentParser(description="Summarize documents with multiple pdf_groups using cyclic shifts.")
    ap.add_argument("--input", required=True, help="Directory containing JSON files")
    ap.add_argument("--out", required=True, help="Output CSV file path")
    ap.add_argument("--pattern", default="*.json", help='Glob for JSON files (default: "*.json")')
    ap.add_argument("--recurse", action="store_true", help="Recurse into subdirectories")
    args = ap.parse_args()

    json_dir = Path(args.input)
    if not json_dir.is_dir():
        print(f"[error] Input directory does not exist or is not a directory: {json_dir}", file=sys.stderr)
        sys.exit(1)

    if args.recurse:
        files = sorted(json_dir.rglob(args.pattern))
    else:
        files = sorted(json_dir.glob(args.pattern))

    if not files:
        print(f"[warn] No JSON files matched {args.pattern} under {json_dir}", file=sys.stderr)
        sys.exit(0)

    rows: List[Dict[str, Any]] = []

    for jp in files:
        try:
            data = json.loads(jp.read_text(encoding="utf-8"))
        except Exception as e:
            print(f"[warn] Skipping {jp}: failed to parse JSON ({e})", file=sys.stderr)
            continue

        if not isinstance(data, dict):
            print(f"[warn] Skipping {jp}: JSON root is not an object", file=sys.stderr)
            continue

        filename = (data.get("filename") or "").strip()
        if not filename:
            continue

        tracking = data.get("tracking") or {}
        groups = as_list(tracking.get("groups"))
        if not groups:
            continue

        # Collect unique pdf_group -> font_group mapping
        group_map: Dict[str, str] = {}
        for g in groups:
            if not isinstance(g, dict):
                continue
            pg = g.get("pdf_group")
            if pg is None:
                continue
            pg_str = str(pg).strip()
            if not pg_str:
                continue
            fg = g.get("font_group")
            fg_str = str(fg).strip() if fg is not None else ""
            if pg_str not in group_map:
                group_map[pg_str] = fg_str

        # Only care about docs that belong to 2 or more distinct pdf_groups
        if len(group_map) < 2:
            continue

        # Sort by pdf_group alphabetically
        sorted_groups: List[Tuple[str, str]] = sorted(group_map.items(), key=lambda kv: kv[0])

        # Keep at most 3 groups (since we only have 3 columns)
        sorted_groups = sorted_groups[:3]
        pdf_groups_base = [g[0] for g in sorted_groups]
        font_groups_base = [g[1] for g in sorted_groups]
        m = len(pdf_groups_base)

        # Produce m cyclic shifts
        for shift in range(m):
            shifted_pdf = pdf_groups_base[shift:] + pdf_groups_base[:shift]
            shifted_font = font_groups_base[shift:] + font_groups_base[:shift]

            # Pad to 3 with empty strings
            while len(shifted_pdf) < 3:
                shifted_pdf.append("")
            while len(shifted_font) < 3:
                shifted_font.append("")

            row = {
                "filename": filename,
                "pdf_group1": shifted_pdf[0],
                "pdf_group2": shifted_pdf[1],
                "pdf_group3": shifted_pdf[2],
                "font_group1": shifted_font[0],
                "font_group2": shifted_font[1],
                "font_group3": shifted_font[2],
            }
            rows.append(row)

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    fieldnames = [
        "filename",
        "pdf_group1", "pdf_group2", "pdf_group3",
        "font_group1", "font_group2", "font_group3",
    ]

    with out_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f"[ok] Wrote {len(rows)} multi-group shift rows to: {out_path.resolve()}")


if __name__ == "__main__":
    main()