#!/usr/bin/env python3
import argparse
import collections
import datetime as dt
import json
import os
import re
import sys
from pathlib import Path

try:
    import ijson
except ImportError:
    print("ERROR: Missing dependency 'ijson'. Install with:", file=sys.stderr)
    print("  /usr/bin/python3 -m pip install ijson", file=sys.stderr)
    sys.exit(1)

TYPE_ORDER = ["null", "bool", "int", "float", "string", "object", "array"]

def infer_type(value):
    if value is None:
        return "null"
    if isinstance(value, bool):
        return "bool"
    if isinstance(value, int) and not isinstance(value, bool):
        return "int"
    if isinstance(value, float):
        return "float"
    if isinstance(value, str):
        return "string"
    if isinstance(value, dict):
        return "object"
    if isinstance(value, list):
        return "array"
    return type(value).__name__

def scalar_preview(value, max_len=120):
    if isinstance(value, (dict, list)):
        return None
    s = str(value)
    s = re.sub(r"\s+", " ", s).strip()
    return s[:max_len]

def walk(value, path="", stats=None, record_seen=None, sample_limit=3):
    if stats is None:
        stats = {}
    if record_seen is None:
        record_seen = set()

    if path:
        t = infer_type(value)
        entry = stats.setdefault(path, {
            "records": 0,
            "occurrences": 0,
            "types": collections.Counter(),
            "max_string_length": 0,
            "examples": [],
        })
        entry["occurrences"] += 1
        entry["types"][t] += 1
        if path not in record_seen:
            entry["records"] += 1
            record_seen.add(path)
        if isinstance(value, str):
            entry["max_string_length"] = max(entry["max_string_length"], len(value))
            if len(entry["examples"]) < sample_limit:
                pv = scalar_preview(value)
                if pv is not None and pv not in entry["examples"]:
                    entry["examples"].append(pv)
        elif not isinstance(value, (dict, list)):
            if len(entry["examples"]) < sample_limit:
                pv = scalar_preview(value)
                if pv is not None and pv not in entry["examples"]:
                    entry["examples"].append(pv)

    if isinstance(value, dict):
        for k, v in value.items():
            child = f"{path}.{k}" if path else k
            walk(v, child, stats, record_seen, sample_limit)
    elif isinstance(value, list):
        item_path = f"{path}[]" if path else "[]"
        if not value:
            entry = stats.setdefault(item_path, {
                "records": 0,
                "occurrences": 0,
                "types": collections.Counter(),
                "max_string_length": 0,
                "examples": [],
            })
            if item_path not in record_seen:
                entry["records"] += 1
                record_seen.add(item_path)
            entry["occurrences"] += 1
            entry["types"]["empty_array"] += 1
        else:
            for item in value:
                walk(item, item_path, stats, record_seen, sample_limit)

def mysql_type_for(field):
    types = set(field["types"].keys()) - {"null", "empty_array"}
    maxlen = field.get("max_string_length", 0)

    if not types:
        return "TEXT"
    if types <= {"bool"}:
        return "TINYINT(1)"
    if types <= {"int"}:
        return "BIGINT"
    if types <= {"int", "float"}:
        return "DOUBLE"
    if types <= {"string"}:
        if maxlen <= 32:
            return "VARCHAR(32)"
        if maxlen <= 64:
            return "VARCHAR(64)"
        if maxlen <= 128:
            return "VARCHAR(128)"
        if maxlen <= 255:
            return "VARCHAR(255)"
        if maxlen <= 1000:
            return "VARCHAR(1024)"
        return "TEXT"
    return "JSON"

def column_name(path):
    p = path.replace("[]", "_item")
    p = p.replace(".", "_")
    p = re.sub(r"[^A-Za-z0-9_]+", "_", p)
    return p.lower().strip("_")[:64]

