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

prepare_base_model(cfg, fresh)  -> purge (optional) + authenticated fresh HF download -> local folder
profile_system(model_dir, cfg)  -> CPU/RAM/GPU/model-size profile
plan_execution(profile, cfg)    -> tier 1/2/3 plan with exact RAM-vs-Disk split
load_model(model_dir, plan, cfg, row_recipe=None) -> model (tier 3 models carry a `_rss_ring` handle)
release_memory()                -> gc + CUDA cache flush
"""
from __future__ import annotations

import gc
from pathlib import Path
from typing import Dict, Optional

from .internal import cache_cleaner_and_downloader as cleaner
from .internal import system_profiler as profiler
from .internal import tier_allocation_planner as planner
from .internal.sequential_ring_loader import SequentialRingLoader


def purge_old_artifacts(cfg: Dict) -> list:
    print("[Hardware] purging old model leftovers ...")
    return cleaner.purge_old_artifacts(Path(cfg["root"]), cfg["base_model_repo"])


def prepare_base_model(cfg: Dict, fresh: bool = False) -> str:
    """Return the local folder that holds BASE_MODEL_REPO (downloads it if needed)."""
    root = Path(cfg["root"])
    if fresh:
        cleaner.purge_old_artifacts(root, cfg["base_model_repo"])
    dest = root / "base_model"
    path = cleaner.download_base_model(cfg["base_model_repo"], dest, cfg.get("hf_token"), force=fresh)
    return str(path)


def profile_system(model_dir: str, cfg: Dict) -> Dict:
    return profiler.profile(model_dir, cfg.get("hf_token") if cleaner.token_is_valid(cfg.get("hf_token")) else None)


def plan_execution(prof: Dict, cfg: Dict) -> Dict:
    force = cfg.get("force_tier")
    plan = planner.make_plan(prof, int(force) if force else None)
    print("[Hardware] system profile:\n" + profiler.describe(prof))
    print("[Hardware] execution plan:\n" + planner.describe_plan(plan, prof))
    return plan


def load_model(model_dir: str, plan: Dict, cfg: Dict, row_recipe: Optional[Dict] = None):
    import torch
    root = Path(cfg["root"])
    torch.set_num_threads(max(1, int(cfg.get("cpu_threads") or (torch.get_num_threads()))))
    print(f"[Hardware] loading model from {model_dir} (tier {plan['tier']}) ...")
    if plan["tier"] == 3:
        loader = SequentialRingLoader(
            model_dir, dtype=getattr(torch, plan["dtype"]), window=plan.get("window") or 2,
            evict_ends=plan.get("evict_ends", True), overlay_dir=root / "offload_disk" / "overlay",
            row_recipe=row_recipe)
        model = loader.build()
        model._rss_ring = loader
    else:
        model = planner.materialize_model(model_dir, plan, root / "offload_disk")
        model._rss_ring = None
    model._rss_plan = plan
    return model


def release_memory() -> None:
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass