#!/usr/bin/env python3
"""Push a PEFT/LoRA adapter checkpoint to the Hugging Face Hub.

Usage examples
--------------
    # Validate everything without pushing:
    python scripts/push_to_hub.py --dry-run

    # Push with defaults (reads HF_TOKEN from environment):
    python scripts/push_to_hub.py

    # Explicit options:
    python scripts/push_to_hub.py \\
        --checkpoint checkpoints/available-lora-0.5b-full/final \\
        --repo-id neuralbroker/blitzkode-lora-0.5b \\
        --commit-message "Add trained adapter v2.1"

    # Private repo push with explicit token:
    python scripts/push_to_hub.py --private --token hf_...
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import textwrap
from pathlib import Path

# ---------------------------------------------------------------------------
# Constants / defaults
# ---------------------------------------------------------------------------
REPO_ROOT = Path(__file__).resolve().parents[1]

DEFAULT_CHECKPOINT = REPO_ROOT / "checkpoints" / "available-lora-0.5b-full" / "final"
DEFAULT_REPO_ID = "neuralbroker/blitzkode-lora-0.5b"
DEFAULT_REPO_TYPE = "model"
DEFAULT_COMMIT_MSG = "Upload BlitzKode LoRA adapter"

# Files that must be present for a valid PEFT adapter
REQUIRED_FILES = ["adapter_config.json", "adapter_model.safetensors"]


# ---------------------------------------------------------------------------
# CLI argument parsing
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=DEFAULT_CHECKPOINT,
        metavar="PATH",
        help=(
            f"Path to the adapter checkpoint directory. "
            f"(default: {DEFAULT_CHECKPOINT})"
        ),
    )
    parser.add_argument(
        "--repo-id",
        default=DEFAULT_REPO_ID,
        metavar="OWNER/REPO",
        help=f"HuggingFace repo to push to. (default: {DEFAULT_REPO_ID})",
    )
    parser.add_argument(
        "--repo-type",
        default=DEFAULT_REPO_TYPE,
        choices=("model", "dataset", "space"),
        help=f"Repository type. (default: {DEFAULT_REPO_TYPE})",
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Create the repository as private. (default: public)",
    )
    parser.add_argument(
        "--token",
        default=None,
        metavar="HF_TOKEN",
        help=(
            "HuggingFace API write token. "
            "Falls back to the HF_TOKEN environment variable if not set."
        ),
    )
    parser.add_argument(
        "--create-repo",
        action=argparse.BooleanOptionalAction,
        default=True,
        help=(
            "Create the HuggingFace repo if it does not exist. "
            "Use --no-create-repo to skip. (default: True)"
        ),
    )
    parser.add_argument(
        "--commit-message",
        default=DEFAULT_COMMIT_MSG,
        metavar="MSG",
        help=f"Commit message for the Hub upload. (default: '{DEFAULT_COMMIT_MSG}')",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help=(
            "Validate the checkpoint and configuration but do NOT push anything "
            "to Hugging Face Hub. Useful for CI or pre-flight checks."
        ),
    )
    return parser.parse_args()


# ---------------------------------------------------------------------------
# Dependency check
# ---------------------------------------------------------------------------
def check_huggingface_hub() -> None:
    """Abort with a helpful message if huggingface_hub is not installed."""
    try:
        import huggingface_hub  # noqa: F401  # type: ignore[import]
    except ImportError:
        print(
            "\n[ERROR] The `huggingface_hub` package is not installed.\n"
            "Install it with one of the following commands:\n\n"
            "    pip install huggingface_hub\n"
            "    pip install -r requirements-training.txt\n",
            file=sys.stderr,
        )
        sys.exit(1)


# ---------------------------------------------------------------------------
# Checkpoint validation
# ---------------------------------------------------------------------------
def validate_checkpoint(checkpoint: Path) -> dict:
    """Ensure the checkpoint directory is valid.

    Checks that the directory exists and contains every file listed in
    REQUIRED_FILES. Returns the parsed ``adapter_config.json`` dict.
    """
    if not checkpoint.exists():
        print(
            f"\n[ERROR] Checkpoint directory not found: {checkpoint}\n"
            "Run training first, e.g.:\n"
            "    python scripts/train_available.py\n",
            file=sys.stderr,
        )
        sys.exit(1)

    if not checkpoint.is_dir():
        print(
            f"\n[ERROR] Checkpoint path is not a directory: {checkpoint}\n",
            file=sys.stderr,
        )
        sys.exit(1)

    missing = [f for f in REQUIRED_FILES if not (checkpoint / f).exists()]
    if missing:
        print(
            f"\n[ERROR] Missing required files in {checkpoint}:\n"
            + "\n".join(f"  - {f}" for f in missing)
            + "\n\nIs this a valid PEFT adapter checkpoint?\n",
            file=sys.stderr,
        )
        sys.exit(1)

    config_path = checkpoint / "adapter_config.json"
    try:
        adapter_config: dict = json.loads(config_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        print(
            f"\n[ERROR] adapter_config.json is not valid JSON: {exc}\n",
            file=sys.stderr,
        )
        sys.exit(1)
    except OSError as exc:
        print(
            f"\n[ERROR] Could not read adapter_config.json: {exc}\n",
            file=sys.stderr,
        )
        sys.exit(1)

    return adapter_config


# ---------------------------------------------------------------------------
# Token resolution
# ---------------------------------------------------------------------------
def resolve_token(args_token: str | None, *, dry_run: bool = False) -> str:
    """Return the HF token, or abort with instructions if none is found."""
    token = args_token or os.environ.get("HF_TOKEN", "")
    if token:
        return token

    if dry_run:
        # A token is not needed for dry runs, return a placeholder.
        return "__dry_run_placeholder__"

    print(
        "\n[ERROR] No HuggingFace API token found.\n"
        "Provide a write token using one of these methods:\n\n"
        "  1. CLI flag:\n"
        "       python scripts/push_to_hub.py --token hf_YOUR_TOKEN\n\n"
        "  2. Environment variable (recommended):\n"
        "       Windows CMD : set HF_TOKEN=hf_YOUR_TOKEN\n"
        "       PowerShell  : $env:HF_TOKEN = 'hf_YOUR_TOKEN'\n"
        "       Linux/macOS : export HF_TOKEN=hf_YOUR_TOKEN\n\n"
        "  3. HuggingFace CLI login (persists across sessions):\n"
        "       pip install huggingface_hub\n"
        "       huggingface-cli login\n\n"
        "Generate a token at: https://huggingface.co/settings/tokens\n"
        "Make sure the token has **write** access to the target repo.\n",
        file=sys.stderr,
    )
    sys.exit(1)


# ---------------------------------------------------------------------------
# Model card / README generation
# ---------------------------------------------------------------------------
def build_model_card(adapter_config: dict, repo_id: str) -> str:
    """Generate the HuggingFace-compatible README.md content for the adapter repo."""
    base_model = adapter_config.get(
        "base_model_name_or_path", "Qwen/Qwen2.5-0.5B-Instruct"
    )
    lora_r = adapter_config.get("r", 16)
    lora_alpha = adapter_config.get("lora_alpha", 32)
    lora_dropout = adapter_config.get("lora_dropout", 0.05)
    target_modules: list = adapter_config.get("target_modules", [])
    modules_str = (
        ", ".join(f"`{m}`" for m in target_modules)
        if target_modules
        else "`q_proj`, `k_proj`, `v_proj`, `o_proj`"
    )

    # YAML frontmatter -------------------------------------------------------
    frontmatter = textwrap.dedent(f"""\
        ---
        language:
          - en
        license: mit
        library_name: peft
        tags:
          - code-generation
          - lora
          - qwen2.5
          - blitzkode
          - coding-assistant
          - fine-tuned
          - peft
        base_model: {base_model}
        pipeline_tag: text-generation
        ---
        """)

    # README body ------------------------------------------------------------
    body = textwrap.dedent(f"""\
        # BlitzKode LoRA Adapter (0.5B)

        **BlitzKode** is a local AI coding assistant fine-tuned from
        **[{base_model}](https://huggingface.co/{base_model})** using LoRA
        (Low-Rank Adaptation). This repository contains the PEFT adapter — the
        research-friendly version that can be hot-loaded on top of the base model.

        > **Creator:** [Sajad (neuralbroker)](https://github.com/neuralbroker)
        > **GitHub:** <https://github.com/neuralbroker/blitzkode>
        > **Production GGUF:** [`neuralbroker/blitzkode`](https://huggingface.co/neuralbroker/blitzkode)

        ---

        ## Model Details

        | Property | Value |
        |---|---|
        | **Adapter version** | 2.1 |
        | **Base model** | `{base_model}` |
        | **LoRA rank (r)** | {lora_r} |
        | **LoRA alpha** | {lora_alpha} |
        | **LoRA dropout** | {lora_dropout} |
        | **Target modules** | {modules_str} |
        | **Training steps** | 50 |
        | **Final loss** | ~0.48 |
        | **Library** | PEFT |
        | **License** | MIT |

        ---

        ## Training Pipeline

        This adapter was produced by a **4-stage fine-tuning pipeline** applied
        to the Qwen2.5 family:

        | Stage | Method | Purpose |
        |---|---|---|
        | 1 | SFT | Supervised fine-tuning on 71 curated algorithmic coding problems |
        | 2 | Reward-SFT | Continued SFT with heuristic reward signals for code correctness and formatting |
        | 3 | DPO | Direct Preference Optimization on handcrafted chosen/rejected pairs |
        | 4 | LoRA SFT (this adapter) | Final LoRA fine-tune (r={lora_r}) on 99 samples; base model Qwen2.5-0.5B |

        ### Training Dataset (199 total samples)

        | Subset | Count | Source | License |
        |---|---|---|---|
        | Curated algorithmic problems | 71 | Custom (local) — arrays, strings, trees, DP, graphs | MIT |
        | MetaMathQA samples | 100 | [`meta-math/MetaMathQA`](https://huggingface.co/datasets/meta-math/MetaMathQA) | CC BY 4.0 |
        | Python/JavaScript patterns | 28 | Custom (local) — decorators, context managers, data classes | MIT |
        | **Total** | **199** | | |

        ---

        ## Usage

        ### Load with PEFT

        ```python
        from peft import PeftModel
        from transformers import AutoModelForCausalLM, AutoTokenizer

        base_model_id = "{base_model}"
        adapter_repo  = "{repo_id}"

        tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            torch_dtype="auto",
            device_map="auto",
            trust_remote_code=True,
        )
        model = PeftModel.from_pretrained(model, adapter_repo)
        model.eval()
        ```

        ### Generate code

        ```python
        prompt = (
            "<|im_start|>system\\n"
            "You are BlitzKode, a precise AI coding assistant created by Sajad.\\n"
            "<|im_end|>\\n"
            "<|im_start|>user\\n"
            "Write a Python function for binary search with full edge-case handling.\\n"
            "<|im_end|>\\n"
            "<|im_start|>assistant\\n"
        )

        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,
            temperature=0.7,
            do_sample=True,
            repetition_penalty=1.1,
        )
        print(tokenizer.decode(outputs[0], skip_special_tokens=True))
        ```

        ### Merge adapter into base model (for export)

        ```python
        merged = model.merge_and_unload()
        merged.save_pretrained("blitzkode-0.5b-merged")
        tokenizer.save_pretrained("blitzkode-0.5b-merged")
        ```

        ---

        ## Prompt Format

        BlitzKode uses the **ChatML** template standard for Qwen models:

        ```
        <|im_start|>system
        You are BlitzKode, a precise AI coding assistant created by Sajad.<|im_end|>
        <|im_start|>user
        {{your question}}<|im_end|>
        <|im_start|>assistant
        ```

        ---

        ## Limitations

        - **Text-only** — no image/multimodal support.
        - **0.5B parameters** — smaller and faster than the 1.5B GGUF variant; may be
          less accurate on complex algorithmic tasks.
        - **2048-token context** — not suitable for long repository-level analysis.
        - **Review all outputs** — generated code must be tested before use in production.
        - **Not security-audited** — do not use for cryptographic or safety-critical code
          without thorough expert review.
        - **Math reasoning** — MetaMathQA training improves basic reasoning but does not
          substitute a dedicated math model.

        ---

        ## Relation to the Production Model

        | Variant | Repo | Size | Runtime | Use case |
        |---|---|---|---|---|
        | GGUF (1.5B, F16) | [`neuralbroker/blitzkode`](https://huggingface.co/neuralbroker/blitzkode) | ~3 GB | llama.cpp / llama-cpp-python | Production; CPU/GPU, no Python ML stack needed |
        | LoRA adapter (0.5B) | `{repo_id}` (this repo) | ~100 MB | PEFT + Transformers | Research; merging, further fine-tuning, quantization |

        ---

        ## License

        **MIT** — see [LICENSE](https://github.com/neuralbroker/blitzkode/blob/main/LICENSE).

        You must also comply with the upstream
        [{base_model}](https://huggingface.co/{base_model}) license
        when redistributing any derived weights.

        ---

        ## Citation

        ```bibtex
        @software{{blitzkode2025,
          author  = {{Sajad}},
          title   = {{BlitzKode: A Local AI Coding Assistant}},
          year    = {{2025}},
          url     = {{https://github.com/neuralbroker/blitzkode}}
        }}
        ```
        """)

    return frontmatter + "\n" + body


# ---------------------------------------------------------------------------
# Main push routine
# ---------------------------------------------------------------------------
def push(args: argparse.Namespace) -> None:  # noqa: C901
    check_huggingface_hub()

    # Import here so the check above can give a clean error first.
    from huggingface_hub import HfApi  # type: ignore[import]
    from huggingface_hub.utils import HfHubHTTPError  # type: ignore[import]

    sep = "=" * 70
    print(sep)
    print("BlitzKode — Push LoRA Adapter to Hugging Face Hub")
    if args.dry_run:
        print("(DRY RUN — nothing will be pushed)")
    print(sep)

    # ------------------------------------------------------------------
    # Step 1: Validate checkpoint
    # ------------------------------------------------------------------
    print(f"\n[1/5] Validating checkpoint directory …")
    print(f"      Path: {args.checkpoint}")
    adapter_config = validate_checkpoint(args.checkpoint)

    base_model = adapter_config.get("base_model_name_or_path", "unknown")
    lora_r = adapter_config.get("r", "?")
    lora_alpha = adapter_config.get("lora_alpha", "?")
    target_modules = adapter_config.get("target_modules", [])
    files_found = sorted(p.name for p in args.checkpoint.iterdir() if p.is_file())

    print(f"      base_model     : {base_model}")
    print(f"      lora r / alpha : {lora_r} / {lora_alpha}")
    print(f"      target_modules : {target_modules}")
    print(f"      files          : {files_found}")
    print("      [OK] Checkpoint is valid.")

    # ------------------------------------------------------------------
    # Step 2: Resolve token
    # ------------------------------------------------------------------
    print("\n[2/5] Resolving HuggingFace token …")
    token = resolve_token(args.token, dry_run=args.dry_run)
    if args.dry_run:
        print("      [OK] Token check skipped (dry run).")
    else:
        masked = token[:8] + "..." if len(token) > 8 else "***"
        print(f"      [OK] Token resolved (starts with: {masked})")

    # ------------------------------------------------------------------
    # Dry-run exit
    # ------------------------------------------------------------------
    if args.dry_run:
        print()
        print(sep)
        print("DRY RUN COMPLETE — all validations passed, nothing was pushed.")
        print(f"  Checkpoint  : {args.checkpoint}")
        print(f"  Target repo : https://huggingface.co/{args.repo_id}")
        print(f"  Repo type   : {args.repo_type}")
        print(f"  Private     : {args.private}")
        print(f"  Files ready : {files_found}")
        print(sep)
        return

    api = HfApi(token=token)

    # ------------------------------------------------------------------
    # Step 3: Create repo (if requested)
    # ------------------------------------------------------------------
    if args.create_repo:
        print(f"\n[3/5] Creating / verifying repo: {args.repo_id} …")
        try:
            repo_url = api.create_repo(
                repo_id=args.repo_id,
                repo_type=args.repo_type,
                private=args.private,
                exist_ok=True,   # silently succeed if repo already exists
            )
            print(f"      [OK] Repo ready: {repo_url}")
        except HfHubHTTPError as exc:
            print(
                f"\n[ERROR] Failed to create / access repo '{args.repo_id}':\n"
                f"  {exc}\n"
                "Check that your token has write access and the repo name is correct.\n",
                file=sys.stderr,
            )
            sys.exit(1)
    else:
        print("\n[3/5] Skipping repo creation (--no-create-repo).")

    # ------------------------------------------------------------------
    # Step 4: Upload checkpoint folder
    # ------------------------------------------------------------------
    print(f"\n[4/5] Uploading checkpoint folder → {args.repo_id} …")
    print(f"      Commit message: \"{args.commit_message}\"")
    try:
        commit_info = api.upload_folder(
            folder_path=str(args.checkpoint),
            repo_id=args.repo_id,
            repo_type=args.repo_type,
            commit_message=args.commit_message,
        )
        commit_ref = getattr(commit_info, "oid", None) or str(commit_info)
        print(f"      [OK] Folder uploaded. Commit: {commit_ref}")
    except HfHubHTTPError as exc:
        print(
            f"\n[ERROR] Folder upload failed:\n  {exc}\n",
            file=sys.stderr,
        )
        sys.exit(1)

    # ------------------------------------------------------------------
    # Step 5: Upload model card README.md
    # ------------------------------------------------------------------
    print("\n[5/5] Uploading model card (README.md) …")
    readme_content = build_model_card(adapter_config, args.repo_id)
    try:
        api.upload_file(
            path_or_fileobj=readme_content.encode("utf-8"),
            path_in_repo="README.md",
            repo_id=args.repo_id,
            repo_type=args.repo_type,
            commit_message="Update model card README.md",
        )
        print("      [OK] README.md uploaded.")
    except HfHubHTTPError as exc:
        # Non-fatal: the adapter files are already uploaded.
        print(
            f"\n[WARN] Could not upload README.md (adapter files were uploaded OK):\n"
            f"  {exc}\n"
            "You can upload the model card manually from the Hub web interface.\n",
            file=sys.stderr,
        )

    # ------------------------------------------------------------------
    # Summary
    # ------------------------------------------------------------------
    repo_url = f"https://huggingface.co/{args.repo_id}"
    print()
    print(sep)
    print("PUSH COMPLETE")
    print(f"  Repo URL      : {repo_url}")
    print(f"  Checkpoint    : {args.checkpoint}")
    print(f"  Files pushed  : {files_found}")
    print(f"  Base model    : {base_model}")
    print(f"  LoRA r/alpha  : {lora_r}/{lora_alpha}")
    print(f"  Commit msg    : {args.commit_message}")
    print(sep)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def main() -> None:
    args = parse_args()
    push(args)


if __name__ == "__main__":
    main()