from __future__ import annotations

import errno
import os
from pathlib import Path

_UNUSABLE_STAT_ERRNOS = {
    errno.ENOENT,
    errno.ENOTDIR,
    errno.EIO,
    errno.ESTALE,
    5,
    116,
}

REMOVABLE_MOUNT_PARENTS = (Path("/media"), Path("/mnt"), Path("/run/media"))
FOUR_TB_DISK = Path("/media/nvidia/4TB")
FOUR_TB2_DISK = Path("/media/nvidia/4TB2")


def is_usable_path(path: Path, *, directory: bool = False) -> bool:
    """True when `path` can be stat'd. NTFS ESTALE/EIO dentries are False."""

    try:
        return path.is_dir() if directory else path.is_file() or path.exists()
    except OSError as error:
        if error.errno in _UNUSABLE_STAT_ERRNOS:
            return False
        raise


def is_usable_dir(path: Path) -> bool:
    return is_usable_path(path, directory=True)


def is_usable_file(path: Path) -> bool:
    try:
        return path.is_file()
    except OSError as error:
        if error.errno in _UNUSABLE_STAT_ERRNOS:
            return False
        raise


def resolve_data_path(path: str | Path) -> Path:
    """Resolve a config data path, skipping dead NTFS project symlinks.

    Configs keep portable values like ``data/tokenized/stage3-sft/manifest.json``.
    When that project symlink is ESTALE/EIO, use ``TINY_GDN_DATA_ROOT`` instead.
    """

    raw = Path(path)
    candidates: list[Path] = [raw]
    parts = raw.parts
    if len(parts) >= 2 and parts[0] == "data":
        candidates.append(data_root().joinpath(*parts[1:]))
    elif "tokenized" in parts:
        index = parts.index("tokenized")
        candidates.append(data_root().joinpath(*parts[index:]))
    for candidate in candidates:
        if is_usable_file(candidate) or is_usable_dir(candidate):
            try:
                return candidate.resolve()
            except OSError:
                return candidate
    if len(parts) >= 2 and parts[0] == "data":
        return data_root().joinpath(*parts[1:])
    return raw


