#!/usr/bin/python3

# The MIT License (MIT)
#
# Copyright © 2020-2025 pacman64
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the “Software”), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


from math import isinf, isnan, modf, sqrt
from sys import argv, exit, stderr
from typing import List


info = '''
primes [options...] [count...] [numbers...]


Find or check prime numbers.

By default it finds the first n prime numbers, emitting each one on its own
line, with a default of 100 thousand when not given a count.

When using one of the check-prime options, all later cmd-line arguments are
expected to be numbers, and are checked for primality.

All (optional) leading options start with either single or double-dash, or
even no dashes at all. Some of the options are, shown in their single-dash
form:

    -h          show this help message
    -help       show this help message

    -check      check all numbers given as cmd-line args for primality
    -isprime    same as option -check
    -is-prime   same as option -check
    -is         same as option -check
'''

# handle standard help cmd-line options, quitting right away in that case
if len(argv) > 1 and argv[1].lstrip('-') in ('h', 'help'):
    print(info.strip(), file=stderr)
    exit(0)


def show_primes(left: int) -> None:
    'Show the number of primes given.'

    # 2 is the only even prime number
    if left > 0:
        print(2)
        left -= 1

    n = 3
    while left > 0:
        if is_prime(n):
            print(n)
            left -= 1
        n += 2


def is_prime(n: int) -> bool:
    'Check if a number is prime, with O(n**0.5) time-complexity.'

    # handle non-sense inputs
    if n < 2:
        return False

    # 2 is the only even prime number
    if n % 2 == 0:
        return n == 2

    # check only odd divisors up to the square-root
    d = 3
    limit = int(sqrt(n))

    while d <= limit:
        if n % d == 0:
            return False
        d += 2
    return True


def is_float_str(s: str) -> bool:
    'Simplify control-flow for func check_primes.'

    try:
        float(s)
        return True
    except Exception:
        return False


def check_primes(args: List[str]) -> bool:
    '''
    Check all the numbers given (as strings) for primality: its output
    is lines of tab-separated values (TSV), starting with a header line.
    The return value is this func's success, namely its lack of errors.
    '''

    errors = 0
    print('number\tprime')
    for s in args:
        try:
            f = float(s)
            if isnan(f) or isinf(f) or modf(f)[0] != 0.0:
                print(f'{f:,}\tno')
            else:
                prime = 'yes' if is_prime(int(f)) else 'no'
                print(f'{f:,}\t{prime}')
        except Exception:
            print(f'\x1b[31m{s} isn\'t even a number\x1b[0m', file=stderr)
            errors += 1
    return errors == 0


try:
    check_options = ('check', 'is', 'isprime', 'is-prime', 'is_prime')

    # handle prime-checking mode
    if len(argv) > 1 and argv[1].lstrip('-').lower() in check_options:
        if len(argv) < 2:
            raise ValueError('no numbers given to check')
        if not check_primes(argv[2:]):
            exit(1)
        exit(0)

    # handle main mode, which finds/shows prime numbers: its default count
    # is 100 thousand, unless a leading cmd-line arg overrides it
    n = int(1e5)
    if len(argv) > 1:
        try:
            n = int(float(argv[1]))
        except Exception:
            raise ValueError(f'unsupported leading option {argv[1]}')

    print(f'showing the first {n:,} primes', file=stderr)
    # running inside a func speeds things up in older versions of python
    show_primes(n)
except BrokenPipeError:
    # quit quietly, instead of showing a confusing error message
    stderr.close()
except KeyboardInterrupt:
    print('\x1b[31mscript was force-quit early\x1b[0m', file=stderr)
    exit(2)
except Exception as e:
    print(f'\x1b[31m{e}\x1b[0m', file=stderr)
    exit(1)