# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Centralized path resolution for Magenta RT.

All paths resolve under MAGENTA_HOME/magenta-rt-v2 (where MAGENTA_HOME defaults to ~/Documents/Magenta).
Override with the MAGENTA_HOME environment variable.
"""

import os
import pathlib
from typing import Union


# Configurable root for all downloaded assets and models.
_MAGENTA_BASE = pathlib.Path(
    os.environ.get("MAGENTA_HOME", pathlib.Path.home() / "Documents" / "Magenta")
)
_MAGENTA_HOME = _MAGENTA_BASE / "magenta-rt-v2"

# Default model directory name (under ~/Documents/Magenta/magenta-rt-v2/models/).
DEFAULT_MODEL_NAME = "mrt2_base"
DEFAULT_CHECKPOINT = "mrt2_base.safetensors"


def magenta_home() -> pathlib.Path:
    """Returns the magenta home directory (default: ~/Documents/Magenta/magenta-rt-v2)."""
    return _MAGENTA_HOME


def set_magenta_home(path: Union[pathlib.Path, str]) -> None:
    """Override the magenta home directory at runtime."""
    global _MAGENTA_HOME
    if isinstance(path, str):
        path = pathlib.Path(path)
    _MAGENTA_HOME = path


# ---------------------------------------------------------------------------
# Resource directories
# ---------------------------------------------------------------------------


def resources_dir() -> pathlib.Path:
    """~/Documents/Magenta/magenta-rt-v2/resources — shared resource files (musiccoca, spectrostream)."""
    return _MAGENTA_HOME / "resources"


def musiccoca_dir() -> pathlib.Path:
    """~/Documents/Magenta/magenta-rt-v2/resources/musiccoca — MusicCoCa TFLite models."""
    return resources_dir() / "musiccoca"


def spectrostream_dir() -> pathlib.Path:
    """~/Documents/Magenta/magenta-rt-v2/resources/spectrostream — SpectroStream weights."""
    return resources_dir() / "spectrostream"


def models_dir() -> pathlib.Path:
    """~/Documents/Magenta/magenta-rt-v2/models — exported .mlxfn model directories."""
    return _MAGENTA_HOME / "models"


def default_model_dir() -> pathlib.Path:
    """~/Documents/Magenta/magenta-rt-v2/models/<DEFAULT_MODEL_NAME> — the default model to load."""
    return models_dir() / DEFAULT_MODEL_NAME


def outputs_dir() -> pathlib.Path:
    """~/Documents/Magenta/magenta-rt-v2/outputs — generation and export outputs."""
    d = _MAGENTA_HOME / "outputs"
    d.mkdir(parents=True, exist_ok=True)
    return d


def checkpoints_dir() -> pathlib.Path:
    """~/Documents/Magenta/magenta-rt-v2/checkpoints — full safetensors from Linen models."""
    d = _MAGENTA_HOME / "checkpoints"
    d.mkdir(parents=True, exist_ok=True)
    return d


def resolve_checkpoint(filename: str) -> pathlib.Path:
    """Resolve a checkpoint file path.

    First check if literal filepath exists; fallback to ~/Documents/Magenta/magenta-rt-v2/checkpoints/<filename>.

    Args:
        filename: Checkpoint filename ending in `.safetensors`

    Returns:
        Path to the checkpoint file (may not exist yet).
    """
    if os.path.isfile(filename):
        return filename
    return checkpoints_dir() / filename

# ---------------------------------------------------------------------------
# Path resolution without fallbacks — everything in ~/Documents/Magenta/magenta-rt-v2)
# ---------------------------------------------------------------------------

def resolve_encoder_weights() -> pathlib.Path:
    """Returns ~/Documents/Magenta/magenta-rt-v2/resources/spectrostream/encoder.safetensors."""
    return spectrostream_dir() / "encoder.safetensors"


def resolve_decoder_weights() -> pathlib.Path:
    """Returns ~/Documents/Magenta/magenta-rt-v2/resources/spectrostream/decoder.safetensors."""
    return spectrostream_dir() / "decoder.safetensors"