"""Installation API for comfy-env."""

import inspect
import os
from pathlib import Path
from typing import Callable, List, Optional, Set, Union

from .config import ComfyEnvConfig, NodeDependency, load_config, discover_config, CONFIG_FILE_NAME, ROOT_CONFIG_FILE_NAME

USE_COMFY_ENV_VAR = "USE_COMFY_ENV"


def _is_comfy_env_enabled() -> bool:
    return os.environ.get(USE_COMFY_ENV_VAR, "1").lower() not in ("0", "false", "no", "off")


def install(
    config: Optional[Union[str, Path]] = None,
    node_dir: Optional[Path] = None,
    log_callback: Optional[Callable[[str], None]] = None,
    dry_run: bool = False,
) -> bool:
    """Install dependencies from comfy-env-root.toml or comfy-env.toml."""
    if node_dir is None:
        node_dir = Path(inspect.stack()[1].filename).parent.resolve()

    log = log_callback or print

    if config is not None:
        config_path = Path(config)
        if not config_path.is_absolute():
            config_path = node_dir / config_path
        cfg = load_config(config_path)
    else:
        cfg = discover_config(node_dir, root=True)

    if cfg is None:
        raise FileNotFoundError(f"No {ROOT_CONFIG_FILE_NAME} or {CONFIG_FILE_NAME} found in {node_dir}")

    if cfg.apt_packages: _install_apt_packages(cfg.apt_packages, log, dry_run)
    if cfg.env_vars: _set_persistent_env_vars(cfg.env_vars, log, dry_run)
    if cfg.node_reqs:
        _install_node_dependencies(cfg.node_reqs, node_dir, log, dry_run)
        _reinstall_main_requirements(node_dir, log, dry_run)

    if _is_comfy_env_enabled():
        _install_via_pixi(cfg, node_dir, log, dry_run)
        _install_isolated_subdirs(node_dir, log, dry_run)
    else:
        log("\n[comfy-env] Isolation disabled (USE_COMFY_ENV=0)")
        _install_to_host_python(cfg, node_dir, log, dry_run)

    log("\nInstallation complete!")
    return True


def _install_apt_packages(packages: List[str], log: Callable[[str], None], dry_run: bool) -> None:
    from .packages.apt import apt_install
    import platform
    if platform.system() != "Linux":
        return
    log(f"\n[apt] Installing: {', '.join(packages)}")
    if not dry_run:
        success = apt_install(packages, log)
        if not success:
            log("[apt] WARNING: Some apt packages failed to install. This may cause issues.")


def _set_persistent_env_vars(env_vars: dict, log: Callable[[str], None], dry_run: bool) -> None:
    import platform, subprocess
    if not env_vars: return

    log(f"\n[env] Setting {len(env_vars)} env var(s)")
    for k, v in env_vars.items(): log(f"  {k}={v}")
    if dry_run: return

    system = platform.system()
    if system == "Windows":
        for k, v in env_vars.items():
            subprocess.run(["setx", k, v], capture_output=True)
    elif system == "Darwin":
        for k, v in env_vars.items():
            subprocess.run(["launchctl", "setenv", k, v], capture_output=True)
        _add_to_shell_profile(env_vars, log)
    else:
        _add_to_shell_profile(env_vars, log)


def _add_to_shell_profile(env_vars: dict, log: Callable[[str], None]) -> None:
    shell = os.environ.get("SHELL", "/bin/bash")
    rc_file = Path.home() / (".zshrc" if "zsh" in shell else ".bashrc")
    profile_file = Path.home() / ".comfy-env-profile"

    with open(profile_file, "w") as f:
        f.write("# Generated by comfy-env\n")
        for k, v in env_vars.items():
            f.write(f'export {k}="{v}"\n')

    source_line = f'source "{profile_file}"'
    existing = rc_file.read_text() if rc_file.exists() else ""
    if source_line not in existing:
        with open(rc_file, "a") as f:
            f.write(f'\n# comfy-env\n{source_line}\n')
    log(f"  [env] Wrote {profile_file}")


def _install_node_dependencies(node_reqs: List[NodeDependency], node_dir: Path, log: Callable[[str], None], dry_run: bool) -> None:
    from .packages.node_dependencies import install_node_dependencies
    custom_nodes_dir = node_dir.parent
    log(f"\nInstalling {len(node_reqs)} node dependencies...")
    if dry_run:
        for req in node_reqs:
            log(f"  {req.name}: {'exists' if (custom_nodes_dir / req.name).exists() else 'would clone'}")
        return
    install_node_dependencies(node_reqs, custom_nodes_dir, log, {node_dir.name})


def _reinstall_main_requirements(node_dir: Path, log: Callable[[str], None], dry_run: bool) -> None:
    """Re-install main package's requirements.txt after node_reqs to restore correct versions."""
    from .packages.node_dependencies import install_requirements
    req_file = node_dir / "requirements.txt"
    if not req_file.exists():
        return
    log(f"\n[requirements] Re-installing main package requirements...")
    if not dry_run:
        install_requirements(node_dir, log)


