"""
Lightning AI'dan HuggingFace'e V5 verisi + checkpoint yukle.

3 ayri repo kullanir (clean separation):
  - musabc/nanogpt-tr-v5-data   (dataset, ~30GB binaries + tokenizer)
  - musabc/nanogpt-tr-v5-ckpts  (model, latest_ckpt.pt + best_ckpt.pt)
  - musabc/nanogpt-tr-v5-code   (model, scripts — Thunder'da clone edilir)

Kullanim:
    huggingface-cli login                  # bir kere
    python hf_push_v5.py --all             # her seyi yukle
    python hf_push_v5.py --data            # sadece binaries + tokenizer
    python hf_push_v5.py --ckpt            # sadece checkpoints
    python hf_push_v5.py --code            # sadece scriptler
    python hf_push_v5.py --user musabc     # user/org override

NOT: ilk yuklemede 30+ GB upload, internet hizina gore 30-60 dk.
     huggingface_hub multipart upload otomatik kullanir.
"""

import argparse
import os
import sys
from pathlib import Path

try:
    from huggingface_hub import HfApi, create_repo, upload_file, upload_folder
except ImportError:
    print("! huggingface_hub yok. Yukle: pip install -U huggingface_hub")
    sys.exit(1)

REPO_BASE = "nanogpt-tr-v5"
DEFAULT_USER = "musabc"

ROOT = Path(__file__).parent
DATA_DIR = ROOT / "data"
RUN_DIR = ROOT / "runs" / "tr-200m-v5"

# Veri repo'suna gidecekler (dataset)
DATA_FILES = [
    DATA_DIR / "v5_stage1.bin",
    DATA_DIR / "v5_stage2.bin",
    DATA_DIR / "v5_stage3.bin",
    DATA_DIR / "v5_val.bin",
    DATA_DIR / "v5_val_stage1.bin",  # opsiyonel — yoksa atlanir
    DATA_DIR / "v5_val_stage2.bin",
    DATA_DIR / "v5_val_stage3.bin",
    DATA_DIR / "tokenizer-tr-v5.json",
]

# Checkpoint repo'suna gidecekler (model)
CKPT_FILES = [
    RUN_DIR / "latest_ckpt.pt",
    RUN_DIR / "best_ckpt.pt",
    RUN_DIR / "train.log",
]

# Kod repo'suna gidecekler (model — kod da model repo'sunda ok)
CODE_FILES = [
    "model_v5.py",
    "muon.py",
    "05_train_v5.py",
    "06_sample.py",
    "04_tokenize.py",
    "04b_make_val.py",
    "hf_push_v5.py",
    "hf_pull_v5.py",      # asagida olusturulacak
]


def fmt_size(b):
    for u in ["B", "KB", "MB", "GB"]:
        if b < 1024:
            return f"{b:.1f} {u}"
        b /= 1024
    return f"{b:.1f} TB"


def ensure_repo(api: HfApi, repo_id: str, repo_type: str, private: bool):
    try:
        api.repo_info(repo_id, repo_type=repo_type)
        print(f"  ✓ repo var: {repo_id} ({repo_type})")
    except Exception:
        print(f"  + repo olusturuluyor: {repo_id} ({repo_type}, private={private})")
        create_repo(repo_id, repo_type=repo_type, private=private, exist_ok=True)


def push_files(api: HfApi, repo_id: str, repo_type: str,
               files: list, target_subdir: str = ""):
    total_size = 0
    pushed = 0
    skipped = 0
    for f in files:
        f = Path(f)
        if not f.exists():
            print(f"  - atlandı (yok): {f.name}")
            skipped += 1
            continue
        size = f.stat().st_size
        total_size += size
        target = f"{target_subdir}/{f.name}" if target_subdir else f.name
        print(f"  → {f.name}  ({fmt_size(size)})  upload...", flush=True)
        api.upload_file(
            path_or_fileobj=str(f),
            path_in_repo=target,
            repo_id=repo_id,
            repo_type=repo_type,
            commit_message=f"upload {f.name}",
        )
        pushed += 1
    print(f"\n  ✓ {pushed} dosya yuklendi ({fmt_size(total_size)}), "
          f"{skipped} atlandi")


def write_data_readme():
    """Dataset repo icin README olustur."""
    content = """---
language: tr
license: cc-by-4.0
size_categories:
  - 10B<n<100B
tags:
  - turkish
  - pretraining
  - language-modeling
---

# nanogpt-tr-v5 Data

V5 (200M Türkçe LM) eğitimi için tokenize edilmiş veri.

## Dosyalar

- `v5_stage1.bin` — Web tier (OSCAR, mC4, forum, FineWeb-HQ) ~2.94B token
- `v5_stage2.bin` — Medium tier (BellaTurca, Cosmos, CulturaX, Havadis, Cosmopedia) ~9.03B token
- `v5_stage3.bin` — Premium tier (Wiki, Wikisource, Tezler, Akademik, FinePDFs, Özenli) ~2.97B token
- `v5_val.bin` — Validation (3 stage'in son %1'i, ~150M token)
- `tokenizer-tr-v5.json` — BPE tokenizer, 32K vocab, Stage3 üzerinde eğitildi

## Format

- uint16 token id'leri (vocab=32000 < 65535)
- Numpy memmap ile okunur:
  ```python
  import numpy as np
  data = np.memmap("v5_stage1.bin", dtype=np.uint16, mode="r")
  ```

## Üretim

Bkz. [code repo](https://huggingface.co/{user}/nanogpt-tr-v5-code).
"""
    return content


