#!/usr/bin/env python3
"""Environment manager for EE 508 / DS 537.

Download this file together with ``requirements.yml`` into your project folder,
then run:

    python dev.py setup

This creates a conda/mamba environment with everything the course needs,
verifies the install, and (optionally) adds a one-word command to launch
Jupyter. Other commands: update, clean, verify, launcher, list.
"""

import os
import shutil
import subprocess
import sys
from pathlib import Path

DEFAULT_ENV_NAME = 'ee508'
REQUIREMENTS = 'requirements.yml'

# Packages the labs import directly — checked by `verify`.
VERIFY_MODULES = [
    'numpy',
    'pandas',
    'geopandas',
    'shapely',
    'pyproj',
    'rasterio',
    'rioxarray',
    'rasterstats',
    'exactextract',
    'statsmodels',
    'sklearn',
    'xgboost',
    'lightgbm',
    'catboost',
    'matplotlib',
    'folium',
]


def get_package_manager():
    """Return the full path to mamba if available, otherwise conda."""
    # Full path avoids picking up shell functions that shadow the executable.
    mamba_path = shutil.which('mamba')
    conda_path = shutil.which('conda')

    if mamba_path:
        return mamba_path
    elif conda_path:
        return conda_path
    else:
        print('✗ Error: neither mamba nor conda found in PATH')
        print('\nInstall one of the following first, then re-run this script:')
        print('  - Miniforge (includes mamba): '
              'https://github.com/conda-forge/miniforge')
        print('  - Miniconda (conda): '
              'https://docs.conda.io/en/latest/miniconda.html')
        sys.exit(1)


PKG_MGR = get_package_manager()


def get_env_name():
    """Prompt for an environment name (default: ee508)."""
    env_name = input(
        f'Environment name (press Enter for "{DEFAULT_ENV_NAME}"): '
    ).strip()
    return env_name if env_name else DEFAULT_ENV_NAME


def run(cmd, check=True):
    """Run a shell command, printing it first."""
    print(f'→ {cmd}')
    result = subprocess.run(cmd, shell=True)
    if check and result.returncode != 0:
        print(f'✗ Command failed with code {result.returncode}')
        sys.exit(1)
    return result.returncode == 0


def _require_requirements_file():
    """Exit with a helpful message if requirements.yml is not next to dev.py."""
    req = Path(__file__).parent / REQUIREMENTS
    if not req.exists():
        print(f'✗ {REQUIREMENTS} not found next to dev.py ({req.parent}).')
        print('  Download requirements.yml into the same folder and try again.')
        sys.exit(1)
    return req


def verify(env_name=None):
    """Import the course packages inside the env and report ✓/✗ per module."""
    env_name = env_name or get_env_name()
    print(f'\nVerifying packages in environment: {env_name}')

    modules = ','.join(f'"{m}"' for m in VERIFY_MODULES)
    checker = (
        'import importlib\n'
        f'mods = [{modules}]\n'
        'failed = []\n'
        'for m in mods:\n'
        '    try:\n'
        '        importlib.import_module(m)\n'
        '        print("  \\u2713", m)\n'
        '    except Exception as e:\n'
        '        print("  \\u2717", m, "-", e)\n'
        '        failed.append(m)\n'
        'import geopandas\n'
        'print("\\ngeopandas", geopandas.__version__)\n'
        'raise SystemExit(1 if failed else 0)\n'
    )
    # Pass the script via stdin to avoid shell-quoting headaches.
    proc = subprocess.run(
        f'{PKG_MGR} run -n {env_name} python -',
        input=checker,
        shell=True,
        text=True,
    )
    if proc.returncode == 0:
        print('✓ All course packages import cleanly.')
        return True
    else:
        print('✗ Some packages failed to import (see above).')
        return False


def install_launcher(env_name):
    """Add a one-word command that activates the env and launches Jupyter."""
    notebooks_dir = Path(__file__).parent.resolve() / 'notebooks'
    notebooks_dir.mkdir(parents=True, exist_ok=True)
    sentinel = f'# ee508-launcher:{env_name}'

    if sys.platform == 'win32':
        bat_path = Path.home() / f'{env_name}.bat'
        if bat_path.exists() and sentinel in bat_path.read_text():
            print(f'✓ Launcher already installed: {bat_path}')
            return
        content = (
            f'@echo off\n'
            f'REM {sentinel}\n'
            f'call conda activate {env_name}\n'
            f'cd /d "{notebooks_dir}"\n'
            f'jupyter notebook\n'
        )
        try:
            bat_path.write_text(content)
            print(f'✓ Launcher created: {bat_path}')
            print('  Open an Anaconda Prompt and type '
                  f'`{env_name}` to launch Jupyter.')
        except OSError as e:
            print(f'✗ Could not write launcher: {e}')

    else:
        func = (
            f'\n{sentinel}\n'
            f'{env_name}() {{\n'
            f'  cd "{notebooks_dir}" || return\n'
            f'  conda activate {env_name}\n'
            f'  jupyter notebook\n'
            f'}}\n'
        )
        if sys.platform == 'darwin':
            rc_file = Path('~/.zshrc').expanduser()
        else:
            shell = os.path.basename(os.environ.get('SHELL', 'bash'))
            if shell == 'zsh':
                rc_file = Path('~/.zshrc').expanduser()
            elif shell == 'bash':
                rc_file = Path('~/.bashrc').expanduser()
            else:
                print(f'✗ Shell {shell!r} not supported for automatic setup.')
                print('  Add this function to your shell rc file manually:')
                print(func)
                return

        if rc_file.exists() and sentinel in rc_file.read_text():
            print(f'✓ Launcher already installed in {rc_file}')
            return
        try:
            with rc_file.open('a') as f:
                f.write(func)
            print(f'✓ Shell function `{env_name}()` added to {rc_file}')
            print(f'  Open a new terminal (or `source {rc_file}`), then type '
                  f'`{env_name}`.')
        except OSError as e:
            print(f'✗ Could not write to {rc_file}: {e}')


