"""
kaggle_train.py — Automate Kaggle Training
============================================
Full automation of the Kaggle training pipeline:

  Step 1  Zip crops_dataset/ and upload (or update) it as a Kaggle Dataset
  Step 2  Patch the notebook's DATA_DIR to match the uploaded dataset path
  Step 3  Push the notebook as a Kaggle Kernel and trigger a run
  Step 4  Poll the kernel status every 30 s until complete or failed
  Step 5  Download best.pt + class_names.json into your local runs/classify/

Requirements
------------
  pip install kaggle
  
  Set up credentials (one-time, pick one):
  
  Option A — New token file (what Kaggle now issues):
    Go to https://www.kaggle.com/settings/api → "Create New Token"
    Run the command Kaggle shows you, e.g.:
      mkdir -p ~/.kaggle && echo YOUR_TOKEN > ~/.kaggle/access_token
      chmod 600 ~/.kaggle/access_token
  
  Option B — Environment variable:
    export KAGGLE_API_TOKEN=your_token
  
  Option C — Legacy kaggle.json (still supported):
    Settings → API → "Create Legacy API Key" → save to ~/.kaggle/kaggle.json
    chmod 600 ~/.kaggle/kaggle.json

Usage
-----
# First time — creates the Kaggle dataset and notebook from scratch
python kaggle_train.py \
    --dataset_dir   crops_dataset \
    --notebook      2_product-classifier.ipynb \
    --dataset_name  my-product-crops \
    --kernel_name   product-classifier-train \
    --output_dir    runs/classify

# Subsequent runs — detects existing dataset/kernel and updates them
python kaggle_train.py \
    --dataset_dir   crops_dataset \
    --notebook      2_product-classifier.ipynb \
    --dataset_name  my-product-crops \
    --kernel_name   product-classifier-train

Arguments
---------
--dataset_dir   Local crops_dataset/ folder to upload      (default: crops_dataset)
--notebook      Kaggle notebook .ipynb file                 (default: 2_product-classifier.ipynb)
--dataset_name  Kaggle dataset slug (lowercase, hyphens)   (default: product-dataset)
--kernel_name   Kaggle kernel slug  (lowercase, hyphens)   (default: product-predection)
--output_dir    Where to save downloaded weights            (default: runs/classify)
--poll_interval Seconds between status checks              (default: 30)
--timeout       Max minutes to wait for kernel             (default: 180)
--skip_upload   Skip dataset upload (use existing version) (flag)
--skip_push     Skip kernel push   (use last run)          (flag)
--gpu           Request GPU accelerator                    (default: True)
--public        Make dataset/kernel public                 (default: False)
"""

import argparse
import json
import os
import shutil
import sys
import tempfile
import time
import zipfile
from pathlib import Path


# ─────────────────────────── args ────────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser(
        description="Automate Kaggle dataset upload + notebook run + output download",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--dataset_dir",   default="crops_dataset")
    p.add_argument("--notebook",      default="2_product-classifier.ipynb")
    p.add_argument("--dataset_name",  default="product-dataset",
                   help="Kaggle dataset slug — lowercase, hyphens, no spaces")
    p.add_argument("--kernel_name",   default="product-predection",
                   help="Kaggle kernel slug — lowercase, hyphens, no spaces")
    p.add_argument("--output_dir",    default="runs/classify")
    p.add_argument("--poll_interval", type=int, default=30,
                   help="Seconds between kernel status polls")
    p.add_argument("--timeout",       type=int, default=180,
                   help="Max minutes to wait for the kernel to finish")
    p.add_argument("--skip_upload",   action="store_true",
                   help="Skip dataset upload — use the existing Kaggle dataset version")
    p.add_argument("--skip_push",     action="store_true",
                   help="Skip kernel push — just poll and download last run")
    p.add_argument("--gpu",           action="store_true", default=True,
                   help="Request GPU (T4) accelerator for the kernel")
    p.add_argument("--no_gpu",        action="store_true",
                   help="Override --gpu: run on CPU instead")
    p.add_argument("--public",        action="store_true",
                   help="Make dataset and kernel public (default: private)")
    return p.parse_args()


