"""
HuggingFace Hub upload functionality for BADAS models.

Handles checkpoint slimming (strips optimizer states, cleans keys) and
upload to HuggingFace Hub with an auto-generated model card.

Example:
    >>> from badas.export.huggingface import upload_to_huggingface
    >>>
    >>> upload_to_huggingface(
    ...     checkpoint_path="training/logs/.../model.ckpt",
    ...     hf_repo="nexar-ai/badas-1.5-flash-lite",
    ...     backbone_name="nexar-ai/vjepa2-vits-fpc16-256-nexar",
    ...     model_type="vjepa2",
    ...     display_name="BADAS-1.5-Flash-Lite",
    ...     params="22M",
    ...     arch="ViT-S (SSL)",
    ...     description="Smallest distilled model, 5.9ms on Jetson Thor.",
    ... )
"""

import tempfile
from collections import OrderedDict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import torch


# ── Slim checkpoint ────────────────────────────────────────────────────────────

def slim_checkpoint(
    checkpoint_path: str,
    backbone_name: str,
    model_type: str,
) -> dict:
    """
    Load a Lightning checkpoint and return a slim inference-only dict.

    - Strips optimizer_states, lr_schedulers, loops, callbacks, etc.
    - Drops feature_projector keys (distillation-only, not needed for inference)
    - Pre-cleans state_dict keys: removes model./student. prefixes
    - Patches hyper_parameters.model.model_name to backbone_name
    - Adds model_type key

    Args:
        checkpoint_path: Path to Lightning .ckpt file.
        backbone_name: HuggingFace backbone model ID to store in the checkpoint.
        model_type: Model type string (e.g. "vjepa2").

    Returns:
        Slim checkpoint dict ready for torch.save().
    """
    ckpt = torch.load(checkpoint_path, map_location="cpu")

    # Patch backbone name
    hp = ckpt["hyper_parameters"]
    hp["model"]["model_name"] = backbone_name

    # Clean state dict
    raw_sd = ckpt["state_dict"]
    clean_sd = OrderedDict()
    for key, value in raw_sd.items():
        if key.startswith("feature_projector"):
            continue
        if key.startswith("model."):
            key = key[len("model."):]
        elif key.startswith("student."):
            key = key[len("student."):]
        clean_sd[key] = value

    return {
        "state_dict":       clean_sd,
        "hyper_parameters": hp,
        "model_type":       model_type,
    }


# ── Model card ─────────────────────────────────────────────────────────────────

@dataclass
class ModelCardConfig:
    """Configuration for generating a BADAS model card."""
    hf_repo:      str
    display_name: str
    backbone_name: str
    params:       str
    arch:         str
    description:  str
    frame_count:       int   = 16
    img_size:          int   = 256
    fps:               int   = 8
    family_description: str  = "the [BADAS-1.5 model family](https://huggingface.co/nexar-ai)"
    family_repos: list  = field(default_factory=lambda: [
        ("badas-1.5",           "300M", "Cloud / high-accuracy"),
        ("badas-1.5-flash",     "86M",  "Real-time GPU"),
        ("badas-1.5-flash-lite","22M",  "Edge (Jetson)"),
    ])


def make_model_card(cfg: ModelCardConfig) -> str:
    """Generate a README.md model card string for a BADAS model."""
    family_rows = "\n".join(
        f"| [nexar-ai/{repo}](https://huggingface.co/nexar-ai/{repo}) | {params} | {use} |"
        for repo, params, use in cfg.family_repos
    )

    return f"""\
---
license: other
tags:
  - video-classification
  - collision-anticipation
  - dashcam
  - autonomous-driving
  - vjepa2
  - nexar
pipeline_tag: video-classification
---

# {cfg.display_name}

**{cfg.description}**

Part of {cfg.family_description} — Nexar's dashcam-based
collision anticipation system built on V-JEPA2 backbones.

| | Value |
|---|---|
| Architecture | {cfg.arch} |
| Parameters | {cfg.params} |
| Input | {cfg.frame_count} frames × {cfg.img_size}×{cfg.img_size}, {cfg.fps} fps |
| Output | Risk probability (0–1) |
| Backbone | [{cfg.backbone_name}](https://huggingface.co/{cfg.backbone_name}) |

## Model Family

| Model | Params | Use case |
|---|---|---|
{family_rows}

## Usage

```python
from badas.inference import BADAS

predictor = BADAS.from_pretrained("{cfg.hf_repo}")

# Predict on a video file
results = predictor.predict_video("dashcam.mp4")

# Streaming (frame by frame)
for pred in predictor.predict_stream("dashcam.mp4"):
    print(f"[{{pred['timestamp']:.2f}}s] risk={{pred['probability']:.3f}}")
```

Or download the checkpoint manually:

```python
from huggingface_hub import hf_hub_download
from badas.inference import BADAS

ckpt = hf_hub_download("{cfg.hf_repo}", "model.ckpt")
predictor = BADAS(ckpt)
```

## Export to TensorRT (Jetson)

```python
from huggingface_hub import hf_hub_download
from badas.export import export_to_onnx, convert_to_tensorrt

ckpt = hf_hub_download("{cfg.hf_repo}", "model.ckpt")
export_to_onnx(ckpt, "model.onnx")
convert_to_tensorrt("model.onnx", "model.trt")
```

## Citation

```bibtex
@techreport{{nexar2025badas15,
  title       = {{BADAS-1.5: Scalable Collision Anticipation with Distilled Vision Transformers}},
  author      = {{Nexar AI}},
  year        = {{2025}},
  institution = {{Nexar}},
}}
```
"""


