"""
Upload a local evaluation dataset of Solidity files to Hugging Face Hub.

Reads .sol files organized by vulnerability type in subfolders, converts them
into a structured HF Dataset, and pushes to the Hub.

Expected folder structure:
    Evaluation_Dataset/
        Longer_example/
            Reentrancy/
                3.sol
                4.sol
                ...
            Access Control/
                1.sol
                2.sol
                ...
            Integer Overflow/
                ...
            Timestamp Dependence/
                ...
            Unchecked Low-Level Calls/
                ...
        Short_example/          (optional — additional source folders)
            Reentrancy/
                ...

Output dataset columns:
    - code (str):               Full Solidity source code
    - vulnerability_type (str): Folder name (e.g. "Reentrancy")
    - filename (str):           Original filename (e.g. "3.sol")
    - source (str):             Parent folder (e.g. "Longer_example")
    - filepath (str):           Relative path from root (e.g. "Longer_example/Reentrancy/3.sol")
    - num_lines (int):          Number of lines in the file
    - num_chars (int):          Number of characters in the file

Usage:
    # Upload with defaults:
    python upload_eval_dataset.py --data_dir ./Evaluation_Dataset

    # Specify Hub dataset name:
    python upload_eval_dataset.py --data_dir ./Evaluation_Dataset --hub_dataset jhsu12/solidity-eval-dataset

    # Dry run — see what would be uploaded:
    python upload_eval_dataset.py --data_dir ./Evaluation_Dataset --dry_run

    # Make it private:
    python upload_eval_dataset.py --data_dir ./Evaluation_Dataset --private
"""

import argparse
import os
import sys
from datasets import Dataset
from huggingface_hub import create_repo


def parse_args():
    parser = argparse.ArgumentParser(
        description="Upload Solidity evaluation dataset to Hugging Face Hub."
    )
    parser.add_argument(
        "--data_dir", type=str, required=True,
        help="Path to the root evaluation dataset folder"
    )
    parser.add_argument(
        "--hub_dataset", type=str, default="jhsu12/solidity-vulnerability-eval-dataset",
        help="Hub dataset ID (default: jhsu12/solidity-vulnerability-eval-dataset)"
    )
    parser.add_argument(
        "--private", action="store_true", default=False,
        help="Create as private dataset"
    )
    parser.add_argument(
        "--dry_run", action="store_true", default=False,
        help="Show what would be uploaded without actually uploading"
    )
    return parser.parse_args()


def scan_sol_files(data_dir):
    """
    Walk the directory tree and collect all .sol files.

    Handles two levels of nesting:
        data_dir/VulnType/*.sol           → source = "" (flat)
        data_dir/Source/VulnType/*.sol     → source = Source (nested)
    """
    records = []
    data_dir = os.path.abspath(data_dir)

    for entry in sorted(os.listdir(data_dir)):
        entry_path = os.path.join(data_dir, entry)
        if not os.path.isdir(entry_path):
            continue

        # Check if this directory contains .sol files directly (flat: VulnType/*.sol)
        sol_files_here = [f for f in os.listdir(entry_path) if f.endswith(".sol")]
        subdirs_here = [d for d in os.listdir(entry_path)
                        if os.path.isdir(os.path.join(entry_path, d))]

        if sol_files_here and not subdirs_here:
            # Flat: data_dir/VulnType/*.sol
            vuln_type = entry
            for sol_file in sorted(sol_files_here):
                filepath = os.path.join(entry_path, sol_file)
                records.append({
                    "vulnerability_type": vuln_type,
                    "filename": sol_file,
                    "source": "",
                    "filepath": os.path.join(entry, sol_file),
                    "abs_path": filepath,
                })
        else:
            # Nested: data_dir/Source/VulnType/*.sol
            source_name = entry
            for vuln_dir in sorted(os.listdir(entry_path)):
                vuln_path = os.path.join(entry_path, vuln_dir)
                if not os.path.isdir(vuln_path):
                    continue

                vuln_type = vuln_dir
                for sol_file in sorted(os.listdir(vuln_path)):
                    if not sol_file.endswith(".sol"):
                        continue
                    filepath = os.path.join(vuln_path, sol_file)
                    records.append({
                        "vulnerability_type": vuln_type,
                        "filename": sol_file,
                        "source": source_name,
                        "filepath": os.path.join(source_name, vuln_dir, sol_file),
                        "abs_path": filepath,
                    })

    return records


