#!/usr/bin/env python3
"""
MCRO PDF Group Report
=====================

Scans a directory of JSON document descriptors, filters to documents with
  tracking.tracked_flag == true, explodes tracking.groups (multi-group
membership), and emits a set of CSV "data dumps" plus per-group statistics.

FOCUS:
  • Only JSON files where tracking.tracked_flag is true (bool True or "true").
  • Grouping is by EACH entry in tracking.groups[*].pdf_group.
    (A single document can belong to multiple pdf_groups.)
  • Everything is driven by the JSON structure used in the MCRO pipeline
    (docs, objects, language.terms, signatures.pdfsig, metadata, tracking, case).

EXTRACTED FIELDS
================

From the top-level document and nested structures, for each *tracked* document
and for each group in tracking.groups:

  filename
  filing_type
  filing_date
  pdf_page_count

  tracking.groups[*]            -> per-group fields (pdf_group, font_group,
                                   font_hash_date, font1_sha256, etc.)
  objects[*]                    -> flattened (one row per (group, object))
  language.terms[*]             -> search_group, search_term, quantity
                                   (one row per (group, term))
  signatures.pdfsig             -> all keys (one row per (group, doc))
  metadata.*                    -> all keys as columns
  case.*                        -> all keys EXCEPT case.cluster_size

ATTORNEYS:
  • case.attorney_defense and case.attorney_prosecution are emitted into a
    dedicated table: pdf_group_attorneys.csv (one row per (group, attorney)).
  • pdf_group_documents.csv does NOT contain raw attorney JSON blobs.

LANGUAGE:
  • language.instances is NOT included.

FLAGS:
  • flags.* is NOT included (handled separately via CSV).

SIGNATURES:
  • signatures.pdfsig is included (all keys).
  • signatures.full_report and signatures.mcro_authentication are NOT included
    in the data-dump tables.

OUTPUT FILES (all CSV, written into --outdir)
=============================================

  1) pdf_group_documents.csv
     One row per (tracked document, pdf_group), with a curated column set
     in this exact order:
       pdf_group, font_group, case_case_id, case_cluster_id, case_cluster_name,
       filename, filing_type, filing_date, pdf_page_count,
       tracking_font_hash_date,
       tracking_font1_sha256, tracking_font2_sha256,
       tracking_font3_sha256, tracking_font4_sha256,
       tracking_font1_name, tracking_font2_name,
       tracking_font3_name, tracking_font4_name,
       tracking_font1_type, tracking_font2_type,
       tracking_font3_type, tracking_font4_type,
       metadata_efile_date,
       metadata_xmp_create_date, metadata_xmp_metadata_date,
       metadata_document_id, metadata_instance_id,
       metadata_xmp_author, metadata_xmp_company,
       metadata_xmp_creator, metadata_xmp_title,
       metadata_xmp_creator_tool, metadata_xmp_toolkit,
       metadata_derived_create_newer_by_days,
       metadata_derived_metadata_newer_by_days,
       case_defendant_name, case_case_status,
       case_case_assigned_judge, case_filing_date,
       case_case_year, case_in_past_flag, case_link_past_flag

  2) pdf_group_attorneys.csv
     One row per (pdf_group, attorney entry) per document:
       • pdf_group, doc_sha256, filename, case_id, cluster_id, cluster_name
       • attorney_role ("defense" or "prosecution")
       • attorney_index (position in the list)
       • attorney_<field> for every key in the original attorney dict.

  3) pdf_group_objects.csv
     One row per (pdf_group, doc_sha256, object_sha256).
     Includes all object-level keys plus doc-level context.

  4) pdf_group_terms.csv
     One row per (pdf_group, doc_sha256, search_group, search_term).
     Includes quantity plus doc and case context.

  5) pdf_group_signatures_pdfsig.csv
     One row per (pdf_group, tracked document).
     All signatures.pdfsig keys flattened as columns, plus pdf_group/doc context.

  6) pdf_group_overview.csv
     One row per pdf_group, with high-level statistics:
       - pdf_group
       - n_docs
       - n_cases
       - n_clusters
       - case_ids
       - cluster_ids
       - cluster_names
       - min_filing_date, max_filing_date
       - min_pdf_page_count, max_pdf_page_count, avg_pdf_page_count
       - n_unique_filing_types, filing_types
       - n_unique_terms, n_term_occurrences
       - n_unique_search_groups
       - n_objects, n_unique_object_sha256, n_unique_object_font_names
       - n_docs_with_pdfsig

  7) pdf_group_term_stats.csv
     One row per (pdf_group, search_group, search_term):
       - total_quantity
       - doc_count
       - case_count
       - cluster_count
       - first_seen_date
       - last_seen_date

  8) README.txt
     Brief description of what's in the folder.

USAGE
=====

  python mcro_pdf_group_report.py \
      --input /path/to/json_dir \
      --outdir /path/to/pdf_group_report \
      [--pattern "*.json"] \
      [--recurse]

The script is robust to missing keys and will simply leave those cells blank.
"""

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List
import sys
import pandas as pd


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


