"""Auto-push trained model to the Hugging Face Hub and clear local training cache (PRIVATE)."""
from __future__ import annotations

import gc
import shutil
from pathlib import Path
from typing import Optional

PLACEHOLDER_MARK = "PUT_YOUR_TOKEN"

_CARD = """---
language: [hi, en]
library_name: transformers
base_model: {base}
tags: [rss-ai, hinglish, thinking-modes, lora]
---
# RSS AI

Trained with the RSS AI Engine (continual pre-training + rational SFT).
Modes: `thinking`, `non_thinking`, `max_thinking` (selected through `[MODE: ...]` in the system prompt).
Base model: `{base}`
"""


def _valid(token: Optional[str]) -> bool:
    return bool(token) and PLACEHOLDER_MARK not in token


def upload_model_folder(folder: Path, repo_id: str, token: Optional[str], base_model: str,
                        private: bool = True) -> bool:
    folder = Path(folder)
    if not _valid(token):
        print("   [upload] HF_TOKEN missing/placeholder -> upload skipped (model stays in "
              f"{folder}). Put a WRITE token in .env and run: main.py --mode train_all --upload-only")
        return False
    if not (folder / "rss_meta.json").exists():
        print(f"   [upload] {folder} has no rss_meta.json -> refusing to upload an unfinished model")
        return False
    from huggingface_hub import HfApi
    api = HfApi(token=token)
    try:
        api.create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
        (folder / "README.md").write_text(_CARD.format(base=base_model), encoding="utf-8")
        print(f"   [upload] pushing {folder} -> https://huggingface.co/{repo_id} ...")
        api.upload_folder(folder_path=str(folder), repo_id=repo_id, repo_type="model",
                          commit_message="RSS AI: trained model + tokenizer",
                          ignore_patterns=[".rss_overlay/*", "*.incomplete"])
        print("   [upload] SUCCESS - model is safe on the Hub")
        return True
    except Exception as exc:
        print(f"   [upload] FAILED ({type(exc).__name__}: {exc}). Local copy kept in {folder}")
        return False


def download_trained(repo_id: str, dest: Path, token: Optional[str]) -> Path:
    from huggingface_hub import snapshot_download
    dest = Path(dest)
    print(f"   [download] pulling trained model {repo_id} -> {dest}")
    snapshot_download(repo_id=repo_id, local_dir=str(dest), token=token if _valid(token) else None)
    if not (dest / "rss_meta.json").exists():
        raise RuntimeError(f"{repo_id} does not look like an RSS AI export (rss_meta.json missing)")
    return dest


def clear_training_cache(root: Path) -> None:
    """Free disk + RAM after a successful upload. The final trained_model/ folder is kept."""
    root = Path(root)
    for name in ("checkpoints", "offload_disk"):
        target = root / name
        if target.exists():
            shutil.rmtree(target, ignore_errors=True)
            print(f"   [cleanup] removed {target}/")
    tmp = root / "cleaned_data" / "_repo_text.tmp.jsonl"
    if tmp.exists():
        tmp.unlink()
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass
    print("   [cleanup] RAM/VRAM released")