def write_ckpt_readme():
    """Checkpoint repo icin README."""
    content = """---
language: tr
license: apache-2.0
tags:
  - turkish
  - pretrained
  - gpt
---

# nanogpt-tr-v5 Checkpoints

V5 200M Türkçe pretrained LM, multi-stage curriculum eğitimi.

## Mimari

- 18 layer, 14 head, 896 embd
- 32K vocab, 2048 context
- RoPE (theta=100K) + RMSNorm + SwiGLU + QK-norm
- Logit soft-cap (30) + tied embeddings
- 210M parametre

## Eğitim

- 21.6B token, multi-stage curriculum (web → medium → premium annealing)
- Muon (2D weights) + AdamW (1D + embed)
- bf16 mixed precision, torch.compile
- Lightning AI → Thunder Compute migration

## Yükleme

```python
import torch
from model_v5 import GPTV5, GPTConfigV5

ckpt = torch.load("best_ckpt.pt", weights_only=False)
cfg = GPTConfigV5(**ckpt["config"])
model = GPTV5(cfg)
state = {k.replace("_orig_mod.", ""): v for k, v in ckpt["model"].items()}
model.load_state_dict(state)
```

## Sample

`code` repo'sundaki `06_sample.py` kullanın.
"""
    return content


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--all", action="store_true")
    parser.add_argument("--data", action="store_true")
    parser.add_argument("--ckpt", action="store_true")
    parser.add_argument("--code", action="store_true")
    parser.add_argument("--user", type=str, default=DEFAULT_USER,
                        help="HuggingFace user/org adı")
    parser.add_argument("--private", action="store_true",
                        help="Repo'ları private yap (varsayılan: public)")
    parser.add_argument("--token", type=str, default=None,
                        help="HF token (yoksa env HF_TOKEN veya cache)")
    args = parser.parse_args()

    if not (args.all or args.data or args.ckpt or args.code):
        print("! Hiçbir hedef seçilmedi. --all / --data / --ckpt / --code")
        sys.exit(1)

    api = HfApi(token=args.token or os.environ.get("HF_TOKEN"))
    # Token kontrol
    try:
        whoami = api.whoami()
        print(f"HF user: {whoami['name']}")
    except Exception as e:
        print(f"! HF login problemi: {e}")
        print("  huggingface-cli login  ile bir kere giris yap.")
        sys.exit(1)

    data_repo = f"{args.user}/{REPO_BASE}-data"
    ckpt_repo = f"{args.user}/{REPO_BASE}-ckpts"
    code_repo = f"{args.user}/{REPO_BASE}-code"

    # DATA
    if args.all or args.data:
        print(f"\n{'='*60}\nDATA upload → {data_repo}\n{'='*60}")
        ensure_repo(api, data_repo, "dataset", args.private)
        push_files(api, data_repo, "dataset", DATA_FILES)
        # README
        readme = write_data_readme().replace("{user}", args.user)
        api.upload_file(
            path_or_fileobj=readme.encode(),
            path_in_repo="README.md",
            repo_id=data_repo, repo_type="dataset",
            commit_message="add README",
        )
        print(f"  ✓ README yazildi")

    # CKPT
    if args.all or args.ckpt:
        print(f"\n{'='*60}\nCKPT upload → {ckpt_repo}\n{'='*60}")
        ensure_repo(api, ckpt_repo, "model", args.private)
        push_files(api, ckpt_repo, "model", CKPT_FILES)
        readme = write_ckpt_readme()
        api.upload_file(
            path_or_fileobj=readme.encode(),
            path_in_repo="README.md",
            repo_id=ckpt_repo, repo_type="model",
            commit_message="add README",
        )
        print(f"  ✓ README yazildi")

    # CODE
    if args.all or args.code:
        print(f"\n{'='*60}\nCODE upload → {code_repo}\n{'='*60}")
        ensure_repo(api, code_repo, "model", args.private)
        code_paths = [ROOT / f for f in CODE_FILES]
        push_files(api, code_repo, "model", code_paths)

    print(f"\n{'='*60}\n✓ TAMAMLANDI\n{'='*60}")
    print(f"\nThunder Compute'da indirmek için:")
    print(f"  python hf_pull_v5.py --user {args.user}")


if __name__ == "__main__":
    main()