def join_unique(vals):
    vals = [v for v in vals if pd.notna(v)]
    return "; ".join(sorted(set(map(str, vals))))


def safe_get(d: Dict[str, Any], key: str, default=None):
    if isinstance(d, dict):
        return d.get(key, default)
    return default


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True, help="Directory containing JSON files")
    ap.add_argument("--outdir", required=True, help="Output directory for CSVs")
    ap.add_argument("--pattern", default="*.json", help="Glob pattern for JSON files (default: *.json)")
    ap.add_argument("--recurse", action="store_true", help="Recurse into subdirectories")
    args = ap.parse_args()

    in_dir = Path(args.input)
    if args.recurse:
        files = sorted(in_dir.rglob(args.pattern))
    else:
        files = sorted(in_dir.glob(args.pattern))

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

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    # Data collectors
    docs_rows: List[Dict[str, Any]] = []
    objects_rows: List[Dict[str, Any]] = []
    terms_rows: List[Dict[str, Any]] = []
    sig_rows: List[Dict[str, Any]] = []
    attorney_rows: List[Dict[str, Any]] = []

    # For overview/term stats
    all_terms_for_stats: List[Dict[str, Any]] = []

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

        tracking = data.get("tracking") or {}
        tracked_flag = tracking.get("tracked_flag")
        # Accept True, "true", "TRUE"
        if not (tracked_flag is True or (isinstance(tracked_flag, str) and tracked_flag.strip().lower() == "true")):
            continue  # skip untracked docs

        # Multi-group: we now expect tracking["groups"] to be a list of group dicts
        groups = as_list(tracking.get("groups"))
        if not groups:
            # no group membership -> out of scope for this report
            continue

        doc_sha256 = data.get("doc_sha256") or data.get("sha256") or ""
        filename = data.get("filename", "")
        filing_type = data.get("filing_type", "")
        filing_date = data.get("filing_date", "")
        pdf_page_count = data.get("pdf_page_count", None)
        size_bytes = data.get("size_bytes", None)

        metadata = data.get("metadata") or {}
        case = data.get("case") or {}
        signatures = data.get("signatures") or {}

        # Case flatten (excluding cluster_size and attorneys)
        case_flat: Dict[str, Any] = {}
        case_keys = [k for k in case.keys() if k not in ("attorney_defense", "attorney_prosecution", "cluster_size")]
        for k in case_keys:
            case_flat[f"case_{k}"] = case.get(k)

        # Metadata flatten (all keys) – same for all groups
        metadata_flat: Dict[str, Any] = {}
        for k, v in metadata.items():
            metadata_flat[f"metadata_{k}"] = v

        # --- Per-group expansion ---
        for group in groups:
            if not isinstance(group, dict):
                continue
            raw_pg = group.get("pdf_group")
            if raw_pg is None or str(raw_pg).strip() == "":
                continue

            pdf_group = str(raw_pg).strip()
            font_group = group.get("font_group")

            # Group-specific tracking: prefix keys with tracking_
            group_tracking_flat: Dict[str, Any] = {}
            for k, v in group.items():
                group_tracking_flat[f"tracking_{k}"] = v

            # Document-level row for this (doc, group)
            doc_row: Dict[str, Any] = {
                "pdf_group": pdf_group,
                "font_group": font_group,
                "doc_sha256": doc_sha256,
                "filename": filename,
                "filing_type": filing_type,
                "filing_date": filing_date,
                "pdf_page_count": pdf_page_count,
                "size_bytes": size_bytes,
            }
            doc_row.update(group_tracking_flat)
            doc_row.update(metadata_flat)
            doc_row.update(case_flat)
            docs_rows.append(doc_row)

            # Attorney rows for this group
            for side in ("attorney_defense", "attorney_prosecution"):
                side_list = as_list(case.get(side))
                role = "defense" if side == "attorney_defense" else "prosecution"
                for idx, att in enumerate(side_list):
                    if not isinstance(att, dict):
                        continue
                    a_row: Dict[str, Any] = {
                        "pdf_group": pdf_group,
                        "doc_sha256": doc_sha256,
                        "filename": filename,
                        "case_id": case.get("case_id"),
                        "cluster_id": case.get("cluster_id"),
                        "cluster_name": case.get("cluster_name"),
                        "attorney_role": role,
                        "attorney_index": idx,
                    }
                    for k, v in att.items():
                        a_row[f"attorney_{k}"] = v
                    attorney_rows.append(a_row)

            # Objects rows for this group
            for obj in as_list(data.get("objects")):
                if not isinstance(obj, dict):
                    continue
                o_row: Dict[str, Any] = {
                    "pdf_group": pdf_group,
                    "doc_sha256": doc_sha256,
                    "filename": filename,
                    "case_id": case.get("case_id"),
                }
                for k, v in obj.items():
                    o_row[k] = v
                objects_rows.append(o_row)

            # Terms rows (language.terms only) for this group
            lang = data.get("language") or {}
            for term in as_list(lang.get("terms")):
                if not isinstance(term, dict):
                    continue
                search_group = term.get("search_group")
                search_term = term.get("search_term")
                quantity = term.get("quantity")
                t_row = {
                    "pdf_group": pdf_group,
                    "doc_sha256": doc_sha256,
                    "filename": filename,
                    "case_id": case.get("case_id"),
                    "cluster_id": case.get("cluster_id"),
                    "cluster_name": case.get("cluster_name"),
                    "filing_date": filing_date,
                    "search_group": search_group,
                    "search_term": search_term,
                    "quantity": quantity,
                }
                terms_rows.append(t_row)
                all_terms_for_stats.append(t_row)

            # Signatures.pdfsig row for this group
            pdfsig = safe_get(signatures, "pdfsig", {})
            if isinstance(pdfsig, dict):
                s_row: Dict[str, Any] = {
                    "pdf_group": pdf_group,
                    "doc_sha256": doc_sha256,
                    "filename": filename,
                    "case_id": case.get("case_id"),
                }
                for k, v in pdfsig.items():
                    s_row[f"pdfsig_{k}"] = v
                sig_rows.append(s_row)

    # Convert to DataFrames
    docs_df = pd.DataFrame(docs_rows)
    objects_df = pd.DataFrame(objects_rows)
    terms_df = pd.DataFrame(terms_rows)
    sig_df = pd.DataFrame(sig_rows)
    attorneys_df = pd.DataFrame(attorney_rows)
    stats_terms_df = pd.DataFrame(all_terms_for_stats)

    # ---- 1) pdf_group_documents.csv — curated column set only ----
    if not docs_df.empty:
        desired_cols = [
            "pdf_group", "font_group",
            "case_case_id", "case_cluster_id", "case_cluster_name",
            "filename", "filing_type", "filing_date", "pdf_page_count",
            "tracking_font_hash_date",
            "tracking_font1_sha256", "tracking_font2_sha256",
            "tracking_font3_sha256", "tracking_font4_sha256",
            "tracking_font1_name", "tracking_font2_name",
            "tracking_font3_name", "tracking_font4_name",
            "tracking_font1_type", "tracking_font2_type",
            "tracking_font3_type", "tracking_font4_type",
            "metadata_efile_date",
            "metadata_xmp_create_date", "metadata_xmp_metadata_date",
            "metadata_document_id", "metadata_instance_id",
            "metadata_xmp_author", "metadata_xmp_company",
            "metadata_xmp_creator", "metadata_xmp_title",
            "metadata_xmp_creator_tool", "metadata_xmp_toolkit",
            "metadata_derived_create_newer_by_days",
            "metadata_derived_metadata_newer_by_days",
            "case_defendant_name", "case_case_status",
            "case_case_assigned_judge", "case_filing_date",
            "case_case_year", "case_in_past_flag", "case_link_past_flag",
        ]
        # Ensure all desired columns exist (fill missing with NA)
        for col in desired_cols:
            if col not in docs_df.columns:
                docs_df[col] = pd.NA
        docs_sel = docs_df[desired_cols]
        docs_sel.to_csv(outdir / "pdf_group_documents.csv", index=False)
    else:
        (outdir / "pdf_group_documents.csv").write_text("", encoding="utf-8")

    # ---- 2) pdf_group_attorneys.csv — dedicated attorney table ----
    if not attorneys_df.empty:
        attorneys_df.to_csv(outdir / "pdf_group_attorneys.csv", index=False)
    else:
        (outdir / "pdf_group_attorneys.csv").write_text("", encoding="utf-8")

    # ---- 3) pdf_group_objects.csv ----
    if not objects_df.empty:
        objects_df.to_csv(outdir / "pdf_group_objects.csv", index=False)
    else:
        (outdir / "pdf_group_objects.csv").write_text("", encoding="utf-8")

    # ---- 4) pdf_group_terms.csv ----
    if not terms_df.empty:
        terms_df.to_csv(outdir / "pdf_group_terms.csv", index=False)
    else:
        (outdir / "pdf_group_terms.csv").write_text("", encoding="utf-8")

    # ---- 5) pdf_group_signatures_pdfsig.csv ----
    if not sig_df.empty:
        sig_df.to_csv(outdir / "pdf_group_signatures_pdfsig.csv", index=False)
    else:
        (outdir / "pdf_group_signatures_pdfsig.csv").write_text("", encoding="utf-8")

    # ---- Per-group overview ----
    if not docs_df.empty:
        # For overview, we only care about unique (pdf_group, doc_sha256) combos.
        docs_overview = docs_df.drop_duplicates(subset=["pdf_group", "doc_sha256"]).copy()

        # Merge term stats for counts
        if not stats_terms_df.empty:
            stats_terms_df["quantity"] = pd.to_numeric(stats_terms_df["quantity"], errors="coerce").fillna(0).astype(int)
            term_group = stats_terms_df.groupby("pdf_group", as_index=False).agg(
                n_term_occurrences=("quantity", "sum"),
                n_unique_terms=("search_term", lambda s: int(pd.Series(s).dropna().nunique())),
                n_unique_search_groups=("search_group", lambda s: int(pd.Series(s).dropna().nunique())),
            )
        else:
            term_group = pd.DataFrame(columns=["pdf_group", "n_term_occurrences", "n_unique_terms", "n_unique_search_groups"])

        # object stats
        if not objects_df.empty:
            obj_group = objects_df.groupby("pdf_group", as_index=False).agg(
                n_objects=("object_sha256", "count") if "object_sha256" in objects_df.columns else ("doc_sha256", "count"),
                n_unique_object_sha256=(
                    "object_sha256",
                    lambda s: int(pd.Series(s).dropna().nunique())
                ) if "object_sha256" in objects_df.columns else ("doc_sha256", lambda s: int(pd.Series(s).dropna().nunique())),
                n_unique_object_font_names=(
                    "object_font_name",
                    lambda s: int(pd.Series(s).dropna().nunique())
                ) if "object_font_name" in objects_df.columns else ("doc_sha256", lambda s: 0),
            )
        else:
            obj_group = pd.DataFrame(columns=["pdf_group", "n_objects", "n_unique_object_sha256", "n_unique_object_font_names"])

        # pdfsig stats
        if not sig_df.empty:
            sig_group = sig_df.groupby("pdf_group", as_index=False).agg(
                n_docs_with_pdfsig=("doc_sha256", "nunique")
            )
        else:
            sig_group = pd.DataFrame(columns=["pdf_group", "n_docs_with_pdfsig"])

        # assemble overview
        overview = (
            docs_overview.groupby("pdf_group", as_index=False)
            .agg(
                n_docs=("doc_sha256", "nunique"),
                n_cases=("case_case_id", "nunique") if "case_case_id" in docs_overview.columns else ("doc_sha256", "count"),
                n_clusters=("case_cluster_id", "nunique") if "case_cluster_id" in docs_overview.columns else ("doc_sha256", "count"),
                case_ids=("case_case_id", join_unique) if "case_case_id" in docs_overview.columns else ("doc_sha256", join_unique),
                cluster_ids=("case_cluster_id", join_unique) if "case_cluster_id" in docs_overview.columns else ("doc_sha256", join_unique),
                cluster_names=("case_cluster_name", join_unique) if "case_cluster_name" in docs_overview.columns else ("doc_sha256", join_unique),
                min_filing_date=("filing_date", lambda s: s.dropna().min() if len(s.dropna()) else pd.NA),
                max_filing_date=("filing_date", lambda s: s.dropna().max() if len(s.dropna()) else pd.NA),
                min_pdf_page_count=("pdf_page_count", lambda s: pd.to_numeric(s, errors="coerce").min()),
                max_pdf_page_count=("pdf_page_count", lambda s: pd.to_numeric(s, errors="coerce").max()),
                avg_pdf_page_count=("pdf_page_count", lambda s: float(pd.to_numeric(s, errors="coerce").mean() or 0.0)),
                n_unique_filing_types=("filing_type", lambda s: int(pd.Series(s).dropna().nunique())),
                filing_types=("filing_type", join_unique),
            )
        )

        # merge term / object / sig summary
        overview = overview.merge(term_group, on="pdf_group", how="left")
        overview = overview.merge(obj_group, on="pdf_group", how="left")
        overview = overview.merge(sig_group, on="pdf_group", how="left")

        # fill NaNs with sensible defaults
        for col in [
            "n_term_occurrences",
            "n_unique_terms",
            "n_unique_search_groups",
            "n_objects",
            "n_unique_object_sha256",
            "n_unique_object_font_names",
            "n_docs_with_pdfsig",
        ]:
            if col in overview.columns:
                overview[col] = overview[col].fillna(0).astype(int)

        overview.to_csv(outdir / "pdf_group_overview.csv", index=False)
    else:
        (outdir / "pdf_group_overview.csv").write_text("", encoding="utf-8")

    # ---- Term stats per group ----
    if not stats_terms_df.empty:
        stats_terms_df["quantity"] = pd.to_numeric(stats_terms_df["quantity"], errors="coerce").fillna(0).astype(int)
        term_stats = (
            stats_terms_df.groupby(["pdf_group", "search_group", "search_term"], as_index=False)
            .agg(
                total_quantity=("quantity", "sum"),
                doc_count=("doc_sha256", "nunique"),
                case_count=("case_id", "nunique"),
                cluster_count=("cluster_id", "nunique"),
                first_seen_date=("filing_date", lambda s: s.dropna().min() if len(s.dropna()) else pd.NA),
                last_seen_date=("filing_date", lambda s: s.dropna().max() if len(s.dropna()) else pd.NA),
            )
        )
        term_stats.to_csv(outdir / "pdf_group_term_stats.csv", index=False)
    else:
        (outdir / "pdf_group_term_stats.csv").write_text("", encoding="utf-8")

    # ---- README ----
    readme = """
    MCRO PDF Group Report
    =====================
    This folder was generated by mcro_pdf_group_report.py.

    Scope:
      • Only JSON files where tracking.tracked_flag == true.
      • Only documents that have at least one entry in tracking.groups[*].pdf_group.
      • A document may belong to multiple pdf_groups; it is included in each group's stats.

    Main dumps:
      • pdf_group_documents.csv         — curated column set; 1 row per (doc, pdf_group).
      • pdf_group_attorneys.csv         — 1 row per (pdf_group, attorney entry).
      • pdf_group_objects.csv           — 1 row per (pdf_group, object) with doc context.
      • pdf_group_terms.csv             — 1 row per (pdf_group, doc, term) with quantities.
      • pdf_group_signatures_pdfsig.csv — 1 row per (pdf_group, doc) with pdfsig fields flattened.

    Stats:
      • pdf_group_overview.csv          — high-level metrics per pdf_group.
      • pdf_group_term_stats.csv        — per pdf_group/search_term aggregates.

    You can now use these CSVs to drive deeper analyses:
      • Per-group flag overlays (using your flags CSV).
      • Font hash uniqueness/probabilities per group.
      • Cross-group language, signature, and metadata comparisons.
    """.strip() + "\n"
    (outdir / "README.txt").write_text(readme, encoding="utf-8")

    print(f"[ok] Wrote PDF group report to: {outdir.resolve()}")


if __name__ == "__main__":
    main()