# ─────────────────────────── helpers ─────────────────────────────────────────

def log(msg: str, icon: str = "▸"):
    ts = time.strftime("%H:%M:%S")
    print(f"[{ts}] {icon}  {msg}", flush=True)

def fatal(msg: str):
    print(f"\n[ERROR] {msg}", file=sys.stderr)
    sys.exit(1)

# ── Credential paths for all 4 Kaggle auth methods ──────────────────────────
_KAGGLE_DIR        = Path.home() / ".kaggle"
_ACCESS_TOKEN_PATH = _KAGGLE_DIR / "access_token"   # new  (Option 3)
_KAGGLE_JSON_PATH  = _KAGGLE_DIR / "kaggle.json"    # legacy (Option 4)


def _detect_auth_method() -> str:
    """
    Return which auth method is available, in priority order:
      'env_token'   — KAGGLE_API_TOKEN env var        (Option 2, new)
      'env_legacy'  — KAGGLE_USERNAME + KAGGLE_KEY    (Option 2, legacy)
      'access_token'— ~/.kaggle/access_token file     (Option 3, new)
      'kaggle_json' — ~/.kaggle/kaggle.json           (Option 4, legacy)
      None          — nothing found
    """
    if os.getenv("KAGGLE_API_TOKEN"):
        return "env_token"
    if os.getenv("KAGGLE_USERNAME") and os.getenv("KAGGLE_KEY"):
        return "env_legacy"
    if _ACCESS_TOKEN_PATH.exists():
        #print(f"this is the PATH:  {_ACCESS_TOKEN_PATH}")
        return "access_token"
    if _KAGGLE_JSON_PATH.exists():
        return "kaggle_json"
    return None


def check_credentials():
    """
    Detect which Kaggle auth method is present and print a clear error if none.
    Supports all 4 methods documented at https://github.com/Kaggle/kaggle-cli:
      Option 2: KAGGLE_API_TOKEN env var          (new single-token style)
      Option 2: KAGGLE_USERNAME + KAGGLE_KEY      (legacy env vars)
      Option 3: ~/.kaggle/access_token            (new token file — what Kaggle now issues)
      Option 4: ~/.kaggle/kaggle.json             (legacy JSON file)
    """
    method = _detect_auth_method()
    if method is None:
        fatal(
            "Kaggle credentials not found. Choose one of these options:\n\n"
            "Option 1 — access_token file (what Kaggle now issues):\n"
            "  1. Go to https://www.kaggle.com/settings/api\n"
            "  2. Click 'Create New Token' and copy the command shown, e.g.:\n"
            "       mkdir -p ~/.kaggle && echo YOUR_TOKEN > ~/.kaggle/access_token\n"
            "       chmod 600 ~/.kaggle/access_token\n\n"
            "Option 2 — environment variable (new style):\n"
            "  export KAGGLE_API_TOKEN=your_token\n\n"
            "Option 3 — environment variables (legacy style):\n"
            "  export KAGGLE_USERNAME=your_username\n"
            "  export KAGGLE_KEY=your_api_key\n\n"
            "Option 4 — kaggle.json (legacy file, still supported):\n"
            "  Go to Kaggle → Settings → API → 'Create Legacy API Key'\n"
            "  Move kaggle.json to ~/.kaggle/kaggle.json\n"
            "  chmod 600 ~/.kaggle/kaggle.json\n"
        )

    labels = {
        "env_token":    "KAGGLE_API_TOKEN env var",
        "env_legacy":   "KAGGLE_USERNAME + KAGGLE_KEY env vars",
        "access_token": f"~/.kaggle/access_token  (new token file)",
        "kaggle_json":  f"~/.kaggle/kaggle.json   (legacy)",
    }
    log(f"Auth method: {labels[method]}", "✓")