def _has_isolated_subdirs(node_dir: Path) -> bool:
    """Check if there are any comfy-env.toml files in subdirectories."""
    for config_file in node_dir.rglob(CONFIG_FILE_NAME):
        if config_file.parent != node_dir:
            return True
    return False


def _install_via_pixi(cfg: ComfyEnvConfig, node_dir: Path, log: Callable[[str], None], dry_run: bool, is_root: bool = True) -> None:
    from .packages.pixi import ensure_pixi
    from .packages.toml_generator import write_pixi_toml
    from .packages.cuda_wheels import get_wheel_url, CUDA_TORCH_MAP
    from .detection import get_recommended_cuda_version
    from .environment.cache import get_cache_dir, get_junction_name, create_junction, get_env_name, cleanup_orphaned_envs
    import shutil, subprocess, sys

    # Clean up orphaned environments before installing
    cleanup_orphaned_envs(log)

    deps = cfg.pixi_passthrough.get("dependencies", {})
    pypi_deps = cfg.pixi_passthrough.get("pypi-dependencies", {})
    if not cfg.cuda_packages and not deps and not pypi_deps:
        # Only log "no packages" if this is root AND there are no isolated subdirs
        if not is_root or not _has_isolated_subdirs(node_dir):
            log("No packages to install")
        return

    log(f"\nInstalling via pixi:")
    if cfg.cuda_packages: log(f"  CUDA: {', '.join(cfg.cuda_packages)}")
    if deps: log(f"  Conda: {len(deps)}")
    if pypi_deps: log(f"  PyPI: {len(pypi_deps)}")
    if dry_run: return

    # Find config file path
    config_path = node_dir / CONFIG_FILE_NAME
    if not config_path.exists():
        config_path = node_dir / ROOT_CONFIG_FILE_NAME

    # Find main node dir (for naming)
    main_node_dir = node_dir
    for parent in node_dir.parents:
        if parent.parent.name == "custom_nodes":
            main_node_dir = parent
            break

    # Central env path and local junction
    central_env = get_cache_dir() / get_env_name(main_node_dir, config_path)
    junction_name = get_junction_name(config_path)
    junction_path = node_dir / junction_name

    # Create env at short central path from the start
    env_work_dir = central_env.parent / f"{central_env.name}_build"
    log(f"[DEBUG] env_work_dir={env_work_dir}")
    log(f"[DEBUG] junction={junction_path}")

    # Clean up any previous build
    if env_work_dir.exists():
        shutil.rmtree(env_work_dir)
    env_work_dir.mkdir(parents=True, exist_ok=True)

    log("[DEBUG] ensure_pixi...")
    pixi_path = ensure_pixi(log=log)
    log(f"[DEBUG] pixi_path={pixi_path}")

    cuda_version = torch_version = None
    if cfg.has_cuda and sys.platform != "darwin":
        log("[DEBUG] get_recommended_cuda_version...")
        cuda_version = get_recommended_cuda_version()
        log(f"[DEBUG] cuda_version={cuda_version}")
        if cuda_version:
            torch_version = CUDA_TORCH_MAP.get(".".join(cuda_version.split(".")[:2]), "2.8")
            log(f"[DEBUG] torch_version={torch_version}")

    # Write pixi.toml to short path and run pixi there
    log("[DEBUG] write_pixi_toml...")
    write_pixi_toml(cfg, env_work_dir, log)
    log("Running pixi install...")
    result = subprocess.run([str(pixi_path), "install"], cwd=env_work_dir, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"pixi install failed:\nstderr: {result.stderr}\nstdout: {result.stdout}")

    if cfg.cuda_packages:
        # Get Python from the build dir
        pixi_env = env_work_dir / ".pixi" / "envs" / "default"
        python_path = pixi_env / ("python.exe" if sys.platform == "win32" else "bin/python")
        if not python_path.exists():
            raise RuntimeError(f"No Python in pixi env: {python_path}")

        result = subprocess.run([str(python_path), "-c", "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"],
                               capture_output=True, text=True)
        py_version = result.stdout.strip() if result.returncode == 0 else f"{sys.version_info.major}.{sys.version_info.minor}"

        # PyTorch packages install from official PyTorch index
        pytorch_packages = {"torch", "torchvision", "torchaudio"}
        # torchvision has different versioning scheme
        torchvision_map = {"2.8": "0.23", "2.4": "0.19"}

        if cuda_version:
            # GPU mode - use CUDA index
            pytorch_index = f"https://download.pytorch.org/whl/cu{cuda_version.replace('.', '')[:3]}"
            pin_torch_version = torch_version
            log(f"Installing CUDA packages from {pytorch_index}")
        else:
            # CPU mode - use CPU index
            pytorch_index = "https://download.pytorch.org/whl/cpu"
            pin_torch_version = "2.8"  # Default torch version for CPU
            log(f"Installing CPU packages from {pytorch_index}")

        for package in cfg.cuda_packages:
            if package in pytorch_packages:
                # Install from PyTorch index with version pinning
                if package == "torch":
                    pin_version = pin_torch_version
                elif package == "torchvision":
                    pin_version = torchvision_map.get(pin_torch_version, "0.23")
                else:  # torchaudio follows torch versioning
                    pin_version = pin_torch_version
                pkg_spec = f"{package}=={pin_version}.*"
                pip_cmd = [str(python_path), "-m", "pip", "install", "--no-cache-dir",
                          "--extra-index-url", pytorch_index, pkg_spec]
                log(f"  {' '.join(pip_cmd)}")
                result = subprocess.run(pip_cmd, capture_output=True, text=True)
                if result.returncode != 0:
                    raise RuntimeError(f"Failed to install {package}:\nstderr: {result.stderr}\nstdout: {result.stdout}")
            elif cuda_version:
                # Install GPU-only packages from cuda-wheels (skip on CPU)
                wheel_url = get_wheel_url(package, torch_version, cuda_version, py_version)
                if not wheel_url:
                    raise RuntimeError(f"No wheel for {package}")
                log(f"  {package} from {wheel_url}")
                result = subprocess.run([str(python_path), "-m", "pip", "install", "--no-deps", "--no-cache-dir", wheel_url],
                                       capture_output=True, text=True)
                if result.returncode != 0:
                    raise RuntimeError(f"Failed to install {package}:\nstderr: {result.stderr}\nstdout: {result.stdout}")
            else:
                # CPU mode - skip GPU-only packages
                log(f"  {package} (skipped - GPU only)")

    # Move env from build dir to final location and create junction
    pixi_env = env_work_dir / ".pixi" / "envs" / "default"
    if pixi_env.exists():
        if central_env.exists():
            shutil.rmtree(central_env)
        shutil.move(str(pixi_env), str(central_env))
        create_junction(junction_path, central_env)
        shutil.rmtree(env_work_dir, ignore_errors=True)
        # Clean up any old .pixi in node_dir
        shutil.rmtree(node_dir / ".pixi", ignore_errors=True)
        log(f"Env: {central_env}")
        log(f"Junction: {junction_path} -> {central_env}")