def setup():
    """Create the environment from requirements.yml, verify, add a launcher."""
    _require_requirements_file()
    pkg_mgr = PKG_MGR.split(os.sep)[-1]

    print('This will create the EE 508 course environment.\n')
    print(f'Found package manager: {PKG_MGR}\n')

    env_name = get_env_name()
    launcher_response = input(
        f'\nInstall a `{env_name}` command to launch Jupyter from the terminal? '
        '[Y/n] '
    ).strip().lower()

    print(f'\nUsing package manager: {pkg_mgr}')
    print(f'Creating environment: {env_name}')
    print('(this downloads several GB and can take a while)\n')

    if not run(f'{PKG_MGR} env create -f {REQUIREMENTS} -n {env_name} -y'):
        print('✗ Failed to create environment')
        return

    print('\nVerifying the installation...')
    verify(env_name)

    if launcher_response in ('', 'y'):
        install_launcher(env_name)

    print('\n✓ Environment ready!')
    print('\nNext steps:')
    if launcher_response in ('', 'y'):
        if sys.platform == 'win32':
            print(f'  1. Open an Anaconda Prompt and type `{env_name}`.')
        else:
            print(f'  1. Open a new terminal and type `{env_name}`.')
        print('  2. Start coding!')
    else:
        print(f'  1. conda activate {env_name}')
        print('  2. cd notebooks')
        print('  3. jupyter notebook')


def update():
    """Update an existing environment from requirements.yml, then verify."""
    _require_requirements_file()
    env_name = get_env_name()

    print(f'\nUsing package manager: {PKG_MGR}')
    print(f'Updating environment: {env_name}')
    run(f'{PKG_MGR} env update -f {REQUIREMENTS} -n {env_name} --prune')

    print('\nVerifying the installation...')
    verify(env_name)
    print('\n✓ Environment updated!')


def clean():
    """Remove the environment."""
    env_name = get_env_name()

    response = input(f'Remove the "{env_name}" environment? [y/N] ')
    if response.lower() != 'y':
        print('Cancelled.')
        return

    print(f'\nRemoving environment: {env_name}')
    print('(this may take a minute...)')

    # Capture output to filter the harmless mamba_trash.txt warning on Windows.
    result = subprocess.run(
        f'{PKG_MGR} env remove -n {env_name} -y',
        shell=True,
        capture_output=True,
        text=True,
    )
    if result.stdout:
        print(result.stdout)
    if result.stderr:
        lines = [
            line
            for line in result.stderr.split('\n')
            if 'mamba_trash.txt' not in line.lower()
            and 'error opening for writing' not in line.lower()
        ]
        if any(line.strip() for line in lines):
            print('\n'.join(lines))

    if (
        result.returncode == 0
        or 'Environment removed' in result.stdout
        or 'Environment removed' in result.stderr
    ):
        print('\n✓ Environment removed!')
    else:
        print(f'\n✗ Failed to remove environment (exit code {result.returncode})')


def list_envs():
    """List all conda/mamba environments."""
    print(f'\nListing all {PKG_MGR} environments:')
    run(f'{PKG_MGR} env list')


def main():
    commands = {
        'setup': ('Create the course environment from requirements.yml', setup),
        'update': ('Update the environment from requirements.yml', update),
        'clean': ('Remove the environment', clean),
        'verify': ('Check that course packages import', lambda: verify()),
        'launcher': (
            'Install the terminal launcher command',
            lambda: install_launcher(get_env_name()),
        ),
        'list': ('List all conda/mamba environments', list_envs),
    }

    if len(sys.argv) < 2 or sys.argv[1] not in commands:
        print(f'EE 508 environment manager (using {PKG_MGR})')
        print('\nUsage: python dev.py <command>')
        print('\nCommands:')
        for cmd, (desc, _) in commands.items():
            print(f'  {cmd:9} - {desc}')
        sys.exit(1)

    commands[sys.argv[1]][1]()


if __name__ == '__main__':
    main()