def read_sol_files(records):
    """Read the actual Solidity code from each file."""
    data = {
        "code": [],
        "vulnerability_type": [],
        "filename": [],
        "source": [],
        "filepath": [],
        "num_lines": [],
        "num_chars": [],
    }

    failed = 0
    for rec in records:
        try:
            with open(rec["abs_path"], "r", encoding="utf-8", errors="replace") as f:
                code = f.read()
        except Exception as e:
            print(f"  ⚠️  Failed to read {rec['filepath']}: {e}")
            failed += 1
            continue

        data["code"].append(code)
        data["vulnerability_type"].append(rec["vulnerability_type"])
        data["filename"].append(rec["filename"])
        data["source"].append(rec["source"])
        data["filepath"].append(rec["filepath"])
        data["num_lines"].append(len(code.splitlines()))
        data["num_chars"].append(len(code))

    return data, failed


def main():
    args = parse_args()
    data_dir = os.path.abspath(args.data_dir)

    print("=" * 60)
    print("  Upload Solidity Evaluation Dataset")
    print("=" * 60)

    if not os.path.isdir(data_dir):
        print(f"\n❌ Directory not found: {data_dir}")
        sys.exit(1)

    # ── Scan ──────────────────────────────────────────────────────────────────
    print(f"\n🔍 Scanning: {data_dir}")
    records = scan_sol_files(data_dir)

    if not records:
        print(f"\n❌ No .sol files found!")
        print(f"   Expected structure:")
        print(f"   {data_dir}/")
        print(f"     Longer_example/")
        print(f"       Reentrancy/")
        print(f"         3.sol")
        print(f"         4.sol")
        print(f"       Access Control/")
        print(f"         ...")
        sys.exit(1)

    # ── Summary ───────────────────────────────────────────────────────────────
    # Count by vulnerability type
    from collections import Counter
    vuln_counts = Counter(r["vulnerability_type"] for r in records)
    source_counts = Counter(r["source"] for r in records)

    print(f"\n   Found {len(records)} .sol files\n")

    if any(r["source"] for r in records):
        print(f"   By source:")
        for source, count in sorted(source_counts.items()):
            label = source if source else "(root)"
            print(f"     {label}: {count} files")
        print()

    print(f"   By vulnerability type:")
    for vuln, count in sorted(vuln_counts.items()):
        print(f"     {vuln}: {count} files")

    # ── Read files ────────────────────────────────────────────────────────────
    print(f"\n📖 Reading Solidity files...")
    data, failed = read_sol_files(records)
    total = len(data["code"])

    if failed:
        print(f"   ⚠️  {failed} files failed to read")
    print(f"   ✅ {total} files read successfully")

    # Stats
    avg_lines = sum(data["num_lines"]) / total if total else 0
    avg_chars = sum(data["num_chars"]) / total if total else 0
    max_lines = max(data["num_lines"]) if total else 0
    min_lines = min(data["num_lines"]) if total else 0
    print(f"\n   📊 Stats:")
    print(f"      Lines: min={min_lines}, avg={avg_lines:.0f}, max={max_lines}")
    print(f"      Chars: avg={avg_chars:.0f}")

    # ── Create dataset ────────────────────────────────────────────────────────
    print(f"\n📦 Creating HF Dataset...")
    dataset = Dataset.from_dict(data)
    print(f"   {dataset}")

    # Show a sample
    print(f"\n   Sample row:")
    sample = dataset[0]
    print(f"     vulnerability_type: {sample['vulnerability_type']}")
    print(f"     filename:           {sample['filename']}")
    print(f"     source:             {sample['source']}")
    print(f"     num_lines:          {sample['num_lines']}")
    print(f"     code:               {sample['code'][:100]}...")

    if args.dry_run:
        print(f"\n🔍 DRY RUN — would push to {args.hub_dataset}")
        print(f"   {total} samples, {len(vuln_counts)} vulnerability types")
        sys.exit(0)

    # ── Push to Hub ───────────────────────────────────────────────────────────
    print(f"\n🚀 Pushing to {args.hub_dataset}...")

    # Create repo first
    try:
        create_repo(
            args.hub_dataset,
            repo_type="dataset",
            private=args.private,
            exist_ok=True,
        )
    except Exception as e:
        print(f"   ❌ Failed to create repo: {e}")
        sys.exit(1)

    dataset.push_to_hub(args.hub_dataset, private=args.private)

    print(f"\n   ✅ Pushed to https://huggingface.co/datasets/{args.hub_dataset}")

    # ── Done ──────────────────────────────────────────────────────────────────
    print(f"\n{'=' * 60}")
    print(f"  Dataset uploaded!")
    print(f"  URL:     https://huggingface.co/datasets/{args.hub_dataset}")
    print(f"  Samples: {total}")
    print(f"  Types:   {', '.join(sorted(vuln_counts.keys()))}")
    print(f"\n  Load it with:")
    print(f'    from datasets import load_dataset')
    print(f'    ds = load_dataset("{args.hub_dataset}")')
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()