def _install_to_host_python(cfg: ComfyEnvConfig, node_dir: Path, log: Callable[[str], None], dry_run: bool) -> None:
    import shutil, subprocess, sys
    from .packages.cuda_wheels import get_wheel_url, CUDA_TORCH_MAP
    from .detection import get_recommended_cuda_version

    pypi_deps = cfg.pixi_passthrough.get("pypi-dependencies", {})
    if not pypi_deps and not cfg.cuda_packages:
        log("No packages to install")
        return

    pip_packages = []
    for pkg, spec in pypi_deps.items():
        if isinstance(spec, str):
            pip_packages.append(pkg if spec == "*" else f"{pkg}{spec}")
        elif isinstance(spec, dict):
            extras = spec.get("extras", [])
            version = spec.get("version", "*")
            name = f"{pkg}[{','.join(extras)}]" if extras else pkg
            pip_packages.append(name if version == "*" else f"{name}{version}")

    log(f"\nInstalling to {sys.executable}")
    if dry_run: return

    use_uv = shutil.which("uv")
    if pip_packages:
        cmd = ["uv", "pip", "install", "--python", sys.executable] + pip_packages if use_uv else [sys.executable, "-m", "pip", "install"] + pip_packages
        subprocess.run(cmd, capture_output=True)

    if cfg.cuda_packages:
        cuda_version = get_recommended_cuda_version()
        if not cuda_version: return
        torch_version = CUDA_TORCH_MAP.get(".".join(cuda_version.split(".")[:2]), "2.8")
        py_version = f"{sys.version_info.major}.{sys.version_info.minor}"

        for package in cfg.cuda_packages:
            wheel_url = get_wheel_url(package, torch_version, cuda_version, py_version)
            if wheel_url:
                cmd = ["uv", "pip", "install", "--python", sys.executable, "--no-deps", wheel_url] if use_uv else [sys.executable, "-m", "pip", "install", "--no-deps", wheel_url]
                subprocess.run(cmd, capture_output=True)


def _install_isolated_subdirs(node_dir: Path, log: Callable[[str], None], dry_run: bool) -> None:
    """Find and install comfy-env.toml in subdirectories (isolated folders only)."""
    for config_file in node_dir.rglob(CONFIG_FILE_NAME):
        if config_file.parent == node_dir: continue  # Skip root
        log(f"\n[isolated] {config_file.parent.relative_to(node_dir)}")
        if not dry_run:
            _install_via_pixi(load_config(config_file), config_file.parent, log, dry_run, is_root=False)


def verify_installation(packages: List[str], log: Callable[[str], None] = print) -> bool:
    all_ok = True
    for package in packages:
        import_name = package.replace("-", "_").split("[")[0]
        try:
            __import__(import_name)
            log(f"  {package}: OK")
        except ImportError as e:
            log(f"  {package}: FAILED ({e})")
            all_ok = False
    return all_ok
