import argparse
import json
from pathlib import Path
from .distribution import generate_distribution_report

def main():
    parser = argparse.ArgumentParser(description="Analyze token distribution shift between two marco dataset versions.")
    parser.add_argument("v1_dir", type=str, help="Path to version 1 directory")
    parser.add_argument("v2_dir", type=str, help="Path to version 2 directory")
    
    args = parser.parse_args()
    
    v1_path = Path(args.v1_dir)
    v2_path = Path(args.v2_dir)
    
    if not v1_path.exists() or not v2_path.exists():
        print("Error: Both version directories must exist.")
        return

    try:
        report = generate_distribution_report(v1_path, v2_path)
        print(json.dumps(report, indent=2))
    except Exception as e:
        print(f"Error generating report: {e}")

if __name__ == "__main__":
    main()