def write_reports(stats, total_records, out_dir):
    out_dir.mkdir(parents=True, exist_ok=True)

    ordered = sorted(stats.items(), key=lambda kv: kv[0])

    report_json = {
        "generated_at_utc": dt.datetime.utcnow().isoformat(timespec="seconds") + "Z",
        "total_records": total_records,
        "fields": {}
    }
    for path, e in ordered:
        report_json["fields"][path] = {
            "records": e["records"],
            "record_percentage": round((e["records"] / total_records * 100), 6) if total_records else 0,
            "occurrences": e["occurrences"],
            "types": dict(e["types"]),
            "max_string_length": e["max_string_length"],
            "examples": e["examples"],
            "suggested_mysql_type": mysql_type_for(e),
            "suggested_column": column_name(path),
        }

    (out_dir / "field_inventory.json").write_text(
        json.dumps(report_json, indent=2, ensure_ascii=False),
        encoding="utf-8"
    )

    with (out_dir / "field_inventory.csv").open("w", encoding="utf-8", newline="") as f:
        import csv
        w = csv.writer(f)
        w.writerow([
            "field_path", "records", "record_percentage", "occurrences",
            "types", "max_string_length", "suggested_mysql_type",
            "suggested_column", "examples"
        ])
        for path, e in ordered:
            w.writerow([
                path,
                e["records"],
                round((e["records"] / total_records * 100), 6) if total_records else 0,
                e["occurrences"],
                json.dumps(dict(e["types"]), sort_keys=True),
                e["max_string_length"],
                mysql_type_for(e),
                column_name(path),
                " | ".join(e["examples"]),
            ])

    # Conservative proposed public/master schema:
    # top-level scalar fields and address scalar fields only.
    # Nested arrays/objects are intentionally excluded from automatic import.
    preferred = [
        "firstName", "middleName", "lastName", "name", "title",
        "address.streetNumber", "address.street", "address.unit",
        "address.city", "address.state", "address.zipCode",
        "address.county", "address.country",
        "county", "voterStatus"
    ]
    existing = [p for p in preferred if p in stats]

    # Also surface candidate scalar fields not in preferred for review, but do not auto-import them.
    candidates = []
    for path, e in ordered:
        if path in existing:
            continue
        if "[]" in path:
            continue
        types = set(e["types"].keys()) - {"null"}
        if "object" in types or "array" in types:
            continue
        candidates.append(path)

    schema = {
        "table": "voters_public",
        "generated_from_audit": True,
        "include_fields": existing,
        "review_candidates_not_auto_included": candidates,
        "excluded_by_policy_examples": [
            "_id", "__v", "compositeKey", "createdAt", "updatedAt",
            "processedAt", "importBatch", "embedding", "flags",
            "dateOfBirth", "partyAffiliation", "voterId", "signatures"
        ]
    }
    (out_dir / "mysql_import_map.json").write_text(
        json.dumps(schema, indent=2),
        encoding="utf-8"
    )

    cols = []
    for path in existing:
        e = stats[path]
        cname = column_name(path)
        ctype = mysql_type_for(e)
        cols.append((cname, ctype, path))

    sql = []
    sql.append("CREATE DATABASE IF NOT EXISTS voter_data CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;")
    sql.append("USE voter_data;")
    sql.append("")
    sql.append("CREATE TABLE IF NOT EXISTS voters_public (")
    sql.append("  id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,")
    for cname, ctype, path in cols:
        sql.append(f"  `{cname}` {ctype} NULL,")
    sql.append("  PRIMARY KEY (id)")
    sql.append(") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;")
    sql.append("")
    sql.append("-- Recommended indexes for the public search product")
    available_cols = {c for c,_,_ in cols}
    def add_index(name, cols_):
        if all(c in available_cols for c in cols_):
            sql.append(f"CREATE INDEX `{name}` ON voters_public ({', '.join('`'+c+'`' for c in cols_)});")
    add_index("idx_last_name", ["lastname"])
    add_index("idx_zip_code", ["address_zipcode"])
    add_index("idx_city", ["address_city"])
    add_index("idx_street", ["address_street"])
    add_index("idx_last_city", ["lastname", "address_city"])
    add_index("idx_zip_street", ["address_zipcode", "address_street"])
    add_index("idx_city_street_number", ["address_city", "address_street", "address_streetnumber"])

    ft_candidates = [c for c in [
        "firstname","middlename","lastname","name","address_street","address_city","address_county"
    ] if c in available_cols]
    if ft_candidates:
        sql.append(
            "ALTER TABLE voters_public ADD FULLTEXT INDEX `ft_public_search` (" +
            ", ".join(f"`{c}`" for c in ft_candidates) + ");"
        )

    (out_dir / "generated_mysql_schema.sql").write_text("\n".join(sql) + "\n", encoding="utf-8")

def main():
    ap = argparse.ArgumentParser(description="Audit every field/type in a giant JSON array without loading it into RAM.")
    ap.add_argument("--input", required=True, help="Path to giant JSON array file")
    ap.add_argument("--output-dir", default="./voter_audit", help="Directory for audit outputs")
    ap.add_argument("--progress-every", type=int, default=100000)
    ap.add_argument("--sample-limit", type=int, default=3)
    args = ap.parse_args()

    src = Path(args.input)
    if not src.is_file():
        print(f"ERROR: Input file not found: {src}", file=sys.stderr)
        sys.exit(1)

    stats = {}
    total = 0

    with src.open("rb") as f:
        for doc in ijson.items(f, "item"):
            total += 1
            seen = set()
            walk(doc, stats=stats, record_seen=seen, sample_limit=args.sample_limit)
            if total % args.progress_every == 0:
                print(f"Scanned {total:,} records; discovered {len(stats):,} field paths", flush=True)

    out_dir = Path(args.output_dir)
    write_reports(stats, total, out_dir)

    print("")
    print("AUDIT COMPLETE")
    print(f"Records scanned: {total:,}")
    print(f"Field paths:     {len(stats):,}")
    print(f"Output directory: {out_dir.resolve()}")
    print("")
    print("Created:")
    print("  field_inventory.json")
    print("  field_inventory.csv")
    print("  mysql_import_map.json")
    print("  generated_mysql_schema.sql")

if __name__ == "__main__":
    main()