# ── Upload ─────────────────────────────────────────────────────────────────────

@dataclass
class HFUploadResult:
    """Result of a HuggingFace upload operation."""
    success:      bool
    hf_repo:      str
    ckpt_size_mb: Optional[float] = None
    url:          Optional[str]   = None
    error:        Optional[str]   = None


def upload_to_huggingface(
    checkpoint_path: str,
    hf_repo: str,
    backbone_name: str,
    model_type: str,
    display_name: str,
    params: str,
    arch: str,
    description: str,
    private: bool = True,
    card_config: Optional[ModelCardConfig] = None,
    dry_run: bool = False,
    verbose: bool = True,
) -> HFUploadResult:
    """
    Slim a Lightning checkpoint and upload it to HuggingFace Hub.

    Args:
        checkpoint_path: Path to the source .ckpt file.
        hf_repo: HuggingFace repo ID, e.g. "nexar-ai/badas-1.5-flash".
        backbone_name: HuggingFace backbone model ID to embed in the checkpoint.
        model_type: Model type string, e.g. "vjepa2".
        display_name: Human-readable model name for the model card.
        params: Parameter count string, e.g. "86M".
        arch: Architecture string, e.g. "ViT-B (SSL)".
        description: One-line model description for the model card.
        private: Create a private repo (default: True).
        card_config: Optional ModelCardConfig override for fine-grained control.
        dry_run: If True, prepare but do not upload anything.
        verbose: Print progress messages.

    Returns:
        HFUploadResult with success status and details.

    Example:
        >>> result = upload_to_huggingface(
        ...     checkpoint_path="model.ckpt",
        ...     hf_repo="nexar-ai/badas-1.5-flash-lite",
        ...     backbone_name="nexar-ai/vjepa2-vits-fpc16-256-nexar",
        ...     model_type="vjepa2",
        ...     display_name="BADAS-1.5-Flash-Lite",
        ...     params="22M",
        ...     arch="ViT-S (SSL)",
        ...     description="Edge model, 5.9ms on Jetson Thor.",
        ... )
        >>> print(result.url)
    """
    def log(msg):
        if verbose:
            print(msg, flush=True)

    try:
        # Build model card config
        if card_config is None:
            card_config = ModelCardConfig(
                hf_repo=hf_repo,
                display_name=display_name,
                backbone_name=backbone_name,
                params=params,
                arch=arch,
                description=description,
            )

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

            # Slim checkpoint
            log(f"  Slimming checkpoint ...")
            slim = slim_checkpoint(checkpoint_path, backbone_name, model_type)
            ckpt_out = tmp / "model.ckpt"
            torch.save(slim, ckpt_out)
            size_mb = ckpt_out.stat().st_size / 1024 ** 2
            log(f"  Size: {size_mb:.0f} MB")

            # Model card
            readme = tmp / "README.md"
            readme.write_text(make_model_card(card_config))

            if dry_run:
                log(f"  [DRY RUN] Would create {hf_repo} (private={private})")
                log(f"  [DRY RUN] Would upload model.ckpt ({size_mb:.0f} MB) + README.md")
                return HFUploadResult(success=True, hf_repo=hf_repo, ckpt_size_mb=size_mb)

            from huggingface_hub import HfApi
            api = HfApi()

            log(f"  Creating repo {hf_repo} ...")
            api.create_repo(hf_repo, private=private, exist_ok=True)

            log(f"  Uploading model.ckpt ({size_mb:.0f} MB) ...")
            api.upload_file(
                path_or_fileobj=str(ckpt_out),
                path_in_repo="model.ckpt",
                repo_id=hf_repo,
                commit_message=f"Add {display_name} inference checkpoint",
            )

            log(f"  Uploading README.md ...")
            api.upload_file(
                path_or_fileobj=str(readme),
                path_in_repo="README.md",
                repo_id=hf_repo,
                commit_message="Add model card",
            )

            url = f"https://huggingface.co/{hf_repo}"
            log(f"  Done → {url}")
            return HFUploadResult(success=True, hf_repo=hf_repo, ckpt_size_mb=size_mb, url=url)

    except Exception as e:
        return HFUploadResult(success=False, hf_repo=hf_repo, error=str(e))


__all__ = [
    "slim_checkpoint",
    "ModelCardConfig",
    "make_model_card",
    "HFUploadResult",
    "upload_to_huggingface",
]