def get_kaggle_username() -> str:
    """
    Retrieve the Kaggle username for constructing dataset/kernel refs.
    With the new access_token auth, the username is fetched from the API
    itself (whoami) instead of being embedded in the credential file.
    """
    # Legacy env var — username is explicit
    if os.getenv("KAGGLE_USERNAME"):
        return os.getenv("KAGGLE_USERNAME")

    # Legacy kaggle.json — username is in the file
    if _KAGGLE_JSON_PATH.exists():
        try:
            return json.loads(_KAGGLE_JSON_PATH.read_text())["username"]
        except Exception:
            pass

    # New access_token / KAGGLE_API_TOKEN — must ask the API
    # The kaggle package exposes this after authenticate() is called
    try:
        from kaggle.api.kaggle_api_extended import KaggleApiExtended
        _api = KaggleApiExtended()
        _api.authenticate()
        # kaggle SDK stores the authenticated username here
        if hasattr(_api, "get_config_value"):
            uname = _api.get_config_value("username")
            if uname:
                return uname
        # Fallback: call the whoami REST endpoint directly
        resp = _api.process_response(
            _api.kernel_output_with_http_info("", "")
        )
    except Exception:
        pass

    # Last resort: ask the user
    print()
    uname = input("  Could not auto-detect your Kaggle username.\n"
                  "  Enter it manually (found at kaggle.com/your-username): ").strip()
    if not uname:
        fatal("Username is required to build dataset/kernel references.")
    return uname

def slugify(name: str) -> str:
    """Ensure a slug is lowercase with hyphens."""
    return name.lower().replace("_", "-").replace(" ", "-")


# ─────────────────────────── step 1: zip + upload dataset ────────────────────

def zip_dataset(dataset_dir: Path, tmp_dir: Path) -> Path:
    """Zip the dataset folder, preserving the train/val/test structure."""
    zip_path = tmp_dir / "dataset.zip"
    log(f"Zipping {dataset_dir} …", "📦")

    total = sum(1 for p in dataset_dir.rglob("*") if p.is_file())
    done  = 0

    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for file_path in sorted(dataset_dir.rglob("*")):
            if file_path.is_file():
                arcname = file_path.relative_to(dataset_dir.parent)
                zf.write(file_path, arcname)
                done += 1
                if done % 500 == 0 or done == total:
                    pct = done / total * 100
                    print(f"  Zipped {done:,}/{total:,}  ({pct:.0f}%)", end="\r")

    size_mb = zip_path.stat().st_size / 1e6
    print()
    log(f"Zip ready: {zip_path} ({size_mb:.1f} MB)", "✓")
    return zip_path


def upload_dataset(api, username: str, dataset_name: str,
                   dataset_dir: Path, tmp_dir: Path, public: bool) -> str:
    """
    Create or update a Kaggle dataset from dataset_dir.
    Returns the full dataset reference string: username/dataset-name
    """
    slug  = slugify(dataset_name)
    ref   = f"{username}/{slug}"

    # Check if dataset already exists
    existing = None
    try:
        existing = api.dataset_status(username, slug)
    except Exception:
        pass   # 404 = doesn't exist yet

    # Write dataset-metadata.json into a staging folder
    stage_dir = tmp_dir / "dataset_stage"
    stage_dir.mkdir(exist_ok=True)

    # Copy the crops_dataset folder into staging
    dst = stage_dir / dataset_dir.name
    if dst.exists():
        shutil.rmtree(dst)
    shutil.copytree(str(dataset_dir), str(dst))

    meta = {
        "title":    dataset_name.replace("-", " ").title(),
        "id":       ref,
        "licenses": [{"name": "CC0-1.0"}],
    }
    (stage_dir / "dataset-metadata.json").write_text(json.dumps(meta, indent=2))

    if existing is None:
        log(f"Creating new Kaggle dataset: {ref}", "☁")
        api.dataset_create_new(
            folder      = str(stage_dir),
            public      = public,
            quiet       = False,
            convert_to_csv = False,
            dir_mode    = "zip",
        )
        log(f"Dataset created: https://www.kaggle.com/datasets/{ref}", "✓")
    else:
        log(f"Updating existing Kaggle dataset: {ref}", "☁")
        api.dataset_create_version(
            folder          = str(stage_dir),
            version_notes   = f"Updated via kaggle_train.py at {time.strftime('%Y-%m-%d %H:%M')}",
            quiet           = False,
            convert_to_csv  = False,
            delete_old_versions = False,
            dir_mode        = "zip",
        )
        log(f"Dataset updated: https://www.kaggle.com/datasets/{ref}", "✓")

    return ref


