from __future__ import annotations

import json
import shutil
import tempfile
from pathlib import Path
from typing import Any, Dict

import torch

from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.store import MemoryStore

try:  # pragma: no cover - optional dependency
    from huggingface_hub import HfApi, hf_hub_download

    _HF_AVAILABLE = True
except Exception:  # pragma: no cover - optional dependency
    HfApi = None
    hf_hub_download = None
    _HF_AVAILABLE = False


class HubBackend:
    """Abstract hub backend; subclasses implement push/pull mechanics."""

    def push(
        self,
        repo_id: str,
        bundle_path: Path,
        manifest: Dict[str, Any],
        *,
        filename: str = "bundle.pt",
        readme: str | None = None,
    ) -> str:
        raise NotImplementedError

    def pull(
        self,
        repo_id: str,
        *,
        filename: str = "bundle.pt",
        revision: str = "main",
        cache_dir: str | None = None,
        token: str | None = None,
    ) -> Path:
        raise NotImplementedError


class LocalHubBackend(HubBackend):
    """Simple backend that stores artifacts under a directory tree."""

    def __init__(self, *, root_dir: str | Path | None = None) -> None:
        self.root = Path(root_dir or ".hub_artifacts")
        self.root.mkdir(parents=True, exist_ok=True)

    def push(
        self,
        repo_id: str,
        bundle_path: Path,
        manifest: Dict[str, Any],
        *,
        filename: str = "bundle.pt",
        readme: str | None = None,
    ) -> str:
        target = self._target_dir(repo_id)
        target.mkdir(parents=True, exist_ok=True)
        shutil.copy2(bundle_path, target / filename)
        (target / "MANIFEST.json").write_text(json.dumps(manifest, indent=2))
        if readme:
            (target / "README.md").write_text(readme)
        return str(target)

    def pull(
        self,
        repo_id: str,
        *,
        filename: str = "bundle.pt",
        revision: str = "main",
        cache_dir: str | None = None,
        token: str | None = None,
    ) -> Path:
        target = self._target_dir(repo_id) / filename
        if not target.exists():
            raise FileNotFoundError(f"No artifact found for {repo_id} at {target}")
        return target

    def _target_dir(self, repo_id: str) -> Path:
        safe_repo = repo_id.replace("/", "__")
        return self.root / safe_repo


class HuggingFaceHubBackend(HubBackend):  # pragma: no cover - network optional
    """Backend that uploads artifacts to the Hugging Face Hub."""

    def __init__(
        self,
        *,
        token: str | None = None,
        repo_type: str = "model",
        exist_ok: bool = True,
    ) -> None:
        if not _HF_AVAILABLE:
            raise ImportError(
                "huggingface_hub is required for the HuggingFaceHubBackend. "
                "Install hippotorch[hub] to enable."
            )
        self.api = HfApi(token=token)
        self.token = token
        self.repo_type = repo_type
        self.exist_ok = exist_ok

    def push(
        self,
        repo_id: str,
        bundle_path: Path,
        manifest: Dict[str, Any],
        *,
        filename: str = "bundle.pt",
        readme: str | None = None,
    ) -> str:
        tmpdir = Path(tempfile.mkdtemp(prefix="hippo_hub_"))
        try:
            target_bundle = tmpdir / filename
            shutil.copy2(bundle_path, target_bundle)
            (tmpdir / "MANIFEST.json").write_text(json.dumps(manifest, indent=2))
            if readme:
                (tmpdir / "README.md").write_text(readme)
            self.api.create_repo(
                repo_id=repo_id,
                repo_type=self.repo_type,
                token=self.token,
                exist_ok=self.exist_ok,
            )
            self.api.upload_folder(
                repo_id=repo_id,
                repo_type=self.repo_type,
                folder_path=str(tmpdir),
                token=self.token,
            )
        finally:
            shutil.rmtree(tmpdir, ignore_errors=True)
        return f"https://huggingface.co/{repo_id}"

    def pull(
        self,
        repo_id: str,
        *,
        filename: str = "bundle.pt",
        revision: str = "main",
        cache_dir: str | None = None,
        token: str | None = None,
    ) -> Path:
        if hf_hub_download is None:
            raise ImportError("huggingface_hub is required to download from the Hub.")
        resolved = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            revision=revision,
            repo_type=self.repo_type,
            cache_dir=cache_dir,
            token=token or self.token,
        )
        return Path(resolved)


def push_memory_to_hub(
    target,
    repo_id: str,
    *,
    backend: str | HubBackend | None = None,
    local_dir: str | None = None,
    token: str | None = None,
    filename: str = "bundle.pt",
    readme: str | None = None,
    **kwargs: Any,
) -> str:
    """Push a buffer or memory snapshot to the configured hub backend."""

    buffer = _coerce_buffer(target, **kwargs)
    backend_obj = _resolve_backend(backend, token=token, local_dir=local_dir)
    tmpdir = Path(tempfile.mkdtemp(prefix="hippo_push_"))
    try:
        bundle_path = tmpdir / filename
        buffer.save(bundle_path)
        manifest = {
            "buffer": {"mixture_ratio": buffer.mixture_ratio},
            "memory": {
                "embed_dim": buffer.memory.embed_dim,
                "size": len(buffer.memory),
            },
            "encoder": {"type": buffer.encoder.__class__.__name__},
        }
        return backend_obj.push(
            repo_id,
            bundle_path,
            manifest,
            filename=filename,
            readme=readme,
        )
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


def load_memory_from_hub(
    repo_id: str,
    *,
    backend: str | HubBackend | None = None,
    local_dir: str | None = None,
    token: str | None = None,
    revision: str = "main",
    filename: str = "bundle.pt",
    cache_dir: str | None = None,
    device: str | torch.device | None = None,
    return_buffer: bool = False,
) -> MemoryStore | HippocampalReplayBuffer:
    """Download an artifact from the hub and return a memory or buffer."""

    backend_obj = _resolve_backend(backend, token=token, local_dir=local_dir)
    bundle_path = backend_obj.pull(
        repo_id,
        filename=filename,
        revision=revision,
        cache_dir=cache_dir,
        token=token,
    )
    buffer = HippocampalReplayBuffer.load(bundle_path, device=device or "cpu")
    return buffer if return_buffer else buffer.memory


def _resolve_backend(
    backend: str | HubBackend | None,
    *,
    token: str | None,
    local_dir: str | None,
) -> HubBackend:
    if isinstance(backend, HubBackend):
        return backend
    if backend == "hf":
        return HuggingFaceHubBackend(token=token)
    if backend == "local" or backend is None:
        return LocalHubBackend(root_dir=local_dir)
    raise ValueError(f"Unsupported hub backend {backend!r}")


def _coerce_buffer(target, **kwargs: Any) -> HippocampalReplayBuffer:
    if isinstance(target, HippocampalReplayBuffer):
        return target
    if isinstance(target, MemoryStore):
        encoder = kwargs.pop("encoder", None)
        if encoder is None:
            raise ValueError("encoder must be provided when pushing a raw MemoryStore.")
        mixture = float(kwargs.pop("mixture_ratio", 0.0))
        return HippocampalReplayBuffer(
            memory=target, encoder=encoder, mixture_ratio=mixture
        )
    raise TypeError(
        "push_memory_to_hub expects a HippocampalReplayBuffer or MemoryStore."
    )


__all__ = [
    "push_memory_to_hub",
    "load_memory_from_hub",
    "HubBackend",
    "LocalHubBackend",
    "HuggingFaceHubBackend",
]
