#!/usr/bin/env python3

import re
import subprocess
from pathlib import Path
from sys import argv
from tempfile import TemporaryDirectory

SCRIPT = argv[0]
TARGETS = ('', 'dev', 'docs')  # empty string is used for the base requirements
FILE_HEADER = f'# DO NOT EDIT THIS FILE DIRECTLY - use {SCRIPT} to update\n'


def print_packages(packages, heading):
    print(f'{heading}:')
    print('  ', end='')
    print('\n  '.join(packages))


def get_package_name() -> str:
    match = re.search(
        r"name='(?P<pkg>[\w-]+)',", Path('setup.py').read_text(encoding='utf8')
    )
    if not match:
        raise Exception('failed to determine our package name')
    our_package_name = match.group('pkg')
    print(f'our_package_name: {our_package_name}')

    return our_package_name


def run_command(command: str) -> tuple[str, str, int]:
    try:
        proc = subprocess.run(
            command.split(),
            check=True,
            capture_output=True,
            encoding='utf8',
            timeout=30,
        )
    except Exception as exc:
        print(f"can't run command='{command}'")
        raise exc

    return proc.stdout.strip(), proc.stderr.strip(), proc.returncode


def save_base_reqs(our_package_name: str, tmpdir: str) -> list[str]:
    print('installing base')

    run_command(f'python3 -m venv {tmpdir}')
    run_command(f'{tmpdir}/bin/pip install .')
    stdout, _, _ = run_command(f'{tmpdir}/bin/pip freeze')

    frozen = [*{*stdout.splitlines()}]  # dedup items
    frozen = sorted([p for p in frozen if not p.startswith(our_package_name)])

    print_packages(frozen, 'base-frozen')

    Path('requirements.txt').write_text(
        FILE_HEADER + '\n'.join(frozen) + '\n', encoding='utf8'
    )

    return frozen


def save_extra_reqs(
    our_package_name: str, tmpdir: str, base_frozen: list[str], extra: str
) -> None:
    print(f'installing extra: {extra}')

    run_command(f'{tmpdir}/bin/pip install .[{extra}]')
    stdout, _, _ = run_command(f'{tmpdir}/bin/pip freeze')

    frozen = [*{*stdout.splitlines()} - {*base_frozen}]  # remove base deps
    frozen = sorted([p for p in frozen if not p.startswith(our_package_name)])

    print_packages(frozen, f'{extra}-frozen')

    Path(f'requirements-{extra}.txt').write_text(
        FILE_HEADER + '\n'.join(frozen) + '\n', encoding='utf8'
    )


def main() -> None:
    our_package_name = get_package_name()

    # install all extra dependencies in a new venv, so we don't get duplicate/unwanted deps
    for target in TARGETS:
        with TemporaryDirectory() as tmpdir:
            print(f'using tmpdir={tmpdir}')
            base_frozen = save_base_reqs(our_package_name, tmpdir)
            if not target:
                continue  # skip extra requirements for base package
            save_extra_reqs(our_package_name, tmpdir, base_frozen, target)


if __name__ == '__main__':
    main()