# ─────────────────────────── step 2: patch notebook ──────────────────────────

def patch_notebook(notebook_path: Path, dataset_ref: str,
                   tmp_dir: Path) -> Path:
    """
    Patch DATA_DIR in the notebook to point to the uploaded Kaggle dataset,
    and save a modified copy into tmp_dir.
    
    The notebook's Cell 2 contains something like:
        DATA_DIR = Path("/kaggle/input/your-dataset-name")
    We replace the path to match the actual uploaded dataset slug.
    """
    nb = json.loads(notebook_path.read_text(encoding="utf-8"))

    # Extract the dataset slug from "username/slug" → "slug"
    ds_slug = dataset_ref.split("/")[-1]

    patched = False
    for cell in nb.get("cells", []):
        if cell.get("cell_type") != "code":
            continue
        src = "".join(cell["source"])
        if "DATA_DIR" in src and "/kaggle/input/" in src:
            new_lines = []
            for line in cell["source"]:
                if "DATA_DIR" in line and "/kaggle/input/" in line:
                    # Preserve indentation, replace only the path string
                    indent  = len(line) - len(line.lstrip())
                    new_line = " " * indent + f'DATA_DIR = Path("/kaggle/input/{ds_slug}/{Path(ds_slug).stem}")\n'
                    # Try to find the inner folder name (crops_dataset by default)
                    # Use the dataset_dir name if we can detect it from context
                    new_line = " " * indent + f'DATA_DIR = Path("/kaggle/input/{ds_slug}")\n'
                    new_lines.append(new_line)
                    patched = True
                    log(f"Patched DATA_DIR → /kaggle/input/{ds_slug}", "✓")
                else:
                    new_lines.append(line)
            cell["source"] = new_lines

    if not patched:
        log("⚠ DATA_DIR line not found in notebook — using notebook as-is", "⚠")

    out_path = tmp_dir / notebook_path.name
    out_path.write_text(json.dumps(nb, indent=1), encoding="utf-8")
    return out_path


# ─────────────────────────── step 3: push kernel ─────────────────────────────

def push_kernel(api, username: str, kernel_name: str, notebook_path: Path,
                dataset_ref: str, tmp_dir: Path,
                gpu: bool, public: bool) -> str:
    """
    Create or update a Kaggle kernel (notebook) and trigger a run.
    Returns the kernel ref string: username/kernel-name
    """
    slug = slugify(kernel_name)
    ref  = f"{username}/{slug}"

    kernel_dir = tmp_dir / "kernel_push"
    kernel_dir.mkdir(exist_ok=True)

    # Copy patched notebook into kernel dir
    nb_dest = kernel_dir / notebook_path.name
    shutil.copy2(str(notebook_path), str(nb_dest))

    # kernel-metadata.json — controls GPU, dataset attachment, language
    meta = {
        "id":                  ref,
        "title":               kernel_name.replace("-", " ").title(),
        "code_file":           notebook_path.name,
        "language":            "python",
        "kernel_type":         "notebook",
        "is_private":          not public,
        "enable_gpu":          gpu,
        "enable_tpu":          False,
        "enable_internet":     True,
        "dataset_sources":     [dataset_ref],
        "competition_sources": [],
        "kernel_sources":      [],
        "model_sources":       [],
        "machine_shape":       "t4" if gpu else None,
    }
    (kernel_dir / "kernel-metadata.json").write_text(json.dumps(meta, indent=2))

    log(f"Pushing kernel to Kaggle: {ref}  (GPU={'yes' if gpu else 'no'})", "🚀")
    try:
        api.kernels_push(str(kernel_dir), timeout=None, acc="t4" if gpu else None)
    except TypeError:
        api.kernels_push(str(kernel_dir))
    log(f"Kernel pushed — run started: https://www.kaggle.com/code/{ref}", "✓")
    return ref


# ─────────────────────────── step 4: poll ────────────────────────────────────

