# Copyright 2025-2026 Dimensional Inc.
#
# 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.

from functools import cache
import os
from pathlib import Path
import platform
import subprocess
import sys
import tarfile
import tempfile

from dimos.constants import DIMOS_PROJECT_ROOT
from dimos.utils.logging_config import setup_logger

logger = setup_logger()


def _get_user_data_dir() -> Path:
    """Get platform-specific user data directory."""
    system = platform.system()
    # if virtual env is available, use it to keep venv's from fighting over data
    # a better fix for large files will be added later to minimize storage duplication
    if os.environ.get("VIRTUAL_ENV"):
        venv_data_dir = Path(
            f"{os.environ.get('VIRTUAL_ENV')}/lib/python{sys.version_info.major}.{sys.version_info.minor}/site-packages/dimos/data"
        )
        return venv_data_dir

    if system == "Linux":
        # Use XDG_DATA_HOME if set, otherwise default to ~/.local/share
        xdg_data_home = os.environ.get("XDG_DATA_HOME")
        if xdg_data_home:
            return Path(xdg_data_home) / "dimos"
        return Path.home() / ".local" / "share" / "dimos"
    elif system == "Darwin":  # macOS
        return Path.home() / "Library" / "Application Support" / "dimos"
    else:
        # Fallback for other systems
        return Path.home() / ".dimos"


@cache
def _get_repo_root() -> Path:
    # Check if running from git repo
    if (DIMOS_PROJECT_ROOT / ".git").exists():
        return DIMOS_PROJECT_ROOT

    # Running as installed package - clone repo to data dir
    try:
        data_dir = _get_user_data_dir()
        data_dir.mkdir(parents=True, exist_ok=True)
        # Test if writable
        test_file = data_dir / ".write_test"
        test_file.touch()
        test_file.unlink()
        logger.info(f"Using local user data directory at '{data_dir}'")
    except (OSError, PermissionError):
        # Fall back to temp dir if data dir not writable
        data_dir = Path(tempfile.gettempdir()) / "dimos"
        data_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"Using tmp data directory at '{data_dir}'")

    repo_dir = data_dir / "repo"

    # Clone if not already cloned
    if not (repo_dir / ".git").exists():
        try:
            env = os.environ.copy()
            env["GIT_LFS_SKIP_SMUDGE"] = "1"
            subprocess.run(
                [
                    "git",
                    "clone",
                    "--depth",
                    "1",
                    "--branch",
                    "main",
                    "https://github.com/dimensionalOS/dimos.git",
                    str(repo_dir),
                ],
                check=True,
                capture_output=True,
                text=True,
                env=env,
            )
        except subprocess.CalledProcessError as e:
            raise RuntimeError(
                f"Failed to clone dimos repository: {e.stderr}\n"
                f"Make sure you can access https://github.com/dimensionalOS/dimos.git"
            )

    return repo_dir


@cache
def get_data_dir(extra_path: str | None = None) -> Path:
    if extra_path:
        return _get_repo_root() / "data" / extra_path
    return _get_repo_root() / "data"


@cache
def _get_lfs_dir() -> Path:
    return get_data_dir() / ".lfs"


