"""Node-local startup snapshots and inference output naming (no tensor loading)."""

import logging
import os
import tempfile
from pathlib import Path

from omegaconf import OmegaConf


def setup_startup_logging(config, *, training):
    """Append one local file per global rank; the trainer records successful restore."""
    if training:
        directory = Path(config.log_dir)
    else:
        output = Path(config.output_path)
        directory = output.parent / "logs" / output.name
    directory.mkdir(parents=True, exist_ok=True)
    rank = int(os.environ.get("RANK", "0")) if "LOCAL_RANK" in os.environ else 0
    logger = logging.getLogger("looking_ahead")
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)
    logger.setLevel(logging.INFO)
    logger.propagate = False
    for handler in logger.handlers[:]:
        logger.removeHandler(handler)
        if handler in root_logger.handlers:
            root_logger.removeHandler(handler)
        handler.close()
    for handler in (
        logging.FileHandler(directory / f"rank_{rank}.log", mode="a", encoding="utf-8"),
        logging.StreamHandler(),
    ):
        handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
        logger.addHandler(handler)
        root_logger.addHandler(handler)
    logger.info(
        "Starting session config=%s resume_source=%s",
        config.config_name,
        config.resume_checkpoint_path if training else None,
    )
    return logger


def configure_debug_environment():
    if os.environ.get("DEBUG") == "1":
        os.environ["WAN_VIS_CAUSAL_MASK"] = "1"
    else:
        os.environ.setdefault("WAN_VIS_CAUSAL_MASK", "0")


def publish_config(config, destination: Path) -> None:
    """LOCAL_RANK=0 publishes a closed temporary YAML by same-directory replacement."""
    if int(os.environ.get("LOCAL_RANK", "0")) != 0:
        return
    destination = Path(destination)
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            suffix=".yaml",
            prefix=".startup-",
            dir=destination.parent,
            delete=False,
        ) as stream:
            temporary = Path(stream.name)
            OmegaConf.save(config, stream, resolve=True)
        os.replace(temporary, destination)
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)


def resolve_inference_paths(config) -> None:
    """Resolve inference paths without loading or validating checkpoint contents."""
    checkpoint = config.checkpoint_path
    if checkpoint is not None:
        checkpoint_name = Path(checkpoint).parent.name or "checkpoint_model_init"
    else:
        root = Path(config.expdir) / "checkpoints"
        if config.checkpoint_name is not None:
            checkpoint_name = config.checkpoint_name
            checkpoint = str(root / checkpoint_name / "model.pt")
            if not Path(checkpoint).is_file():
                raise FileNotFoundError(
                    f"checkpoint_name: specified checkpoint not found: {checkpoint}"
                )
        else:
            candidates = sorted(root.glob("checkpoint_model_*")) if root.is_dir() else []
            checkpoint = str(candidates[-1] / "model.pt") if candidates else None
            checkpoint_name = candidates[-1].name if candidates else "checkpoint_model_init"
    if checkpoint is None and config.use_ema:
        raise FileNotFoundError("use_ema: no training checkpoint was requested or discovered")
    if checkpoint is None:
        print("[inference] No training checkpoint loaded; retaining pipeline initial weights.")
    config.checkpoint_path = checkpoint
    if config.output_path is not None:
        return
    data_path = config.get("image_path") or config.get("video_path") or config.data_path
    if config.use_train_data:
        data_path = next(
            source.data_path
            for source in config.data_sources
            if source.id == config.inference_train_data_source_id
        )
    suffix = ""
    if config.model_kwargs.sink_size > 0:
        suffix += f"-sink={config.model_kwargs.sink_size}"
    if config.use_ema:
        suffix += "-ema"
    if config.use_train_data:
        suffix += "-train_data"
    if config.get("v2v", False):
        context = config.v2v_context_latents
        suffix += f"-v2v-k={context if context is not None else 'auto'}"
    if config.i2v:
        suffix += "-i2v"
    if config.get("image_path") is not None:
        suffix += "-single_image"
    config.output_path = str(
        Path(config.expdir)
        / "inference"
        / f"{checkpoint_name}-{Path(data_path).stem}"
        / f"test_{config.duration}{suffix}"
    )