"""ps-like script to dump processes information."""

import os
import sys
from itertools import repeat
from time import sleep
from argparse import (
    ArgumentParser,
    ArgumentTypeError)

from toolrack.script import Script

from ..process.collection import (
    Collection,
    Collector)
from ..process.filter import CommandLineFilter
from ..process.formatters import (
    get_formats,
    get_formatter)


class ProcsScript(Script):
    """ps-like utility.

    It supports commandline-based filtering and output in different formats.

    """

    def get_parser(self):
        parser = ArgumentParser(
            description='Dump info about running processes.')

        def pids(pid_list):
            """Comma-separated list of PIDs."""
            try:
                return [int(pid) for pid in pid_list.split(',')]
            except:
                raise ArgumentTypeError('Must specify a list of PIDs')

        parser.add_argument(
            '--available-stats', help='Print a list of available stats.',
            action='store_true')
        parser.add_argument(
            '--fields', '-f',
            help='comma-separated list of fields to display',
            default='pid,stat.state,comm')
        parser.add_argument(
            '--regexp', '-r', help='regexp to filter by process name')
        parser.add_argument(
            '--cmdline-regexp', '-R',
            help='regexp to filter by full command line')
        parser.add_argument(
            '--pids', '-p', help='list specific PIDs', type=pids)
        parser.add_argument(
            '--format', '-F', help='output format', choices=get_formats(),
            default='table')
        parser.add_argument(
            '--interval', '-i',
            help='sample interval in seconds (default %(default)s)', type=int,
            default=5)
        parser.add_argument(
            '--count', '-c',
            help='number of samples to collect (default unlimited)', type=int,
            default=0)
        return parser

    def main(self, args):
        if args.available_stats:
            self._print_available_stats()
            self.exit()

        fields = [field.strip() for field in args.fields.split(',')]

        collector = Collector(pids=args.pids)
        collection = Collection(collector=collector)
        if args.regexp:
            collection.add_filter(CommandLineFilter(args.regexp))
        if args.cmdline_regexp:
            collection.add_filter(
                CommandLineFilter(args.cmdline_regexp, include_args=True))
        formatter_class = get_formatter(args.format)
        formatter = formatter_class(sys.stdout, fields)

        count_iter = range(args.count) if args.count else repeat(None)
        for n in count_iter:
            formatter.format(collection)
            if n != args.count - 1:
                # don't sleep after last iteration
                sleep(args.interval)

    def _print_available_stats(self):
        collector = Collector(pids=[os.getpid()])
        collection = Collection(collector=collector)
        [process] = collection
        for stat in process.available_stats():
            print(stat)


script = ProcsScript()
