#!/usr/bin/env python
"""
Automated installer for SuiteSparse/UMFPACK + scikit-umfpack on Windows.

Key safety features
-------------------
1.   **Explicit confirmation** for every action that modifies the system.
2.   **Dry-run preview**: run with `--dry-run` to see what *would* happen.
3.   Detects existing Conda envs and asks before overwriting.
4.   All shell commands are run with `conda run -n <env> …` to avoid
     brittle `conda activate` tricks and to isolate effects.
5.   Catches and reports subprocess errors without crashing silently.
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import List, Optional

# --------------------------------------------------------------------------- #
# Utility helpers
# --------------------------------------------------------------------------- #

CONTINUE = input('WARNING!! This script is auto generated by ChatGPT. It is UNTESTED and being verified' + 
                    'Continue at your own risk! Do you want to proceed? [y/n]')
if CONTINUE != 'y':
    print('Aborting')
    quit()

def ask_yes_no(prompt: str, default: str | None = "y") -> bool:
    """Prompt the user with a yes/no question."""
    if default not in {"y", "n", None}:
        raise ValueError("default must be 'y', 'n', or None")
    suffix = " [Y/n] " if default == "y" else " [y/N] " if default == "n" else " [y/n] "
    while True:
        ans = input(prompt + suffix).strip().lower()
        if not ans and default:
            ans = default
        if ans in {"y", "yes"}:
            return True
        if ans in {"n", "no"}:
            return False
        print("Please enter y or n.")

def run(cmd: str | List[str], *, dry: bool = False, cwd: str | Path | None = None) -> None:
    """Run a shell command with live output, honor dry-run mode."""
    print(f"\n>>> {cmd}")
    if dry:
        return
    try:
        subprocess.run(cmd, shell=isinstance(cmd, str), check=True, cwd=cwd)
    except subprocess.CalledProcessError as exc:
        print(f"❌ Command failed (exit {exc.returncode}). Aborting.")
        sys.exit(exc.returncode)

def find_conda_root() -> Optional[Path]:
    """Guess common Conda installation directories."""
    guesses = [
        Path(os.environ.get("CONDA_PREFIX", "")),                                 # active env
        Path(os.environ.get("USERPROFILE", "")) / "Miniconda3",
        Path(os.environ.get("USERPROFILE", "")) / "Anaconda3",
    ]
    for p in guesses:
        if (p / "condabin" / "conda.bat").exists():
            return p
    return None

# --------------------------------------------------------------------------- #
# Installer logic
# --------------------------------------------------------------------------- #

def write_nativefile(conda_root: Path, *, dry: bool = False) -> Path:
    """Create a nativefile.ini pointing to SuiteSparse headers/libs in Conda."""
    content = f"""
[binaries]
c = 'x86_64-w64-mingw32-gcc'
cpp = 'x86_64-w64-mingw32-g++'

[properties]
umfpack-libdir = '''{conda_root.as_posix()}/Library/lib'''
umfpack-includedir = '''{conda_root.as_posix()}/Library/include/suitesparse'''
""".lstrip()

    nativefile = Path.cwd() / "nativefile.ini"
    print(f"\n📝 Creating {nativefile} …")
    if dry:
        print(content)
    else:
        nativefile.write_text(content, encoding="utf-8")
    return nativefile

def conda_run(env: str, args: str) -> str:
    """Helper to prepend 'conda run -n <env>'."""
    return f"conda run -n {env} {args}"

def main() -> None:
    parser = argparse.ArgumentParser(description="Safe UMFPACK installer")
    parser.add_argument(
        "--env",
        default="umf-env",
        help="Name of the Conda environment to create/use (default: %(default)s)",
    )
    parser.add_argument("--python", default="3.10", help="Python version for the env")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show commands but do not execute them",
    )
    args = parser.parse_args()
    dry = args.dry_run
    env_name = args.env
    python_ver = args.python

    print(
        "🛠️  UMFPACK/scikit-umfpack automated installer\n"
        "--------------------------------------------------------"
    )

    # 1. Locate Conda
    conda_root = find_conda_root()
    if conda_root:
        print(f"✅ Found Conda at {conda_root}")
    else:
        print("⚠️  Could not auto-detect Conda. Make sure it is on PATH.")
        conda_path_str = input("Paste the full path to your Conda root (e.g. C:/Miniconda3): ").strip()
        conda_root = Path(conda_path_str)
        if not (conda_root / "condabin" / "conda.bat").exists():
            print("❌ The provided path does not look like a Conda installation. Aborting.")
            sys.exit(1)

    # 2. Create or reuse environment
    print(f"\n🔍 Checking if environment '{env_name}' exists …")
    envs_json = subprocess.check_output(["conda", "env", "list", "--json"], text=True)
    envs = json.loads(envs_json)["envs"]
    env_path = next((Path(p) for p in envs if Path(p).name == env_name), None)

    if env_path:
        print(f"ℹ️  Conda env '{env_name}' already exists at {env_path}.")
        if not ask_yes_no("Reuse this environment? Otherwise I can recreate it", default="y"):
            if ask_yes_no(f"Delete existing env '{env_name}' and recreate it?", default="n"):
                run(f"conda remove -n {env_name} --all -y", dry=dry)
            else:
                print("Aborting as per user request.")
                return

    if not env_path:
        if ask_yes_no(f"Create new environment '{env_name}' with Python {python_ver}?", default="y"):
            run(f"conda create -n {env_name} python={python_ver} -y", dry=dry)
        else:
            print("Aborting at user request.")
            return

    # 3. Install core dependencies
    print("\n📦 Installing Conda packages (this may take a while) …")
    conda_pkgs = [
        "anaconda::suitesparse",
        "meson",
        "swig",
        "-c conda-forge compilers",
        "-c conda-forge m2w64-toolchain",
        "-c conda-forge openblas",
    ]
    run(conda_run(env_name, f"conda install {' '.join(conda_pkgs)} -y"), dry=dry)

    # 4. Install pip dependencies
    print("\n📦 Installing pip packages …")
    pip_cmds = [
        "python -m pip install --upgrade pip",
        "python -m pip install meson-python",
        "python -m pip install numpy scipy",
    ]
    for cmd in pip_cmds:
        run(conda_run(env_name, cmd), dry=dry)

    # 5. Write nativefile.ini
    nativefile_path = write_nativefile(conda_root, dry=dry)

    # 6. Build & install scikit-umfpack
    print("\n🔧 Building scikit-umfpack from source …")
    build_cmd = (
        "python -m pip install scikit-umfpack "
        "--no-build-isolation "
        f"-Csetup-args=\"--native-file={nativefile_path}\""
    )
    run(conda_run(env_name, build_cmd), dry=dry)

    # 7. Optional extras
    extras = [
        "python -m pip install numba numba-progress --no-binary :all:",
        "python -m pip install emerge --no-deps",
    ]
    if ask_yes_no("Install optional extras (numba, numba-progress, emerge)?", default="y"):
        for cmd in extras:
            run(conda_run(env_name, cmd), dry=dry)

    print(
        "\n🎉 Installation finished.\n"
        "Test it with:\n"
        f"   conda run -n {env_name} python - <<'PY'\n"
        "import scikits.umfpack, sys; print('UMFPACK OK on', sys.platform)\n"
        "PY"
    )

# --------------------------------------------------------------------------- #

if __name__ == "__main__":
    main()
