"""PUBLIC INTERFACE of Training_Module. Other modules may import ONLY this file.

load_tokenizer(model_dir, recipe)                  -> (tokenizer, surgery)   registers <think> </think> <|im_end|>
finalize_model(model, tokenizer, surgery)          -> surgery                embedding rows + generation ids
run_pretrain(model, tokenizer, plan, cfg, surgery, pretrain_path, checkpoint_dir) -> dict   Stage 1 (CPT)
run_finetune(model, tokenizer, plan, cfg, surgery, sft_path, final_dir, adapter_dir) -> dict Stage 2 (SFT)
read_meta(folder) -> dict | None                   contents of rss_meta.json
upload_model(folder, cfg) -> bool                  push to OUTPUT_HF_REPO (HF_TOKEN)
download_trained(repo, dest, token) -> Path        pull a trained model from the Hub
clear_cache(root)                                  free disk/RAM after a successful upload

cfg keys used: base_model_repo, output_hf_repo, hf_token, conv_ratio, tune (dict of optional overrides)
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, Optional, Tuple

from .internal import elevator_pretrainer as _pretrainer
from .internal import hf_auto_uploader as _uploader
from .internal import rational_sft_trainer as _sft
from .internal import tokenizer_surgeon as _surgeon

MODES = _surgeon.MODES


def load_tokenizer(model_dir: str, recipe: Optional[Dict] = None) -> Tuple[object, Dict]:
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=False)
    surgery = _surgeon.prepare_tokenizer(tokenizer, recipe_override=recipe)
    if surgery["added"]:
        print(f"   [tokenizer] registered new special tokens: {', '.join(surgery['added'])}")
    return tokenizer, surgery


def finalize_model(model, tokenizer, surgery: Dict) -> Dict:
    return _surgeon.finalize_model_embeddings(model, tokenizer, surgery)


def run_pretrain(model, tokenizer, plan: Dict, cfg: Dict, surgery: Dict, pretrain_path, checkpoint_dir) -> Dict:
    return _pretrainer.run_pretrain(model, tokenizer, plan, cfg, surgery, Path(pretrain_path), Path(checkpoint_dir))


def run_finetune(model, tokenizer, plan: Dict, cfg: Dict, surgery: Dict, sft_path, final_dir,
                 adapter_dir: Optional[str] = None) -> Dict:
    return _sft.run_finetune(model, tokenizer, plan, cfg, surgery, Path(sft_path), Path(final_dir),
                             adapter_dir=adapter_dir)


def read_meta(folder) -> Optional[Dict]:
    meta = Path(folder) / "rss_meta.json"
    if not meta.exists():
        return None
    try:
        return json.loads(meta.read_text(encoding="utf-8"))
    except ValueError:
        return None


def upload_model(folder, cfg: Dict) -> bool:
    return _uploader.upload_model_folder(Path(folder), cfg["output_hf_repo"], cfg.get("hf_token"),
                                         cfg["base_model_repo"])


def download_trained(repo: str, dest, token: Optional[str]) -> Path:
    return _uploader.download_trained(repo, Path(dest), token)


def clear_cache(root) -> None:
    _uploader.clear_training_cache(Path(root))