"""
Utilities for uv package manager integration.

This module provides functions for detecting uv projects and exporting dependencies
via ``uv export`` for automatic dependency inference during model logging.
"""

import logging
import os
import re
import shutil
import subprocess
from pathlib import Path
from typing import NamedTuple

from packaging.version import Version

from mlflow.environment_variables import MLFLOW_LOG_UV_FILES

_logger = logging.getLogger(__name__)

# Minimum uv version required for ``uv export`` functionality
_MIN_UV_VERSION = Version("0.6.10")

# File names used for uv project detection and artifacts
_UV_LOCK_FILE = "uv.lock"
_PYPROJECT_FILE = "pyproject.toml"
_PYTHON_VERSION_FILE = ".python-version"


def get_uv_version() -> Version | None:
    """
    Get the installed uv version.

    Returns:
        The uv version as a packaging.version.Version object, or None if uv is not installed
        or version cannot be determined.
    """
    uv_bin = shutil.which("uv")
    if uv_bin is None:
        return None

    try:
        result = subprocess.run(
            [uv_bin, "--version"],
            capture_output=True,
            text=True,
            check=True,
        )
        # Output format: "uv 0.5.0 (abc123 2024-01-01)"
        version_str = result.stdout.strip().split()[1]
        return Version(version_str)
    except (subprocess.CalledProcessError, IndexError, ValueError) as e:
        _logger.debug(f"Failed to determine uv version: {e}")
        return None


def _get_uv_binary() -> str | None:
    """Return the uv binary path if installed and meeting version requirements, else None."""
    uv_bin = shutil.which("uv")
    if uv_bin is None:
        return None

    version = get_uv_version()
    if version is None:
        return None

    if version < _MIN_UV_VERSION:
        _logger.debug(f"uv version {version} is below minimum required version {_MIN_UV_VERSION}")
        return None

    return uv_bin


def is_uv_available() -> bool:
    """
    Check if uv is installed and meets the minimum version requirement.

    Returns:
        True if uv is installed and version >= 0.6.10, False otherwise.
    """
    return _get_uv_binary() is not None


class UVProjectInfo(NamedTuple):
    """Paths for a detected uv project."""

    uv_lock: Path
    pyproject: Path


def detect_uv_project(directory: str | Path | None = None) -> UVProjectInfo | None:
    """
    Detect if the given directory is a uv project.

    A uv project is identified by the presence of BOTH ``uv.lock`` and ``pyproject.toml``
    in the specified directory.

    Args:
        directory: The directory to check. Defaults to the current working directory.

    Returns:
        A UVProjectInfo with paths to uv.lock and pyproject.toml,
        or None if the directory is not a uv project.
    """
    directory = Path.cwd() if directory is None else Path(directory)

    uv_lock_path = directory / _UV_LOCK_FILE
    pyproject_path = directory / _PYPROJECT_FILE

    if uv_lock_path.exists() and pyproject_path.exists():
        _logger.info(
            f"Detected uv project: found {_UV_LOCK_FILE} and {_PYPROJECT_FILE} in {directory}"
        )
        return UVProjectInfo(uv_lock=uv_lock_path, pyproject=pyproject_path)

    return None


def export_uv_requirements(
    directory: str | Path | None = None,
    no_dev: bool = True,
    no_hashes: bool = True,
    frozen: bool = True,
    groups: list[str] | None = None,
    extras: list[str] | None = None,
) -> list[str] | None:
    """
    Export dependencies from a uv project to pip-compatible requirements.

    Runs ``uv export`` to generate a list of pinned dependencies from the uv lock file.
    Environment markers (e.g. ``; python_version < '3.12'``) are preserved in the output
    so that pip can evaluate them at install time.

    Args:
        directory: The uv project directory. Defaults to the current working directory.
        no_dev: Exclude development dependencies. Defaults to True.
        no_hashes: Omit hashes from output. Defaults to True.
        frozen: Use frozen lockfile without updating. Defaults to True.
        groups: Optional list of dependency groups to include (additive with project deps).
            Maps to ``uv export --group <name>``.
        extras: Optional list of optional dependency extras to include.
            Maps to ``uv export --extra <name>``.

    Returns:
        A list of requirement strings (e.g., ``["requests==2.28.0", "numpy==1.24.0"]``),
        or None if export fails.
    """
    uv_bin = _get_uv_binary()
    if uv_bin is None:
        _logger.warning(
            "uv is not available or version is below minimum required. "
            "Falling back to pip-based inference."
        )
        return None
    directory = Path.cwd() if directory is None else Path(directory)

    cmd = [uv_bin, "export"]

    if no_dev:
        cmd.append("--no-dev")
    if no_hashes:
        cmd.append("--no-hashes")
    if frozen:
        cmd.append("--frozen")

    if groups:
        for group in groups:
            cmd.extend(["--group", group])

    if extras:
        for extra in extras:
            cmd.extend(["--extra", extra])

    # Additional flags for cleaner output
    cmd.extend([
        "--no-header",  # Omit the "autogenerated by uv" comment
        "--no-emit-project",  # Don't emit the project itself
        "--no-annotate",  # Omit "# via ..." comment annotations
    ])

    try:
        _logger.debug(f"Running uv export: {' '.join(str(c) for c in cmd)}")
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
            cwd=directory,
        )

        requirements = [
            line.strip()
            for line in result.stdout.strip().splitlines()
            if line.strip() and not line.strip().startswith("#")
        ]

        _logger.info(f"Exported {len(requirements)} dependencies via uv")
        return requirements

    except subprocess.CalledProcessError as e:
        _logger.warning(
            f"uv export failed with exit code {e.returncode}. "
            f"stderr: {e.stderr}. Falling back to pip-based inference."
        )
        return None
    except Exception as e:
        _logger.warning(f"uv export failed: {e}. Falling back to pip-based inference.")
        return None