def project_root() -> Path:
    configured = os.environ.get("TINY_GDN_PROJECT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return Path(__file__).resolve().parents[2]


def _existing_dir(path: Path) -> Path | None:
    try:
        if path.is_dir():
            return path
    except OSError:
        return None
    return None


def data_root() -> Path:
    """Large corpora root (tokenized packs, raw downloads, HF cache).

    Portable layout is ``<drive>/llm-data`` next to ``<drive>/llm``.
    """

    configured = os.environ.get("TINY_GDN_DATA_ROOT")
    if configured:
        return Path(configured).expanduser()
    sibling = project_root().parent / "llm-data"
    for candidate in (
        sibling,
        Path("D:/llm-data"),
        Path("/mnt/d/llm-data"),
        Path("/data/llm-data"),
    ):
        found = _existing_dir(candidate)
        if found is not None:
            return found
    if os.name == "nt":
        return Path("D:/llm-data")
    return sibling


def runs_root() -> Path:
    """Checkpoint root. Prefer an explicit env, then a live home mirror.

    ``~/llm-runs`` wins when present so a laptop WSL tree is not shadowed by
    a portable ``D:/llm-runs`` copy. Native Ubuntu uses the sibling
    ``<drive>/llm-runs`` once home has no runs directory.
    """

    configured = os.environ.get("TINY_GDN_LINUX_RUNS_ROOT")
    if configured:
        return Path(configured).expanduser()
    home_runs = Path.home() / "llm-runs"
    if _existing_dir(home_runs) is not None:
        return home_runs
    for candidate in (
        project_root().parent / "llm-runs",
        Path("D:/llm-runs"),
        Path("/mnt/d/llm-runs"),
        Path("/data/llm-runs"),
    ):
        found = _existing_dir(candidate)
        if found is not None:
            return found
    return home_runs


def haiku_data_root() -> Path:
    """Haiku corpora and tokenizer live on the system NVMe, not the 4TB disk."""

    configured = os.environ.get("HAIKU_DATA_ROOT")
    if configured:
        return Path(configured).expanduser()
    return Path.home() / "llm-data" / "haiku"


def haiku_tokenizer_v2_path() -> Path:
    """The Haiku vocabulary with single-token tool tags (ids 16-21) and `<|mask|>` (id 22).

    Every other id, merge and the vocabulary size match `tokenizer/`.
    """

    return haiku_data_root() / "tokenizer-v2" / "tokenizer.json"


def haiku_runs_root() -> Path:
    configured = os.environ.get("HAIKU_RUNS_ROOT")
    if configured:
        return Path(configured).expanduser()
    return Path.home() / "llm-runs" / "haiku"


def haiku_checkpoint_root() -> Path:
    """Run directories for new Haiku trains, on 4TB2 so checkpoints stay off the root disk."""

    configured = os.environ.get("HAIKU_CHECKPOINT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return FOUR_TB2_DISK / "llm-runs" / "haiku"


def senryu_checkpoint_root() -> Path:
    """Senryu run directories, on 4TB2 beside Haiku's. Senryu reads Haiku's data and tokenizer."""

    configured = os.environ.get("SENRYU_CHECKPOINT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return FOUR_TB2_DISK / "llm-runs" / "senryu"


def senryu2_1b_checkpoint_root() -> Path:
    """Senryu-2-1B run directories, on 4TB2: each checkpoint also carries ~35 GB of n-gram tables."""

    configured = os.environ.get("SENRYU2_1B_CHECKPOINT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return FOUR_TB2_DISK / "llm-runs" / "senryu-2-1b"


def senryu2_1_1b_checkpoint_root() -> Path:
    """Senryu-2.1-1B run directories, on 4TB2 beside Senryu-2-1B's, with the same n-gram tables."""

    configured = os.environ.get("SENRYU2_1_1B_CHECKPOINT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return FOUR_TB2_DISK / "llm-runs" / "senryu-2.1-1b"


def senryu2_2_1b_checkpoint_root() -> Path:
    """Senryu-2.2-1B run directories, on 4TB2 beside Senryu-2.1-1B's, with the same n-gram tables."""

    configured = os.environ.get("SENRYU2_2_1B_CHECKPOINT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return FOUR_TB2_DISK / "llm-runs" / "senryu-2.2-1b"


FUGUE_STAGE_METADATA_FILE = "stage.json"
FUGUE_EXPERIMENT_STATUS_FILE = "experiment.json"


def fugue_experiment_root() -> Path:
    """The Rondo/Monody experiment.

    ``<model>/<model>-train`` holds rolling checkpoints, ``<model>/<model>-base``
    and ``<model>/<model>-pass-N`` frozen stage snapshots (with
    ``FUGUE_STAGE_METADATA_FILE``), and the root ``FUGUE_EXPERIMENT_STATUS_FILE``.
    """

    configured = os.environ.get("FUGUE_EXPERIMENT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return FOUR_TB2_DISK / "llm-runs" / "rondo-monody"


def prompttercet_checkpoint_root() -> Path:
    """PromptTercet run directories, on 4TB2: checkpoints carry ~17 GB of n-gram tables and moments."""

    configured = os.environ.get("PROMPTTERCET_CHECKPOINT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return FOUR_TB2_DISK / "llm-runs" / "prompttercet"


def prompttercet1_1_0_5b_checkpoint_root() -> Path:
    """PromptTercet-1.1-0.5B run directories, on 4TB2 beside PromptTercet's, with the same n-gram tables."""

    configured = os.environ.get("PROMPTTERCET1_1_0_5B_CHECKPOINT_ROOT")
    if configured:
        return Path(configured).expanduser()
    return FOUR_TB2_DISK / "llm-runs" / "prompttercet-1.1-0.5b"


def translate_haiku_data_root() -> Path:
    """TranslateHaiku tokenizer and caches live on NVMe, never on the 9TB corpus disk."""

    configured = os.environ.get("TRANSLATE_HAIKU_DATA_ROOT")
    if configured:
        return Path(configured).expanduser()
    return Path.home() / "llm-data" / "translate-haiku"


def translate_haiku_runs_root() -> Path:
    configured = os.environ.get("TRANSLATE_HAIKU_RUNS_ROOT")
    if configured:
        return Path(configured).expanduser()
    return Path.home() / "llm-runs" / "translate-haiku"


def assert_translate_write_path(path: Path, *, kind: str = "path") -> Path:
    """Refuse corpus-disk writes. The 4TB2 stick is only for the Windows connector."""

    resolved = path.expanduser().resolve()
    tercet = tercet_tokenizer_dir().resolve()
    if resolved == tercet or tercet in resolved.parents:
        raise ValueError(f"Refusing to write TranslateHaiku {kind} over the Tercet tokenizer: {resolved}")
    if resolved == FOUR_TB_DISK or FOUR_TB_DISK in resolved.parents:
        raise ValueError(
            f"TranslateHaiku {kind} must stay off the 4TB corpus disk: {resolved}"
        )
    return resolved


def tercet_tokenizer_dir() -> Path:
    return project_root() / "artifacts" / "tokenizer"


def mount_point(path: Path) -> Path:
    """Mount point of the filesystem that holds `path`, or would hold it once created."""

    current = path.expanduser().resolve()
    while not os.path.ismount(current):
        current = current.parent
    return current


def assert_mounted_write_path(path: Path, *, kind: str = "path") -> Path:
    """Refuse a path under /media, /mnt, or /run/media whose drive is not mounted.

    Writing through the bare mount point would fill the root filesystem instead.
    """

    resolved = path.expanduser().resolve()
    mount = mount_point(resolved)
    for parent in REMOVABLE_MOUNT_PARENTS:
        mounted = mount == parent or parent in mount.parents
        if parent in resolved.parents and not mounted:
            raise ValueError(
                f"{kind} {resolved} is not on a mounted drive and would land on "
                f"the filesystem mounted at {mount}; mount the drive first"
            )
    return resolved


def assert_haiku_write_path(path: Path, *, kind: str = "path") -> Path:
    """Refuse writes over the Tercet vocab, onto the 4TB disk, or onto an unmounted drive.

    The 4TB disk's USB bridge sporadically rejects 16-byte writes and drops them.
    """

    resolved = path.expanduser().resolve()
    tercet = tercet_tokenizer_dir().resolve()
    if resolved == tercet or tercet in resolved.parents:
        raise ValueError(f"Refusing to write Haiku {kind} over the Tercet tokenizer: {resolved}")
    if resolved == FOUR_TB_DISK or FOUR_TB_DISK in resolved.parents:
        raise ValueError(f"Haiku {kind} must stay off the 4TB disk: {resolved}")
    return assert_mounted_write_path(resolved, kind=f"Haiku {kind}")


def load_hf_token() -> None:
    """Populate ``HF_TOKEN`` from the environment or a standard token file."""

    existing = os.environ.get("HF_TOKEN") or os.environ.get(
        "HUGGING_FACE_HUB_TOKEN"
    )
    if existing:
        os.environ.setdefault("HF_TOKEN", existing)
        os.environ.setdefault("HUGGING_FACE_HUB_TOKEN", existing)
        return
    candidates = [
        Path.home() / ".cache" / "huggingface" / "token",
    ]
    hf_home = os.environ.get("HF_HOME")
    if hf_home:
        candidates.append(Path(hf_home) / "token")
    candidates.append(data_root() / ".hf" / "token")
    candidates.append(Path("/mnt/c/Users/roblo/.cache/huggingface/token"))
    for path in candidates:
        try:
            if not path.is_file():
                continue
            token = path.read_text(encoding="utf-8").strip()
        except OSError:
            continue
        if not token:
            continue
        os.environ["HF_TOKEN"] = token
        os.environ.setdefault("HUGGING_FACE_HUB_TOKEN", token)
        return