#!/usr/bin/env python3
import os
import sys
import argparse
from pathlib import Path
from dotenv import load_dotenv
from huggingface_hub import HfApi

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--yes", action="store_true", help="Auto confirm deletion")
    args = parser.parse_args()

    REPO_ROOT = Path(__file__).resolve().parents[1]
    load_dotenv(REPO_ROOT / ".env")

    token = os.getenv("HF_TOKEN")
    repo_id = os.getenv("HF_SPACE_REPO", "ThinkcatLab/LiveHouse-TS")

    if not token:
        print("HF_TOKEN not found in .env", file=sys.stderr)
        return 1

    print(f"Connecting to Hugging Face Space: {repo_id}")
    api = HfApi(token=token)

    # List all files in the repository
    try:
        files = api.list_repo_files(repo_id=repo_id, repo_type="space")
    except Exception as e:
        print(f"Failed to list files: {e}", file=sys.stderr)
        return 1

    allowed_dirs = {
        "aggregates",
        "arima",
        "chronos_bolt",
        "chronos_2",
        "ets",
        "moirai_2_0",
        "moving_average",
        "seasonal_naive",
        "sundial",
        "tabpfn_ts",
        "timesfm_2_5",
        "tirex",
        "toto_1_0",
    }

    allowed_files = {
        "online_status.json",
        "dataset_properties.csv",
        "dataset_eval_intervals.json",
        "eval_history.jsonl",
        "eval_state.json",
        "baseline_rank_history.csv",
        "baseline_rank_history_daily.csv",
        "baseline_rank_history_weekly.csv",
        "model_allowlist.json",
    }

    files_to_delete = []

    for f in files:
        if not f.startswith("results/"):
            continue

        parts = f.split("/")
        if len(parts) < 2:
            continue

        # check if it is a file directly under results/
        if len(parts) == 2:
            filename = parts[1]
            if filename in {"evaluation_metrics.jsonl", "evaluation_metrics_canonical.jsonl"}:
                files_to_delete.append(f)
            elif filename not in allowed_files:
                files_to_delete.append(f)
        else:
            # it is inside a subdirectory under results/
            sub_dir = parts[1]
            if sub_dir == "aggregates" and (
                parts[-1].endswith(".zip")
                or parts[-1] == "evaluation_metrics_canonical.jsonl"
                or "canonical_all_history" in parts
            ):
                files_to_delete.append(f)
            elif sub_dir == "saved_evals" or "forecasts" in parts:
                files_to_delete.append(f)
            elif sub_dir not in allowed_dirs:
                files_to_delete.append(f)

    if not files_to_delete:
        print("No extra files to delete on remote Space.")
        return 0

    print(f"Found {len(files_to_delete)} files to delete from Hugging Face Space:")
    for f in files_to_delete[:20]:
        print(f" - {f}")
    if len(files_to_delete) > 20:
        print(f" ... and {len(files_to_delete) - 20} more.")

    if not args.yes:
        confirm = input("Proceed with deletion? (y/n): ").strip().lower()
        if confirm != "y":
            print("Cancelled.")
            return 0

    # Delete the files using api.delete_file or commit operations
    from huggingface_hub import CommitOperationDelete

    operations = [CommitOperationDelete(path_in_repo=f) for f in files_to_delete]

    # We batch commit in chunks of 500 files to avoid payload limits
    chunk_size = 500
    for i in range(0, len(operations), chunk_size):
        chunk = operations[i:i+chunk_size]
        print(f"Deleting chunk {i // chunk_size + 1} ({len(chunk)} files)...")
        try:
            api.create_commit(
                repo_id=repo_id,
                repo_type="space",
                operations=chunk,
                commit_message=f"Clean remote results: delete extra model files (part {i // chunk_size + 1})",
            )
        except Exception as e:
            print(f"Failed to delete chunk: {e}", file=sys.stderr)
            return 1

    print("Deletion completed successfully.")
    return 0

if __name__ == "__main__":
    sys.exit(main())