def _check_git_lfs_available() -> bool:
    missing = []

    # Check if git is available
    try:
        subprocess.run(["git", "--version"], capture_output=True, check=True, text=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        missing.append("git")

    # Check if git-lfs is available
    try:
        subprocess.run(["git-lfs", "version"], capture_output=True, check=True, text=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        missing.append("git-lfs")

    if missing:
        raise RuntimeError(
            f"Missing required tools: {', '.join(missing)}.\n\n"
            "Git LFS installation instructions: https://git-lfs.github.io/"
        )

    return True


def _is_lfs_pointer_file(file_path: Path) -> bool:
    try:
        # LFS pointer files are small (typically < 200 bytes) and start with specific text
        if file_path.stat().st_size > 1024:  # LFS pointers are much smaller
            return False

        with open(file_path, encoding="utf-8") as f:
            first_line = f.readline().strip()
            return first_line.startswith("version https://git-lfs.github.com/spec/")

    except (UnicodeDecodeError, OSError):
        return False


def _lfs_pull(file_path: Path, repo_root: Path) -> None:
    try:
        relative_path = file_path.relative_to(repo_root)

        env = os.environ.copy()
        env["GIT_LFS_FORCE_PROGRESS"] = "1"

        subprocess.run(
            ["git", "lfs", "pull", "--include", str(relative_path)],
            cwd=repo_root,
            check=True,
            env=env,
        )
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Failed to pull LFS file {file_path}: {e}")

    return None


def _decompress_archive(filename: str | Path) -> Path:
    target_dir = get_data_dir()
    filename_path = Path(filename)
    with tarfile.open(filename_path, "r:gz") as tar:
        tar.extractall(target_dir)
    return target_dir / filename_path.name.replace(".tar.gz", "")


def _pull_lfs_archive(filename: str | Path) -> Path:
    # Check Git LFS availability first
    _check_git_lfs_available()

    # Find repository root
    repo_root = _get_repo_root()

    # Construct path to test data file
    file_path = _get_lfs_dir() / (str(filename) + ".tar.gz")

    # Check if file exists
    if not file_path.exists():
        raise FileNotFoundError(
            f"Test file '{filename}' not found at {file_path}. "
            f"Make sure the file is committed to Git LFS in the tests/data directory."
        )

    # If it's an LFS pointer file, ensure LFS is set up and pull the file
    if _is_lfs_pointer_file(file_path):
        _lfs_pull(file_path, repo_root)

        # Verify the file was actually downloaded
        if _is_lfs_pointer_file(file_path):
            raise RuntimeError(
                f"Failed to download LFS file '{filename}'. The file is still a pointer after attempting to pull."
            )

    return file_path


def get_data(name: str | Path) -> Path:
    """
    Get the path to a test data, downloading from LFS if needed.

    This function will:
    1. Check that Git LFS is available
    2. Locate the file in the tests/data directory
    3. Initialize Git LFS if needed
    4. Download the file from LFS if it's a pointer file
    5. Return the Path object to the actual file or dir

    Supports nested paths like "dataset/subdir/file.jpg" - will download and
    decompress "dataset" archive but return the full nested path.

    Args:
        name: Name of the test file or dir, optionally with nested path
              (e.g., "lidar_sample.bin" or "dataset/frames/001.png")

    Returns:
        Path: Path object to the test file or dir

    Raises:
        RuntimeError: If Git LFS is not available or LFS operations fail
        FileNotFoundError: If the test file doesn't exist

    Usage:
        # Simple file/dir
        file_path = get_data("sample.bin")

        # Nested path - downloads "dataset" archive, returns path to nested file
        frame = get_data("dataset/frames/001.png")
    """
    data_dir = get_data_dir()
    file_path = data_dir / name

    # already pulled and decompressed, return it directly
    if file_path.exists():
        return file_path

    # extract archive root (first path component) and nested path
    path_parts = Path(name).parts
    archive_name = path_parts[0]
    nested_path = Path(*path_parts[1:]) if len(path_parts) > 1 else None

    # download and decompress the archive root
    archive_path = _decompress_archive(_pull_lfs_archive(archive_name))

    # return full path including nested components
    if nested_path:
        return archive_path / nested_path
    return archive_path


class LfsPath(type(Path())):  # type: ignore[misc]
    """
    A Path subclass that lazily downloads LFS data when accessed.

    This is useful for both lazy loading and differentiating between LFS paths and regular paths.

    This class wraps pathlib.Path and ensures that get_data() is called
    before any meaningful filesystem operation, making LFS data lazy-loaded.

    Usage:
        path = LfsPath("sample_data")
        # No download yet

        with path.open('rb') as f:  # Downloads now if needed
            data = f.read()

        # Or use any Path operation:
        if path.exists():  # Downloads now if needed
            files = list(path.iterdir())
    """

    def __new__(cls, filename: str | Path) -> "LfsPath":
        # Create instance with a placeholder path to satisfy Path.__new__
        # We use "." as a dummy path that always exists
        instance: LfsPath = super().__new__(cls, ".")  # type: ignore[call-arg]
        # Store the actual filename as an instance attribute
        object.__setattr__(instance, "_lfs_filename", filename)
        object.__setattr__(instance, "_lfs_resolved_cache", None)
        return instance

    def _ensure_downloaded(self) -> Path:
        """Ensure the LFS data is downloaded and return the resolved path."""
        cache: Path | None = object.__getattribute__(self, "_lfs_resolved_cache")
        if cache is None:
            filename = object.__getattribute__(self, "_lfs_filename")
            cache = get_data(filename)
            object.__setattr__(self, "_lfs_resolved_cache", cache)
        return cache

    def __getattribute__(self, name: str) -> object:
        # During Path.__new__(), _lfs_filename hasn't been set yet.
        # Fall through to normal Path behavior until construction is complete.
        try:
            object.__getattribute__(self, "_lfs_filename")
        except AttributeError:
            return object.__getattribute__(self, name)

        # After construction, allow access to our internal attributes directly
        if name in ("_lfs_filename", "_lfs_resolved_cache", "_ensure_downloaded"):
            return object.__getattribute__(self, name)

        # For all other attributes, ensure download first then delegate to resolved path
        resolved = object.__getattribute__(self, "_ensure_downloaded")()
        return getattr(resolved, name)

    def __str__(self) -> str:
        """String representation returns resolved path."""
        return str(self._ensure_downloaded())

    def __fspath__(self) -> str:
        """Return filesystem path, downloading from LFS if needed."""
        return str(self._ensure_downloaded())

    def __truediv__(self, other: object) -> Path:
        """Path division operator - returns resolved path."""
        return self._ensure_downloaded() / other  # type: ignore[operator, return-value]

    def __rtruediv__(self, other: object) -> Path:
        """Reverse path division operator."""
        return other / self._ensure_downloaded()  # type: ignore[operator, return-value]