def poll_kernel(api, kernel_ref: str, poll_interval: int, timeout_minutes: int) -> bool:
    """
    Poll the kernel until it reaches a terminal status.
    Returns True if successful, False if failed or timed out.
    """
    deadline = time.time() + timeout_minutes * 60
    attempt  = 0

    TERMINAL   = {"complete", "error", "cancelled", "cancel_acknowledged"}
    RUNNING    = {"running", "queued", "initializing", "pending"}

    log(f"Polling kernel status every {poll_interval}s "
        f"(timeout {timeout_minutes} min) …", "⏳")

    while time.time() < deadline:
        attempt += 1
        try:
            # kernels_status takes the full kernel ref "username/slug"
            status_obj = api.kernels_status(kernel_ref)
            # Extract status from the response object
            raw_status = getattr(status_obj, "status", None)
            if raw_status is not None:
                if hasattr(raw_status, 'name'):
                    status = raw_status.name.lower()
                else:
                    status = str(raw_status).lower()
            else:
                status = str(status_obj).lower()
            
            # Clean up enum prefixes
            status = status.replace("kernelworkerstatus.", "")
        except Exception as e:
            log(f"Poll error (will retry): {e}", "⚠")
            time.sleep(poll_interval)
            continue

        elapsed_min = (time.time() + timeout_minutes * 60 - deadline) / 60
        print(f"  [{attempt:>3}]  status={status:<20}  elapsed={elapsed_min:.1f} min",
              end="\r", flush=True)

        if status == "complete":
            print()
            log("Kernel completed successfully!", "✓")
            return True

        if status == "error":
            print()
            log("Kernel run failed — check logs at "
                f"https://www.kaggle.com/code/{kernel_ref}", "✗")
            return False

        if "cancel" in status:
            print()
            log("Kernel was cancelled.", "✗")
            return False

        if status not in RUNNING:
            print()
            log(f"Unknown status: {status!r} — continuing to poll", "⚠")

        time.sleep(poll_interval)

    print()
    log(f"Timeout after {timeout_minutes} minutes. "
        f"Check manually: https://www.kaggle.com/code/{kernel_ref}", "✗")
    return False


# ─────────────────────────── step 5: download outputs ────────────────────────

def download_outputs(api, kernel_ref: str, output_dir: Path) -> list[Path]:
    """
    Download all kernel output files and return paths to the ones we care about.
    Targets: best.pt, last.pt, class_names.json, *.png
    """
    username, slug = kernel_ref.split("/")
    output_dir.mkdir(parents=True, exist_ok=True)

    log(f"Downloading kernel outputs → {output_dir}", "⬇")

    with tempfile.TemporaryDirectory() as tmp:
        try:
            api.kernels_output(username, slug, path=tmp, unzip=True)
        except Exception as e:
            fatal(f"Could not download kernel outputs: {e}\n"
                  f"Try manually at https://www.kaggle.com/code/{kernel_ref}/output")

        tmp_path = Path(tmp)
        downloaded = []

        # Walk everything the kernel produced
        for f in sorted(tmp_path.rglob("*")):
            if not f.is_file():
                continue

            # We want weights, class map, and training plots
            keep = f.suffix in {".pt", ".json", ".png", ".csv"}
            if not keep:
                continue

            dst = output_dir / f.relative_to(tmp_path)
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(str(f), str(dst))
            downloaded.append(dst)
            log(f"  {dst.relative_to(output_dir)}  ({f.stat().st_size/1e6:.2f} MB)", "  ")

    if not downloaded:
        log("No output files found. The kernel may not have saved to /kaggle/working/", "⚠")
    else:
        log(f"{len(downloaded)} file(s) downloaded to {output_dir.resolve()}", "✓")

    return downloaded


# ─────────────────────────── main ────────────────────────────────────────────