def copy_uv_project_files(
    dest_dir: str | Path,
    source_dir: str | Path,
) -> bool:
    """
    Copy uv project files to the model artifact directory.

    Copies uv.lock, pyproject.toml, and .python-version (if present) to preserve
    uv project configuration as model artifacts for reproducibility.

    Can be disabled by setting MLFLOW_LOG_UV_FILES=false environment variable
    for large projects where uv.lock size is a concern.

    Args:
        dest_dir: The destination directory (model artifact directory).
        source_dir: The source directory containing uv project files.

    Returns:
        True if uv files were copied, False otherwise.
    """
    if not MLFLOW_LOG_UV_FILES.get():
        _logger.info("uv file logging disabled via MLFLOW_LOG_UV_FILES environment variable")
        return False

    dest_dir = Path(dest_dir)
    source_dir = Path(source_dir)

    uv_project = detect_uv_project(source_dir)
    if uv_project is None:
        return False

    try:
        shutil.copy2(uv_project.uv_lock, dest_dir / _UV_LOCK_FILE)
        _logger.info(f"Copied {_UV_LOCK_FILE} to model artifacts")

        shutil.copy2(uv_project.pyproject, dest_dir / _PYPROJECT_FILE)
        _logger.info(f"Copied {_PYPROJECT_FILE} to model artifacts")

        python_version_src = source_dir / _PYTHON_VERSION_FILE
        if python_version_src.exists():
            shutil.copy2(python_version_src, dest_dir / _PYTHON_VERSION_FILE)
            _logger.info(f"Copied {_PYTHON_VERSION_FILE} to model artifacts")

        return True
    except Exception as e:
        _logger.warning(f"Failed to copy uv project files: {e}")
        return False


def extract_index_urls_from_uv_lock(uv_lock_path: str | Path) -> list[str]:
    """
    Extract private index URLs from a uv.lock file.

    Parses the uv.lock TOML file to find package sources that use non-PyPI registries.
    These URLs are returned as `--extra-index-url` compatible strings.

    Note: Credentials are NEVER stored in uv.lock. Users must provide credentials
    at restore time via UV_INDEX_* environment variables or .netrc file.

    Args:
        uv_lock_path: Path to the uv.lock file.

    Returns:
        A list of unique index URLs (excluding PyPI). Empty list if none found
        or if parsing fails.
    """
    uv_lock_path = Path(uv_lock_path)
    if not uv_lock_path.exists():
        return []

    try:
        content = uv_lock_path.read_text()

        # Extract registry URLs from uv.lock using regex
        # Format in uv.lock: source = { registry = "https://..." }
        registry_pattern = r'source\s*=\s*\{\s*registry\s*=\s*["\']([^"\']+)["\']'
        urls = set(re.findall(registry_pattern, content))

        # Filter out default PyPI URLs
        pypi_urls = {
            "https://pypi.org/simple",
            "https://pypi.org/simple/",
            "https://pypi.python.org/simple",
            "https://pypi.python.org/simple/",
        }
        private_urls = [url for url in urls if url.lower() not in pypi_urls]

        if private_urls:
            _logger.info(f"Extracted {len(private_urls)} private index URL(s) from uv.lock")
            _logger.warning(
                "Private package indexes detected in uv lockfile. "
                "Ensure credentials are available at model load time via "
                "UV_INDEX_* environment variables or .netrc file."
            )

        return sorted(private_urls)

    except Exception as e:
        _logger.debug(f"Failed to extract index URLs from uv.lock: {e}")
        return []


