#!/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: list[str], heading: str, filename: str):
    print(f'{heading}:')
    print(f'file={filename}:')
    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]:
    print(f"running command: '{command}'")
    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 filter_output(stdout: str, our_package_name: str) -> list[str]:
    output = [*{*stdout.splitlines()}]  # dedup items
    output = [
        p for p in output if not p.startswith(our_package_name)
    ]  # remove our package

    # special handling for click until python 3.9 is gone due to it dropping
    # support for active versions early
    i = 0
    for c, n in enumerate(output):
        if n.startswith('click=='):
            i = c
    if i:
        output.insert(i + 1, f"{output[i]}; python_version>='3.10'")
        output[i] = "click==8.1.8; python_version<'3.10'"

    return output


def freeze_reqs(our_package_name: str, tmpdir: str, target: str) -> None:
    target_selector = '.' if target == '' else f'.[{target}]'
    target_file = (
        'requirements.txt' if target == '' else f'requirements-{target}.txt'
    )

    print(f"installing selector: '{target_selector}'")

    run_command(f'python3 -m venv {tmpdir}')

    base_reqs: list[str] = []
    if target != '':
        # get base deps for extras
        run_command(f'{tmpdir}/bin/pip install .')
        _stdout, _, _ = run_command(f'{tmpdir}/bin/pip freeze')
        base_reqs = filter_output(_stdout, our_package_name)

    run_command(f'{tmpdir}/bin/pip install {target_selector}')
    _stdout, _, _ = run_command(f'{tmpdir}/bin/pip freeze')
    frozen = filter_output(_stdout, our_package_name)

    # remove base deps from extras
    frozen = sorted([*{*frozen} - {*base_reqs}])

    print_packages(frozen, f'frozen: {target_selector}', target_file)

    Path(target_file).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}'")
            freeze_reqs(our_package_name, tmpdir, target)


if __name__ == '__main__':
    main()