#!/usr/bin/env python3
"""Management commands."""

from __future__ import annotations

import os
import shutil
import subprocess
import sys
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Iterator

PYTHON_VERSIONS = os.getenv('PYTHON_VERSIONS', '3.12 3.13').split()
PYTHON_VERSION = os.getenv('PYTHON_VERSION', '3.13')

exe = ''
prefix = ''


def shell(cmd: str, capture_output: bool = False, **kwargs: Any) -> str | None:
    """Run a shell command."""
    if capture_output:
        return subprocess.check_output(cmd, shell=True, text=True, **kwargs)  # noqa: S602
    subprocess.run(cmd, shell=True, check=True, stderr=subprocess.STDOUT, **kwargs)  # noqa: S602
    return None


@contextmanager
def environ(**kwargs: str) -> Iterator[None]:
    """Temporarily set environment variables."""
    original = dict(os.environ)
    os.environ.update(kwargs)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(original)


def uv_install() -> None:
    """Install dependencies using uv."""
    if 'CI' in os.environ:
        # Install pytorch CPU version first
        shell(
            'uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu'
        )
        shell('uv sync')
    else:
        # Install pytorch CPU version first
        shell('uv run pip install torch torchvision torchaudio')
        shell('uv sync')


def setup() -> None:
    """Setup the project."""
    if not shutil.which('uv'):
        raise ValueError(
            'make: setup: uv must be installed, see https://docs.astral.sh/uv/'
        )

    print('Installing dependencies (default environment)')  # noqa: T201
    uv_install()

    if PYTHON_VERSIONS:
        for version in PYTHON_VERSIONS:
            print(f'\nInstalling dependencies (python{version})')  # noqa: T201
            uv_activate(version)
            uv_install()


def uv_activate(python_ver: str = '3.12') -> None:
    """Activate a virtual environment."""
    global exe, prefix  # noqa: PLW060
    shell(f'uv python install {python_ver}')
    shell(f'uv sync')


def run(version: str, cmd: str, *args: str, **kwargs: Any) -> None:
    """Run a command in a virtual environment."""
    kwargs = {'check': True, **kwargs}
    if version == 'default':
        uv_activate(PYTHON_VERSION)
        subprocess.run([f'{prefix}{cmd}{exe}', *args], **kwargs)  # noqa: S603, PLW1510
    else:
        uv_activate(version)
        os.environ['MULTIRUN'] = '1'
        subprocess.run([f'{prefix}{cmd}{exe}', *args], **kwargs)  # noqa: S603, PLW1510


def multirun(cmd: str, *args: str, **kwargs: Any) -> None:
    """Run a command for all configured Python versions."""
    if PYTHON_VERSIONS:
        for version in PYTHON_VERSIONS:
            run(version, cmd, *args, **kwargs)
    else:
        run(PYTHON_VERSION, cmd, *args, **kwargs)


def allrun(cmd: str, *args: str, **kwargs: Any) -> None:
    """Run a command in all virtual environments."""
    run('default', cmd, *args, **kwargs)
    if PYTHON_VERSIONS:
        multirun(cmd, *args, **kwargs)


def clean() -> None:
    """Delete build artifacts and cache files."""
    paths_to_clean = ['build', 'dist', 'htmlcov', 'site', '.coverage*', '.pdm-build']
    for path in paths_to_clean:
        shell(f'rm -rf {path}')

    cache_dirs = [
        '.cache',
        '.pytest_cache',
        '.mypy_cache',
        '.ruff_cache',
        '__pycache__',
    ]
    for dirpath in Path('.').rglob('*'):
        if any(dirpath.match(pattern) for pattern in cache_dirs) and not (
            dirpath.match('.venv') or dirpath.match('.venvs')
        ):
            shutil.rmtree(path, ignore_errors=True)


def main() -> int:
    """Main entry point."""
    args = list(sys.argv[1:])
    if not args or args[0] == 'help':
        print('Available commands')  # noqa: T201
        print('  help                  Print this help. Add task name to print help.')  # noqa: T201
        print(
            '  setup                 Setup all virtual environments (install dependencies).'
        )  # noqa: T201
        print(
            '  run                   Run a command in the default virtual environment.'
        )  # noqa: T201
        print(
            '  multirun              Run a command for all configured Python versions.'
        )  # noqa: T201
        print('  allrun                Run a command in all virtual environments.')  # noqa: T201
        print('  clean                 Delete build artifacts and cache files.')  # noqa: T201
        try:
            run(PYTHON_VERSION, 'python3', '-V', capture_output=True)
        except (subprocess.CalledProcessError, ValueError):
            pass
        else:
            print(f'args: {args}')  # noqa: T201

    while args:
        cmd = args.pop(0)

        if cmd == 'run':
            run(PYTHON_VERSION, *args)
            return 0

        if cmd == 'multirun':
            multirun(*args)
            return 0

        if cmd == 'allrun':
            allrun(*args)
            return 0

        opts = []
        while args and (args[0].startswith('-') or '=' in args[0]):
            opts.append(args.pop(0))

        if cmd == 'clean':
            clean()
        elif cmd == 'setup':
            setup()

    return 0


if __name__ == '__main__':
    try:
        sys.exit(main())
    except subprocess.CalledProcessError as process:
        if process.output:
            print(process.output, file=sys.stderr)  # noqa: T201
        sys.exit(process.returncode)
