from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Mapping

from pdftomd.models import DocumentIdentity


MODEL_CACHE_ENV = "PDFTOMD_MODEL_CACHE"


@dataclass(frozen=True)
class ModelCachePolicy:
    root: Path

    @classmethod
    def from_environment(
        cls,
        *,
        project_root: str | Path,
        explicit_model_cache: str | Path | None = None,
        env: Mapping[str, str] | None = None,
    ) -> "ModelCachePolicy":
        source = dict(os.environ if env is None else env)
        if explicit_model_cache is not None:
            root = Path(explicit_model_cache)
        elif source.get(MODEL_CACHE_ENV):
            root = Path(source[MODEL_CACHE_ENV])
        else:
            root = Path(project_root) / ".models"
        return cls(root=root.expanduser())

    @property
    def marker_dir(self) -> Path:
        return self.root / "marker"

    @property
    def nougat_dir(self) -> Path:
        return self.root / "nougat"

    @property
    def huggingface_home(self) -> Path:
        return self.root / "huggingface"

    @property
    def huggingface_hub_cache(self) -> Path:
        return self.huggingface_home / "hub"

    def to_environment(self, *, offline: bool = False) -> dict[str, str]:
        values = {
            "HF_HOME": str(self.huggingface_home),
            "HUGGINGFACE_HUB_CACHE": str(self.huggingface_hub_cache),
            "TRANSFORMERS_CACHE": str(self.huggingface_home),
        }
        if offline:
            values["HF_HUB_OFFLINE"] = "1"
            values["TRANSFORMERS_OFFLINE"] = "1"
        return values


@dataclass(frozen=True)
class RuntimeArtifactPaths:
    root: Path
    log_file: Path
    resume_state_file: Path

    @classmethod
    def from_output_root(
        cls,
        output_root: str | Path,
        document: DocumentIdentity,
    ) -> "RuntimeArtifactPaths":
        root = Path(output_root) / ".pdftomd-runtime" / document.slug
        return cls(
            root=root,
            log_file=root / "logs" / "conversion.log",
            resume_state_file=root / "state" / "resume-state.json",
        )
