"""Auto-install support for per-tool optional dependencies."""

from __future__ import annotations

import importlib
import importlib.metadata as importlib_metadata
import importlib.util
import os
import shutil
import subprocess
import sys
import tomllib
from functools import cache
from pathlib import Path
from typing import TYPE_CHECKING

from mindroom.vendor_telemetry import vendor_telemetry_env_values

if TYPE_CHECKING:
    from mindroom.constants import RuntimePaths

_PACKAGE_NAME = "mindroom"
_RECEIPT_NAME = "uv-receipt.toml"
_PROJECT_ROOT = Path(__file__).resolve().parents[3]

# Packages where the pip install name differs from the Python import name.
# Only includes cases where replacing dashes with underscores is insufficient.
_PIP_TO_IMPORT: dict[str, str] = {
    "atlassian-python-api": "atlassian",
    "beautifulsoup4": "bs4",
    "e2b-code-interpreter": "e2b",
    "firecrawl-py": "firecrawl",
    "linkup-sdk": "linkup",
    "mem0ai": "mem0",
    "newspaper4k": "newspaper",
    "google-api-python-client": "googleapiclient",
    "google-auth": "google.auth",
    "google-cloud-bigquery": "google.cloud.bigquery",
    "google-genai": "google.genai",
    "google-maps-places": "google.maps",
    "google-search-results": "serpapi",
    "psycopg-binary": "psycopg",
    "pyobjc-framework-applicationservices": "ApplicationServices",
    "pyobjc-framework-cocoa": "AppKit",
    "py-trello": "trello",
    "pygithub": "github",
    "pyyaml": "yaml",
    "tavily-python": "tavily",
    "spider-client": "spider",
}


def _pip_name_to_import(pip_name: str) -> str:
    """Convert a pip package name to its top-level import module name."""
    normalized = pip_name.strip().lower().replace("_", "-")
    # Strip version specifiers
    for sep in (">=", "<=", "==", ">", "<", "~=", "!="):
        if sep in normalized:
            normalized = normalized.split(sep, 1)[0].strip()
            break
    if normalized in _PIP_TO_IMPORT:
        return _PIP_TO_IMPORT[normalized]
    return normalized.replace("-", "_")


def _normalize_extra_name(extra_name: str) -> str:
    """Normalize extra names across pyproject and installed metadata conventions."""
    return extra_name.strip().lower().replace("_", "-")


def check_deps_installed(dependencies: list[str]) -> bool:
    """Check if all dependencies are importable using find_spec (no side effects)."""
    for dep in dependencies:
        module_name = _pip_name_to_import(dep)
        if importlib.util.find_spec(module_name) is None:
            return False
    return True


def auto_install_enabled(runtime_paths: RuntimePaths) -> bool:
    """Return whether automatic tool dependency installation is enabled."""
    raw = runtime_paths.env_value("MINDROOM_NO_AUTO_INSTALL_TOOLS", default="") or ""
    return raw.lower() not in {"1", "true", "yes"}


def _has_lockfile() -> bool:
    """Check if uv.lock is available alongside pyproject.toml."""
    return (_PROJECT_ROOT / "uv.lock").exists()


@cache
def _available_optional_extras() -> set[str]:
    """Discover available optional extras from pyproject or installed metadata."""
    pyproject_path = _PROJECT_ROOT / "pyproject.toml"
    if pyproject_path.exists():
        data = tomllib.loads(pyproject_path.read_text(encoding="utf-8"))
        optional = data.get("project", {}).get("optional-dependencies", {})
        return set(optional.keys())

    try:
        metadata = importlib_metadata.metadata(_PACKAGE_NAME)
    except importlib_metadata.PackageNotFoundError:
        return set()
    return set(metadata.get_all("Provides-Extra") or [])


def _resolve_optional_extra_name(extra_name: str) -> str | None:
    """Resolve an extra name against available extras with normalized matching."""
    available = _available_optional_extras()
    if extra_name in available:
        return extra_name

    normalized_name = _normalize_extra_name(extra_name)
    for available_name in sorted(available):
        if _normalize_extra_name(available_name) == normalized_name:
            return available_name
    return None


def _is_uv_tool_install() -> bool:
    """Check if running from a uv tool environment."""
    return (Path(sys.prefix) / _RECEIPT_NAME).exists()


def _in_virtualenv() -> bool:
    return sys.prefix != sys.base_prefix


def _get_current_uv_tool_extras() -> list[str]:
    receipt = Path(sys.prefix) / _RECEIPT_NAME
    if not receipt.exists():
        return []
    data = tomllib.loads(receipt.read_text(encoding="utf-8"))
    requirements = data.get("tool", {}).get("requirements", [])
    for requirement in requirements:
        if requirement.get("name") == _PACKAGE_NAME:
            return requirement.get("extras", [])
    return []