def main():
    args = parse_args()

    gpu = args.gpu and not args.no_gpu

    check_credentials()

    # Import here so missing kaggle package gives a clean error
    try:
        from kaggle.api.kaggle_api_extended import KaggleApi
        api = KaggleApi()
        api.authenticate()
    except ImportError:
        fatal("kaggle package not installed.\n  Run: pip install kaggle")
    except Exception as e:
        fatal(
            f"Kaggle authentication failed: {e}\n\n"
            "Troubleshooting:\n"
            "  • New token file:  ~/.kaggle/access_token  (chmod 600)\n"
            "  • Env var:         export KAGGLE_API_TOKEN=your_token\n"
            "  • Legacy JSON:     ~/.kaggle/kaggle.json   (chmod 600)\n"
            "  • Or run:          kaggle auth login\n"
        )

    username = get_kaggle_username()
    log(f"Authenticated as: {username}", "✓")

    dataset_dir   = Path(args.dataset_dir)
    notebook_path = Path(args.notebook)
    output_dir    = Path(args.output_dir)

    if not dataset_dir.exists():
        fatal(f"Dataset directory not found: {dataset_dir}")
    if not notebook_path.exists():
        fatal(f"Notebook not found: {notebook_path}")

    print(f"""
╔══════════════════════════════════════════════════════╗
║           Kaggle Training Automation                ║
╠══════════════════════════════════════════════════════╣
  Kaggle user   : {username}
  Dataset       : {args.dataset_name}  ({dataset_dir})
  Kernel        : {args.kernel_name}
  Notebook      : {notebook_path}
  GPU           : {'yes (T4)' if gpu else 'no (CPU)'}
  Output        : {output_dir}
  Poll interval : {args.poll_interval}s
  Timeout       : {args.timeout} min
╚══════════════════════════════════════════════════════╝
""")

    with tempfile.TemporaryDirectory() as tmp_str:
        tmp = Path(tmp_str)

        # ── Step 1: upload dataset ────────────────────────────────────────────
        if not args.skip_upload:
            dataset_ref = upload_dataset(
                api, username, args.dataset_name,
                dataset_dir, tmp, args.public
            )
            # Give Kaggle a moment to process the upload before pushing the kernel
            log("Waiting 10 s for dataset to register on Kaggle…", "⏳")
            time.sleep(10)
        else:
            dataset_ref = f"{username}/{slugify(args.dataset_name)}"
            log(f"Skipping upload — using existing dataset: {dataset_ref}", "↷")

        # ── Step 2: patch notebook ────────────────────────────────────────────
        patched_nb = patch_notebook(notebook_path, dataset_ref, tmp)

        # ── Step 3: push kernel ───────────────────────────────────────────────
        kernel_ref = f"{username}/{slugify(args.kernel_name)}"
        if not args.skip_push:
            kernel_ref = push_kernel(
                api, username, args.kernel_name,
                patched_nb, dataset_ref, tmp,
                gpu, args.public,
            )
        else:
            log(f"Skipping push — polling existing kernel: {kernel_ref}", "↷")

        # ── Step 4: poll ──────────────────────────────────────────────────────
        success = poll_kernel(api, kernel_ref, args.poll_interval, args.timeout)

        if not success:
            log(
                f"Training did not complete successfully.\n"
                f"  View logs : https://www.kaggle.com/code/{kernel_ref}\n"
                f"  Re-run download only:\n"
                f"    python kaggle_train.py --skip_upload --skip_push "
                f"--kernel_name {args.kernel_name}",
                "✗"
            )
            sys.exit(1)

        # ── Step 5: download ──────────────────────────────────────────────────
        downloaded = download_outputs(api, kernel_ref, output_dir)

    # ── Final summary ─────────────────────────────────────────────────────────
    best_pt    = next((f for f in downloaded if f.name == "best.pt"),       None)
    cls_json   = next((f for f in downloaded if f.name == "class_names.json"), None)
    curves_png = next((f for f in downloaded if "curves" in f.name),        None)

    print()
    print("═" * 55)
    print("  Training pipeline complete!")
    if best_pt:
        print(f"  Weights    : {best_pt}")
    if cls_json:
        print(f"  Classes    : {cls_json}")
    if curves_png:
        print(f"  Curves     : {curves_png}")
    print()
    print("  Next steps:")
    print(f"    python 3_inference.py \\")
    print(f"      --classifier_weights {output_dir}/best.pt \\")
    print(f"      --detector_weights detector/best.pt \\")
    print(f"      --source shelf_images/")
    print("═" * 55)


if __name__ == "__main__":
    main()