"""
Module ``dagon.tool.main``
##########################
"""

from __future__ import annotations

import argparse
import asyncio
import importlib
import importlib.util
import logging
import os
import shlex
import subprocess
import sys
import textwrap
import traceback
from typing_extensions import Protocol
import warnings
from pathlib import Path
from typing import Iterable, NoReturn, Optional, Sequence, cast

import pkg_resources

from ..core.result import Cancellation, Failure
from ..ext.exec import ExtAwareExecutor
from ..ext.loader import ExtLoader
from ..task.dag import TaskDAG, populate_dag_context
from ..task.task import DisabledTaskError, InvalidTask
from ..tool.args import ParsedArgs, get_argparser

_g_loading_dag = False


def _format_docstr(s: str) -> Iterable[str]:
    paragraphs = textwrap.dedent(s).split('\n\n')
    for par in paragraphs:
        for line in textwrap.wrap(par):
            yield line.strip()
        yield ''  # Empty line after paragraph


def _list_tasks(dag: TaskDAG, with_docs: bool = True) -> int:
    for t in dag.tasks:
        name_line = t.name
        if t.depends:
            name_line += ' : ' + ' '.join(shlex.quote(dep.dep_name) for dep in t.depends)
        print(name_line)
        if with_docs and t.doc:
            for line in _format_docstr(t.doc):
                print(f'    {line}')
        if t.is_disabled:
            print(f'    [DISABLED: {t.disabled_reason}]')
    return 0


def _configure_logging(level: int) -> None:
    logging.basicConfig(level=level, format='[dagon] %(message)s')


def _print_failure(task_name: str, fail: Failure) -> None:
    if isinstance(fail.exception, asyncio.CancelledError):
        # Don't print exceptions which are just task cancellations
        return

    head = f'Task "{task_name}" FAILED'

    term_width = int(os.environ.get('COLUMNS', '80'))

    if isinstance(fail.exception, subprocess.CalledProcessError):
        print(f'{head}: Subprocess execution returned non-zero', file=sys.stderr)
        traceback.print_exception(fail.exception_type, fail.exception, fail.traceback)
        print(f'{" Process output ":=^{term_width}}', file=sys.stderr)
        print((fail.exception.stdout or b'').decode(), file=sys.stderr, end='')
        print((fail.exception.stderr or b'').decode(), file=sys.stderr, end='')
        print(f'{" End Output ":^^{term_width}}', file=sys.stderr)
    else:
        print(f'{head}: Encountered an exception during execution:', file=sys.stderr)
        traceback.print_exception(fail.exception_type, fail.exception, fail.traceback)


def run_with_args(dag: TaskDAG, exts: ExtLoader, args: ParsedArgs, default_tasks: Optional[Sequence[str]]) -> int:
    if args.list_tasks:
        return _list_tasks(dag, with_docs=not args.no_doc)

    _configure_logging(logging.DEBUG if args.debug else logging.INFO)

    tasks = list(args.tasks or default_tasks or (t.name for t in dag.tasks if t.is_default))
    if not tasks:
        warnings.warn('No tasks are set to be executed')

    try:
        ll_dag = dag.low_level_graph(tasks)
    except DisabledTaskError as e:
        print(f'Unable to create the execution plan: {e}', file=sys.stderr)
        return 1
    except InvalidTask as e:
        if e.candidate:
            print(f'No such task "{e.key}". Did you mean "{e.candidate.name}" ?', file=sys.stderr)
        else:
            print(f'No such task "{e.key}".', file=sys.stderr)
        return 1

    exe = ExtAwareExecutor(exts, ll_dag, catch_signals=True, fail_cancels=args.fail_cancels)

    results = exe.run_all_until_complete()
    failed = False
    for item in results.values():
        if isinstance(item.result, Failure):
            _print_failure(item.task.name, item.result)
            failed = True
        elif isinstance(item.result, Cancellation):
            print(f'Task "{item.task.name}" was cancelled')
            failed = True
        else:
            # Nothing to print for successes
            pass

    return 1 if failed else 0


def run_for_dag(dag: TaskDAG,
                exts: ExtLoader,
                *,
                argv: Optional[Sequence[str]] = None,
                default_tasks: Optional[Sequence[str]] = None) -> int:
    argv = argv if argv is not None else sys.argv[1:]
    parser = get_argparser(exts)
    args = cast(ParsedArgs, parser.parse_args(argv))
    exts.handle_options(cast(argparse.Namespace, args))
    return run_with_args(dag, exts, args, default_tasks)


def get_extensions() -> ExtLoader:
    return ExtLoader.default()


class MainArgParseResults(ParsedArgs, Protocol):
    file: Optional[str]
    module: Optional[str]
    dir: Optional[Path]
    version: bool


def main(argv: Sequence[str]) -> int:
    exts = get_extensions()
    with exts.app_context():
        return _main(argv, exts)


def _main(argv: Sequence[str], exts: ExtLoader) -> int:
    parser = get_argparser(exts)
    parser.add_argument('-f', '--file', help='The file that defines the tasks')
    parser.add_argument('-m', '--module', help='The module to import to define the tasks')
    parser.add_argument('--dir',
                        help='Add the given directory to the import/search path for finding tasks',
                        default=Path.cwd(),
                        type=Path)
    parser.add_argument(
        '--version',
        help='Display the version and exit',
        action='store_true',
    )
    args = cast(MainArgParseResults, parser.parse_args(argv))

    if args.version:
        version = pkg_resources.get_distribution('dagon').version
        print(version)
        return 0

    def_file = Path(args.file).absolute() if args.file else None
    def_mod: Optional[str] = args.module

    if def_mod and def_file:
        print('`--module` and `--file` cannot be used together', file=sys.stderr)
        return 2

    if not def_mod and not def_file:
        def_mod = 'dag'

    global _g_loading_dag  # pylint: disable=global-statement
    _g_loading_dag = True
    prev_path = list(sys.path)
    try:
        sys.path.insert(0, str(args.dir))
        dag = TaskDAG('<main>')
        if def_mod:
            with populate_dag_context(dag):
                try:
                    if def_mod in sys.modules:
                        importlib.reload(sys.modules[def_mod])
                    else:
                        importlib.import_module(def_mod)
                except ModuleNotFoundError:
                    traceback.print_exc(file=sys.stderr)
                    print(f'Failed to import the "{def_mod}" module for task definitions', file=sys.stderr)
                    return 1
        else:
            spec = importlib.util.spec_from_file_location('<dag-definition>', str(def_file))
            if not spec:
                print('Unable to find task definitions', file=sys.stderr)
                return 1
            mod = importlib.util.module_from_spec(spec)
            with populate_dag_context(dag):
                spec.loader.exec_module(mod)  # type: ignore
    finally:
        _g_loading_dag = False
        sys.path = prev_path

    exts.handle_options(cast(argparse.Namespace, args))
    return run_with_args(dag, exts, args, default_tasks=[])


# def execute_default(*, dag: TaskDAG | None = None, default: Sequence[str] | None = None) -> None:
#     if _g_loading_dag:
#         warnings.warn(
#             'Call to execute_default() while running main Dagon executable. execute_default() will do nothing.')
#         return
#     dag = dag or current_dag()
#     rc = run_for_dag(dag, argv=sys.argv[1:], default_tasks=default or [])
#     sys.exit(rc)


def start() -> NoReturn:
    sys.exit(main(sys.argv[1:]))