def _install_via_uv_tool(extras: list[str], *, quiet: bool) -> bool:
    extras_str = ",".join(extras)
    package_spec = f"{_PACKAGE_NAME}[{extras_str}]"
    major, minor = sys.version_info[:2]
    python_version = f"{major}.{minor}"
    cmd = ["uv", "tool", "install", package_spec, "--force", "--python", python_version]
    if quiet:
        cmd.append("-q")
    env = os.environ.copy()
    env.update(vendor_telemetry_env_values())
    result = subprocess.run(cmd, check=False, env=env)
    return result.returncode == 0


def _current_python_has_module(module_name: str) -> bool:
    """Return whether the active interpreter can import a module."""
    return importlib.util.find_spec(module_name) is not None


def install_command_for_current_python() -> list[str]:
    """Build the pip/uv install command for the current interpreter."""
    in_venv = _in_virtualenv()
    if _current_python_has_module("uv"):
        cmd = [sys.executable, "-m", "uv", "pip", "install", "--python", sys.executable]
    elif shutil.which("uv"):
        cmd = ["uv", "pip", "install", "--python", sys.executable]
    else:
        cmd = [sys.executable, "-m", "pip", "install"]
        if not in_venv:
            cmd.append("--user")
        return cmd
    if not in_venv:
        cmd.append("--system")
    return cmd


def _install_via_uv_sync(extras: list[str], *, quiet: bool) -> bool:
    """Install extras using ``uv sync --locked --inexact`` for pinned versions from uv.lock."""
    cmd = ["uv", "sync", "--locked", "--inexact", "--no-dev"]
    env = os.environ.copy()
    env.update(vendor_telemetry_env_values())
    if _in_virtualenv():
        # Ensure uv targets the interpreter that is currently running MindRoom.
        cmd.extend(["--active", "--python", sys.executable])
        env["VIRTUAL_ENV"] = sys.prefix
    for extra in extras:
        cmd.extend(["--extra", extra])
    if quiet:
        cmd.append("-q")
    result = subprocess.run(cmd, check=False, capture_output=quiet, cwd=_PROJECT_ROOT, env=env)
    return result.returncode == 0


def _install_in_environment(extras: list[str], *, quiet: bool) -> bool:
    extras_str = ",".join(extras)
    package_spec = f"{_PACKAGE_NAME}[{extras_str}]"
    cmd = [*install_command_for_current_python(), package_spec]
    env = os.environ.copy()
    env.update(vendor_telemetry_env_values())
    result = subprocess.run(cmd, check=False, capture_output=quiet, env=env)
    return result.returncode == 0


def _install_optional_extras(extras: list[str], *, quiet: bool = False) -> bool:
    """Install one or more optional extras into the current environment.

    Prefers ``uv sync --locked`` when uv.lock is available (exact pinned versions).
    Falls back to ``uv pip install`` or ``pip install`` otherwise.
    """
    if not extras:
        return False
    if _is_uv_tool_install():
        current_extras = _get_current_uv_tool_extras()
        merged_by_name = {_normalize_extra_name(extra): extra for extra in current_extras}
        merged_by_name.update({_normalize_extra_name(extra): extra for extra in extras})
        merged = sorted(merged_by_name.values())
        return _install_via_uv_tool(merged, quiet=quiet)
    if _has_lockfile() and shutil.which("uv") and _in_virtualenv():
        return _install_via_uv_sync(extras, quiet=quiet)
    return _install_in_environment(extras, quiet=quiet)


def _auto_install_optional_extra(extra_name: str, runtime_paths: RuntimePaths) -> bool:
    """Auto-install an optional extra when supported and enabled."""
    if not auto_install_enabled(runtime_paths):
        return False
    resolved_extra_name = _resolve_optional_extra_name(extra_name)
    if resolved_extra_name is None:
        return False
    return _install_optional_extras([resolved_extra_name], quiet=True)


def auto_install_optional_extra_for_import_retry(extra_name: str, runtime_paths: RuntimePaths) -> bool:
    """Auto-install an optional extra and invalidate import caches before a retry."""
    installed = _auto_install_optional_extra(extra_name, runtime_paths)
    if installed:
        importlib.invalidate_caches()
    return installed


def ensure_optional_deps(
    dependencies: list[str],
    extra_name: str,
    runtime_paths: RuntimePaths,
    *,
    missing_message: str | None = None,
) -> bool:
    """Ensure dependencies are installed, auto-installing via optional extra if needed."""
    if check_deps_installed(dependencies):
        return False
    if not auto_install_optional_extra_for_import_retry(extra_name, runtime_paths):
        missing = ", ".join(dependencies)
        if missing_message is None:
            missing_message = f"Missing dependencies: {missing}. Install with: pip install 'mindroom[{extra_name}]'"
        raise ImportError(missing_message)
    return True


def ensure_tool_deps(
    dependencies: list[str],
    tool_extra: str,
    runtime_paths: RuntimePaths,
    *,
    missing_message: str | None = None,
) -> bool:
    """Ensure tool dependencies are installed, auto-installing via tool extra if needed."""
    return ensure_optional_deps(dependencies, tool_extra, runtime_paths, missing_message=missing_message)
