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")
    p_upload.add_argument("-r", "--results", help="Path to training results file (.txt, .csv, etc.)")
    
    p_attach = subparsers.add_parser("attach-results", aliases=["attach"], help="Attach a results file to an existing version")
    p_attach.add_argument("version", help="Version Hash or Tag")
    p_attach.add_argument("results", help="Path to training results file (.txt, .csv, etc.)")
    
    subparsers.add_parser("list", help="List all dataset versions")
    subparsers.add_parser("lineage", help="Show the lineage graph of all 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")

    p_eval = subparsers.add_parser("evaluate", help="Evaluate the structural health of a dataset version")
    p_eval.add_argument("version", help="Version Hash or Tag")

    p_drift = subparsers.add_parser("drift", help="Detect performance shift between two dataset versions by executing an internal test model against new data")
    p_drift.add_argument("v1", help="Baseline Version")
    p_drift.add_argument("v2", help="New Version")

    # 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")

    p_nb = subparsers.add_parser("naive-bayes", help="Run Naive_base_Model.py pipeline against a restored dataset version")
    p_nb.add_argument("version", help="Version hash or tag to run against")

    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)

        # Auto-init if no .marco directory found yet
        if not (repo_path / '.marco').exists():
            print(f"No Marco repository found — initializing one in {repo_path} ...")
            repository.init_repo(repo_path)
            print(f"Initialized Marco repository in {repo_path / '.marco'}\n")

        if args.results:
            if not Path(args.results).exists():
                print(f"Error: Results file {args.results} 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, results_path=args.results)
        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 in ["attach-results", "attach"]:
        try:
            vid, dest = repository.attach_results(args.version, args.results, repo_path)
            print(f"✅ Successfully attached results to version {vid}: {dest.name}")
        except Exception as e:
            print(f"❌ Failed to attach results: {e}")
            sys.exit(1)
            
    elif args.command == "lineage":
        try:
            print(repository.get_lineage(repo_path))
        except Exception as e:
            print(f"Failed to generate lineage: {e}")
            
    elif args.command == "diff":
        from marco.core.comparator import compare_versions, generate_markdown_report, diff_versions_text, generate_diff_report
        
        try:
            # Generate human-readable metric summary
            report = compare_versions(args.v1, args.v2, repo_path)
            print(f"\n📊 Summary: {args.v1} → {args.v2}")
            for k, v in report['metrics'].items():
                name = k.replace('_', ' ').title()
                diff_val = v['diff']
                diff_sign = "+" if diff_val > 0 else ""
                pct_str = v['pct_change']
                
                if pct_str != "N/A":
                    direction = "increase" if diff_val > 0 else "decrease" if diff_val < 0 else "no change"
                    if diff_val != 0:
                        print(f"  ► {name}: {pct_str} {direction} ({diff_sign}{diff_val}) | {v['v1']} → {v['v2']}")
                    else:
                        print(f"  ► {name}: No change | {v['v1']}")
                else:
                    print(f"  ► {name}: N/A | {v['v1']} → {v['v2']}")

            vocab_overlap = report.get('vocab', {})
            if vocab_overlap:
                jaccard = vocab_overlap.get('jaccard', "N/A")
                print(f"  ► Vocab Match (Jaccard): {jaccard}")

            print("─────────────────────────────────────────────────────────────\n")
            
            diff_output = diff_versions_text(
                args.v1, args.v2, repo_path,
                target=args.target,
                context=args.context,
            )

            if args.save:
                if not diff_output:
                    print(f"No differences found between {args.v1} and {args.v2} for target '{args.target}'.")
                    return
                
                save_dir = Path(args.save)
                out_path = generate_diff_report(
                    args.v1, args.v2, repo_path, args.target, diff_output, save_dir
                )
                print(f"✅ Visual diff output saved to: {out_path.resolve()}")
            else:
                if diff_output:
                    print(diff_output, end="")
                else:
                    print(f"No line-by-line differences found in '{args.target}'.")

        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)
            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 == "evaluate":
        from marco.core.evaluate import evaluate_version
        try:
            print(f"\n📊 Evaluating Structural Health of Dataset '{args.version}'...\n")
            report, out_file = evaluate_version(args.version, repo_path)
            
            # Print a neat summary formatted output directly to terminal
            print(f"  ► Pipeline Determinism... [{report['metrics']['pipeline_determinism_grade']} Grade]")
            volatility_pct = int(report['metrics']['data_volatility_index_dvi'] * 100)
            print(f"  ► Computing Data Volatility (vs. Raw)... [{volatility_pct}% Volatile]")
            print(f"  ► Running Reproducibility Proofs... [{report['metrics']['reproducibility_confidence_score_rcs']} Confident]\n")
            
            if report['overall_health_status'] == 'EXCELLENT':
                print(f"✅ Evaluation Complete! Overall Health: {report['overall_health_status']}")
            elif report['overall_health_status'] == 'WARNING':
                print(f"⚠️ Evaluation Complete! Overall Health: {report['overall_health_status']}")
            else:
                print(f"❌ Evaluation Complete! Overall Health: {report['overall_health_status']}")
                
            print(f"📄 Generating full report: {out_file.resolve()}\n")
            for insight in report['insights']:
                print(f"   {insight}")
                
        except Exception as e:
            print(f"❌ Error: {e}")
            sys.exit(1)

    elif args.command == "drift":
        from marco.model_drift import run_full_drift_check
        from marco.core.repository import resolve_version
        try:
            v1_id = resolve_version(args.v1, repo_path)
            v2_id = resolve_version(args.v2, repo_path)
            v1_dir = repo_path / ".marco" / "versions" / v1_id
            v2_dir = repo_path / ".marco" / "versions" / v2_id
            
            run_full_drift_check(v1_dir, v2_dir)
        except Exception as e:
            print(f"❌ Error: {e}")
            sys.exit(1)

    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 == "naive-bayes":
        import subprocess
        import os
        try:
            temp_data = "./my_training_data.tsv"
            result_file = "results.txt"
            
            print(f"🔄 Restoring version {args.version} to {temp_data}...")
            repository.restore_version(args.version, temp_data, repo_path)
            
            print(f"🚀 Running Naive_base_Model.py on {temp_data}...")
            process = subprocess.run(
                ["python", "Naive_base_Model.py", "--data-path", temp_data], 
                capture_output=True, text=True, encoding='utf-8', errors='replace'
            )
            
            # Print the output so the user still sees it!
            print(process.stdout)
            if process.stderr:
                print(process.stderr)
            
            if process.returncode != 0:
                raise RuntimeError("Naive Bayes Model script failed.")
            
            # If the user script didn't explicitly write results.txt, we automatically harvest STDOUT
            if not os.path.exists(result_file):
                with open(result_file, "w", encoding="utf-8") as f:
                    f.write(process.stdout)
                
            print(f"📎 Attaching generated {result_file} to version {args.version}...")
            repository.attach_results(args.version, result_file, repo_path)
            
            # Clean up
            if os.path.exists(temp_data): os.remove(temp_data)
            if os.path.exists(result_file): os.remove(result_file)
            
            print("✅ Naive Bayes Pipeline successfully integrated and results attached!")
        except Exception as e:
            print(f"❌ Pipeline Failed: {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()

