"""Seeding, device/precision selection, JSON I/O, console setup (Windows-friendly)."""
from __future__ import annotations

import contextlib
import json
import os
import random
import sys
from pathlib import Path
from typing import Any

import numpy as np
import torch

os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")  # noisy on Windows
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


def setup_console() -> None:
    """UTF-8 stdout and ANSI colours on Windows terminals (incl. VS Code)."""
    for stream in (sys.stdout, sys.stderr):
        with contextlib.suppress(Exception):
            stream.reconfigure(encoding="utf-8", errors="replace")
    if os.name == "nt":
        os.system("")  # enables VT100 escape processing in cmd/PowerShell


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def get_device(preference: str = "auto") -> torch.device:
    if preference != "auto":
        return torch.device(preference)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def resolve_precision(preference: str, device: torch.device) -> torch.dtype | None:
    """Autocast dtype, or None for plain fp32."""
    if device.type != "cuda":
        return None  # bf16 autocast on MPS/CPU is not worth the trouble here
    if preference == "fp32":
        return None
    if preference == "fp16":
        return torch.float16
    if preference in ("bf16", "auto"):
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
        return torch.float16 if preference == "auto" else None
    raise ValueError(f"Unknown precision {preference!r}")


def autocast(device: torch.device, dtype: torch.dtype | None):
    if dtype is None:
        return contextlib.nullcontext()
    return torch.autocast(device_type=device.type, dtype=dtype)


def describe_device(device: torch.device) -> str:
    if device.type == "cuda":
        props = torch.cuda.get_device_properties(device)
        return f"cuda:{props.name} ({props.total_memory / 2**30:.1f} GB, sm_{props.major}{props.minor})"
    return device.type


def save_json(obj: Any, path: str | Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8", newline="\n") as fh:
        json.dump(obj, fh, indent=2, ensure_ascii=False, default=_json_default)


def load_json(path: str | Path) -> Any:
    with open(path, encoding="utf-8") as fh:
        return json.load(fh)


def _json_default(o: Any) -> Any:
    if isinstance(o, (np.integer,)):
        return int(o)
    if isinstance(o, (np.floating,)):
        return float(o)
    if isinstance(o, Path):
        return str(o)
    raise TypeError(f"Not JSON serialisable: {type(o)}")


def setup_logging(level: int = 20) -> None:
    import logging
    setup_console()
    logging.basicConfig(level=level, format="%(asctime)s %(levelname)s %(name)s: %(message)s",
                        datefmt="%H:%M:%S")
    for noisy in ("transformers", "huggingface_hub", "urllib3", "filelock"):
        logging.getLogger(noisy).setLevel(logging.WARNING)

def unique_dir(base: Path) -> Path:
    """Return `base`, or base-2 / base-3 / ... if it already exists, so runs never overwrite."""
    if not base.exists():
        return base
    n = 2
    while (candidate := base.with_name(f"{base.name}-{n}")).exists():
        n += 1
    return candidate

def experiment_dir(root: Path, group: str) -> Path:
    """Return runs/<group>, or <group>-2, -3 ... if it already exists."""
    candidate = root / group
    n = 1
    while candidate.exists():
        n += 1
        candidate = root / f"{group}-{n}"
    return candidate