def create_uv_sync_pyproject(
    dest_dir: str | Path,
    python_version: str,
    project_name: str = "mlflow-model-env",
) -> Path:
    """
    Create a minimal pyproject.toml for uv sync.

    This is required for `uv sync --frozen` to work when restoring a model
    environment from a uv.lock artifact.

    Args:
        dest_dir: Directory where pyproject.toml will be created.
        python_version: Python version requirement (e.g., "3.11" or "3.11.5").
        project_name: Name for the temporary project. Defaults to "mlflow-model-env".

    Returns:
        Path to the created pyproject.toml file.
    """
    dest_dir = Path(dest_dir)

    # Pin the exact Python version including micro (e.g. "==3.11.5") so uv
    # restores the environment with the same interpreter that was used at
    # log time.  "dependencies" is left empty because all deps come from
    # uv.lock.
    pyproject_content = f"""[project]
name = "{project_name}"
version = "0.0.0"
requires-python = "=={python_version}"
dependencies = []

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
"""

    pyproject_path = dest_dir / _PYPROJECT_FILE
    pyproject_path.write_text(pyproject_content)
    _logger.debug(f"Created minimal pyproject.toml for uv sync at {pyproject_path}")

    return pyproject_path


def setup_uv_sync_environment(
    env_dir: str | Path,
    model_path: str | Path,
    python_version: str,
) -> bool:
    """
    Set up a uv project structure for environment restoration via ``uv sync --frozen``.

    Copies uv.lock from model artifacts and creates a minimal pyproject.toml
    to enable exact environment restoration using uv.

    Args:
        env_dir: The environment directory where the uv project will be set up.
        model_path: Path to the model artifacts containing uv.lock.
        python_version: Python version for the environment.

    Returns:
        True if setup succeeded (uv.lock exists in model artifacts), False otherwise.
    """
    env_dir = Path(env_dir)
    model_path = Path(model_path)

    uv_lock_artifact = model_path / _UV_LOCK_FILE
    if not uv_lock_artifact.exists():
        _logger.debug(f"No uv.lock found in model artifacts at {model_path}")
        return False

    env_dir.mkdir(parents=True, exist_ok=True)

    # Copy uv.lock to environment directory
    shutil.copy2(uv_lock_artifact, env_dir / _UV_LOCK_FILE)
    _logger.debug(f"Copied uv.lock to {env_dir}")

    # Copy pyproject.toml from model if available, otherwise create minimal one
    pyproject_artifact = model_path / _PYPROJECT_FILE
    if pyproject_artifact.exists():
        shutil.copy2(pyproject_artifact, env_dir / _PYPROJECT_FILE)
        _logger.debug(f"Copied pyproject.toml from model artifacts to {env_dir}")
    else:
        create_uv_sync_pyproject(env_dir, python_version)

    # Copy .python-version if it exists in model artifacts
    python_version_artifact = model_path / _PYTHON_VERSION_FILE
    if python_version_artifact.exists():
        shutil.copy2(python_version_artifact, env_dir / _PYTHON_VERSION_FILE)
        _logger.debug(f"Copied .python-version to {env_dir}")

    _logger.info(f"Set up uv sync environment at {env_dir}")
    return True


def run_uv_sync(
    project_dir: str | Path,
    frozen: bool = True,
    no_dev: bool = True,
    capture_output: bool = False,
) -> bool:
    """
    Run `uv sync` to install dependencies from a uv.lock file.

    This provides exact cross-platform environment restoration, significantly
    faster than pip-based installation.

    Args:
        project_dir: Directory containing pyproject.toml and uv.lock.
        frozen: Use frozen lockfile without updating. Defaults to True.
        no_dev: Exclude development dependencies. Defaults to True.
        capture_output: Whether to capture stdout/stderr. Defaults to False.

    Returns:
        True if sync succeeded, False otherwise.
    """
    uv_bin = _get_uv_binary()
    if uv_bin is None:
        _logger.warning("uv is not available for environment sync")
        return False
    project_dir = Path(project_dir)

    cmd = [uv_bin, "sync"]
    if frozen:
        cmd.append("--frozen")
    if no_dev:
        cmd.append("--no-dev")

    try:
        _logger.info(f"Running uv sync in {project_dir}")
        _logger.debug(f"uv sync command: {' '.join(str(c) for c in cmd)}")

        subprocess.run(
            cmd,
            cwd=project_dir,
            capture_output=capture_output,
            check=True,
            text=True,
            # Set UV_PROJECT_ENVIRONMENT to install packages into python environment at
            # `project_dir` path instead of the default `{cwd}/.venv` path.
            env={**os.environ, "UV_PROJECT_ENVIRONMENT": str(project_dir)},
        )

        _logger.info("uv sync completed successfully")
        return True

    except subprocess.CalledProcessError as e:
        _logger.warning(f"uv sync failed with exit code {e.returncode}: {e.stderr}")
        return False
    except Exception as e:
        _logger.warning(f"uv sync failed: {e}")
        return False


def has_uv_lock_artifact(model_path: str | Path) -> bool:
    """Check if a model has a uv.lock artifact."""
    return (Path(model_path) / _UV_LOCK_FILE).exists()