import argparse
import sys
import json
from pathlib import Path
from marco.core import repository
from marco.cli.interactive import build_interactive_config
from marco.cli import web_serve

def main():
    parser = argparse.ArgumentParser(prog="marco", description="Marco Dataset Versioning System")
    subparsers = parser.add_subparsers(dest="command", required=True)
    
    p_init = subparsers.add_parser("init", help="Initialize a marco repository")
    p_init.add_argument("path", nargs="?", default=".", help="Path to initialize")
    
    p_upload = subparsers.add_parser("upload", help="Create a dataset version")
    p_upload.add_argument("file", help="Path to raw dataset string/CSV/TSV")
    p_upload.add_argument("-c", "--config", help="Path to config.json")
    p_upload.add_argument("-t", "--tag", help="Tag for the version")
    p_upload.add_argument("-u", "--user", default="unknown", help="User name")
    
    subparsers.add_parser("list", help="List all dataset versions")
    
    p_diff = subparsers.add_parser("diff", help="Compare two dataset versions")
    p_diff.add_argument("v1", help="Version Hash or Tag 1")
    p_diff.add_argument("v2", help="Version Hash or Tag 2")
    p_diff.add_argument(
        "--target", choices=["raw", "preprocessed", "config"], default="raw",
        help="Which file to diff: 'raw' (raw.txt, default), 'preprocessed' (preprocessed.tsv), 'config' (config.json)"
    )
    p_diff.add_argument(
        "--context", type=int, default=3,
        help="Lines of context around each changed hunk (default: 3, like git diff)"
    )
    p_diff.add_argument(
        "--save", metavar="FOLDER", default=None,
        help="Save a Markdown metrics report to the specified folder instead of printing a diff"
    )
    
    p_restore = subparsers.add_parser("restore", aliases=["checkout"], help="Restore processed data to your workspace")
    p_restore.add_argument("version", help="Version Hash or Tag")
    p_restore.add_argument("-o", "--out", default="data.tsv", help="Where to save the processed file")
    
    p_export = subparsers.add_parser("export", help="Export a version to tar.gz")
    p_export.add_argument("version", help="Version Hash or Tag")
    p_export.add_argument("dest", help="Destination folder")
    
    p_import = subparsers.add_parser("import", help="Import a version from tar.gz")
    p_import.add_argument("tarball", help="Path to exported tar.gz")

    p_delete = subparsers.add_parser("delete", aliases=["rm"], help="Delete a dataset version and update lineage")
    p_delete.add_argument("version", help="Version Hash or Tag to delete")

    p_analytics = subparsers.add_parser("token-analytics", help="Analyze token distribution shift between two versions")
    p_analytics.add_argument("v1", help="Version Hash or Tag 1")
    p_analytics.add_argument("v2", help="Version Hash or Tag 2")

    # subparsers.add_parser("generate-web", help="Serve the Marco React UI")
    subparsers.add_parser("generate-web", help=argparse.SUPPRESS)

    p_verify = subparsers.add_parser("verify", help="Verify reproducibility of a dataset version")
    p_verify.add_argument("version", help="Version hash or tag to verify")

    subparsers.add_parser("verify-all", help="Verify reproducibility of all dataset versions")
    
    args = parser.parse_args()
    repo_path = Path.cwd()
    
    if args.command == "init":
        p = args.path if args.path else "."
        repository.init_repo(p)
        print(f"Initialized Marco repository in {Path(p).resolve() / '.marco'}")
        
    elif args.command == "upload":
        raw_path = Path(args.file)
        if not raw_path.exists():
            print(f"Error: File {raw_path} does not exist.")
            sys.exit(1)
            
        if args.config:
            config = json.loads(Path(args.config).read_text(encoding='utf-8'))
        else:
            config = build_interactive_config()
            
        tags = [args.tag] if args.tag else []
        vid = repository.create_version(raw_path, config, repo_path, user=args.user, tags=tags)
        print(f"Successfully processed and generated version: {vid}")
        
    elif args.command == "list":
        try:
            versions = repository.list_versions(repo_path)
            print(f"{'Version':<10} | {'Created At':<20} | {'User':<10} | {'Tags':<20}")
            print("-" * 68)
            for v in versions:
                short = v['version_id'][:8]
                d = v.get('created_at', '')[:19].replace('T', ' ')
                u = v.get('created_by', 'unknown')
                t = ", ".join(v.get('tags', []))
                print(f"{short:<10} | {d:<20} | {u:<10} | {t:<20}")
        except Exception as e:
            print(f"Failed to list versions: {e}")
            
    elif args.command == "diff":
        from marco.core.comparator import compare_versions, generate_markdown_report, diff_versions_text
        
        try:
            diff_output = diff_versions_text(
                args.v1, args.v2, repo_path,
                target=args.target,
                context=args.context,
            )

            if args.save:
                import re
                if not diff_output:
                    print(f"No differences found between {args.v1} and {args.v2} for target '{args.target}'.")
                    return
                
                # Strip ANSI codes for clear file saving
                ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
                clean_text = ansi_escape.sub('', diff_output)
                
                save_dir = Path(args.save)
                save_dir.mkdir(parents=True, exist_ok=True)
                
                out_path = save_dir / f"diff_{args.v1}_{args.v2}_{args.target}.md"
                
                # Wrap it nicely depending on the format type
                md_content = f"# Diff Report: `{args.target}`\n\n"
                if args.target == "preprocessed":
                    md_content += "```text\n" + clean_text + "```\n"
                else:
                    md_content += "```diff\n" + clean_text + "```\n"
                    
                out_path.write_text(md_content, encoding='utf-8')
                print(f"✅ Visual diff output saved to: {out_path.resolve()}")
            else:
                if diff_output:
                    print(diff_output, end="")
                else:
                    print("No differences found.")
        except Exception as e:
            print(f"Failed to compare versions: {e}")
            
    elif args.command in ["restore", "checkout"]:
        try:
            dest = repository.restore_version(args.version, args.out, repo_path)
            print(f"✅ Successfully restored dataset to: {dest.resolve()}")
        except Exception as e:
            print(f"Failed to restore dataset: {e}")
            
    elif args.command == "export":
        tar_path = repository.export_version(args.version, args.dest, repo_path)
        print(f"Exported to {tar_path}")
        
    elif args.command == "import":
        vid = repository.import_version(args.tarball, repo_path)
        print(f"Successfully imported version {vid}")

    elif args.command in ["delete", "rm"]:
        try:
            vid = repository.delete_version(args.version, repo_path)
            print(f"✅ Successfully deleted version {vid} and updated lineage.")
        except Exception as e:
            print(f"Failed to delete version: {e}")

    elif args.command == "token-analytics":
        from marco.token_analytics.distribution import generate_distribution_report
        try:
            v1_id = repository.resolve_version(args.v1, repo_path)
            v2_id = repository.resolve_version(args.v2, repo_path)
            v1_dir = repo_path / ".marco" / "versions" / v1_id
            v2_dir = repo_path / ".marco" / "versions" / v2_id
            report = generate_distribution_report(v1_dir, v2_dir)
            import json
            print(json.dumps(report, indent=2))
        except Exception as e:
            print(f"Failed to run token analytics: {e}")

    elif args.command == "generate-web":
        web_serve.serve()

    elif args.command == "verify":
        from marco.reproducibility_proof.verify import verify
        from marco.core.repository import resolve_version
        try:
            vid = resolve_version(args.version, repo_path)
            passed = verify(vid, repo_path)
            sys.exit(0 if passed else 1)
        except Exception as e:
            print(f"❌ Error: {e}")
            sys.exit(1)

    elif args.command == "verify-all":
        from marco.reproducibility_proof.verify import verify_all
        results = verify_all(repo_path)
        failed = [v for v, ok in results.items() if not ok]
        sys.exit(1 if failed else 0)

if __name__ == "__main__":
    main()

