from __future__ import annotations

import os
from contextlib import redirect_stderr
from typing import IO, TYPE_CHECKING, Any, Literal, cast

if TYPE_CHECKING:
    from argparse import ArgumentParser
    from collections.abc import Iterator, Mapping, Sequence

    from ..env.task_env import TaskEnv
    from ..io import PoeIO

from ..exceptions import ConfigValidationError, ExecutionError
from ..options import PoeOptions

ArgParams = dict[str, Any]
ArgsDef = list[str] | list[ArgParams] | dict[str, ArgParams]

arg_types: dict[str, type] = {
    "string": str,
    "float": float,
    "integer": int,
    "boolean": bool,
}


class ArgSpec(PoeOptions):
    default: str | int | float | bool | None = None
    help: str = ""
    name: str
    options: Sequence[str]
    positional: bool | str = False
    required: bool = False
    type: Literal["string", "float", "integer", "boolean"] = "string"
    multiple: bool | int = False
    choices: Sequence[str] | Sequence[float] | Sequence[int] | None = None

    @classmethod
    def normalize(
        cls, source: Mapping[str, Any] | list[Mapping[str, Any]], strict: bool = True
    ):
        """
        Because arguments can be declared with different structures
        (i.e. dict or list), this function normalizes the input into a list of
        dictionaries with necessary keys.

        This is also where we do any validation that requires access to the raw
        config.
        """
        if isinstance(source, list):
            for item in source:
                if isinstance(item, str):
                    yield {"name": item, "options": (f"--{item.lstrip('_')}",)}
                elif isinstance(item, dict):
                    yield dict(
                        item,
                        options=cls._get_arg_options_list(item, strict=strict),
                    )
                elif strict:
                    raise ConfigValidationError(
                        f"Argument {item!r} has invalid type, a string or dict is "
                        "expected"
                    )

        elif isinstance(source, dict):
            for name, params in source.items():
                if not isinstance(params, dict):
                    raise ConfigValidationError(
                        f"Invalid configuration for arg {name!r}, expected dict"
                    )
                if strict and "name" in params:
                    raise ConfigValidationError(
                        f"Unexpected 'name' option for argument {name!r}"
                    )
                yield dict(
                    params,
                    name=name,
                    options=cls._get_arg_options_list(params, name, strict),
                )

    @classmethod
    def parse(
        cls,
        source: Mapping[str, Any] | list,
        strict: bool = True,
        extra_keys: Sequence[str] = (),
    ) -> Iterator[ArgSpec]:
        """
        Override parse function to perform validations that require considering all
        argument declarations at once.
        """
        try:
            result: tuple[ArgSpec] = tuple(super().parse(source, strict, extra_keys))  # type: ignore[assignment]
        except ConfigValidationError as error:
            PoeTaskArgs._enrich_config_error(error, cast("ArgsDef", source))
            raise

        if strict:
            arg_names = set()
            option_args: dict[str, str] = {}
            positional_multiple = None
            for arg in result:
                if arg.name in arg_names:
                    raise ConfigValidationError(
                        f"Duplicate argument name {arg.name!r}",
                        context=f"Invalid argument {arg.name!r} declared",
                    )
                arg_names.add(arg.name)

                if not arg.positional:
                    for option in ArgSpec._get_arg_options_list(arg, arg.name, strict):
                        if option in option_args:
                            raise ConfigValidationError(
                                f"Arguments {option_args[option]!r} and"
                                f" {arg.name!r} generate the same CLI"
                                f" option {option!r}",
                            )
                        option_args[option] = arg.name

                if arg.positional:
                    if positional_multiple:
                        raise ConfigValidationError(
                            f"Only the last positional arg of task may accept"
                            f" multiple values (not {positional_multiple!r}).",
                            context=f"Invalid argument {arg.name!r} declared",
                        )
                    if arg.multiple:
                        positional_multiple = arg.name
        yield from result

    @staticmethod
    def _get_arg_options_list(
        arg: Mapping[str, Any] | ArgSpec, name: str | None = None, strict: bool = True
    ):
        positional = arg.get("positional", False)
        name = name or arg.get("name")
        stripped = (name or "").lstrip("_")
        if positional:
            if strict and arg.get("options"):
                raise ConfigValidationError(
                    f"Positional argument {name!r} may not declare options"
                )
            # Fill in the options param in a way that makes sense for argparse
            if isinstance(positional, str):
                return [positional]
            return [name]
        return tuple(arg.get("options", [f"--{stripped}"]))

    def validate(self):
        try:
            return self._validate()
        except ConfigValidationError as error:
            error.context = f"Invalid argument {self.name!r} declared"
            raise

    def _validate(self):
        if not self.name.replace("-", "_").isidentifier():
            raise ConfigValidationError(
                f"Argument name {self.name!r} is not a valid 'identifier',\n"
                f"see the following documentation for details "
                f"https://docs.python.org/3/reference/lexical_analysis.html#identifiers"
            )

        if self.positional:
            if self.type == "boolean":
                raise ConfigValidationError(
                    f"Positional argument {self.name!r} may not have type 'boolean'"
                )

            if isinstance(self.positional, str) and not self.positional.isidentifier():
                raise ConfigValidationError(
                    f"positional name {self.positional!r} for arg {self.name!r} is "
                    "not a valid 'identifier'\n"
                    "see the following documentation for details "
                    "https://docs.python.org/3/reference/lexical_analysis.html#identifiers"
                )
        else:
            for option in self.options:
                if not option.strip():
                    raise ConfigValidationError(
                        "Invalid empty value in CLI options list"
                    )
                if option[0] != "-":
                    suggestion = f"-{option}" if len(option) == 1 else f"--{option}"
                    raise ConfigValidationError(
                        f"Invalid CLI option provided {option!r}, did you mean "
                        f"{suggestion!r}?"
                    )

        if (
            not isinstance(self.multiple, bool)
            and isinstance(self.multiple, int)
            and self.multiple < 2
        ):
            raise ConfigValidationError(
                "The 'multiple' option accepts a boolean or integer >= 2"
            )

        if self.multiple is not False and self.type == "boolean":
            raise ConfigValidationError(
                "Argument with type 'boolean' may not declare option 'multiple'"
            )

        # Ensure choices are compatible with type
        if self.choices is not None:
            arg_type = arg_types.get(self.type, str)
            for choice in self.choices:
                if not isinstance(choice, arg_type) or isinstance(choice, bool):
                    raise ConfigValidationError(
                        f"Argument {self.name!r} has invalid choice value {choice!r} "
                        f"that does not match type the configured {self.type!r}. "
                        "(maybe update the type option on the argument?)"
                    )
            if (
                self.default is not None
                and (not isinstance(self.default, str) or "${" not in self.default)
                and self.default not in self.choices
            ):
                raise ConfigValidationError(
                    f"Argument {self.name!r} has default value {self.default!r} that "
                    f"is not included in the configured choices {self.choices!r}"
                )
            if self.type == "boolean":
                raise ConfigValidationError(
                    f"Argument {self.name!r} with type 'boolean' may not declare "
                    f"option 'choices'"
                )


class PoeTaskArgs:
    _args: tuple[ArgSpec, ...]

    def __init__(self, args_def: ArgsDef, task_name: str, io: PoeIO):
        self._task_name = task_name
        self._args = self._parse_args_def(args_def)
        self._io = io

    def _parse_args_def(self, args_def: ArgsDef) -> tuple[ArgSpec, ...]:
        try:
            return tuple(ArgSpec.parse(args_def))
        except ConfigValidationError as error:
            self._enrich_config_error(error, args_def, self._task_name)
            raise

    @classmethod
    def get_help_content(
        cls, args_def: ArgsDef | None, task_name: str, suppress_errors: bool = False
    ) -> list[tuple[tuple[str, ...], str, str]]:
        if args_def is None:
            return []

        def format_arg_details(arg) -> str:
            parts: list[str] = []
            if default := arg.get("default"):
                parts.append(f"default: {default}")
            if choices := arg.get("choices"):
                parts.append(f"choices: {', '.join(map(repr, choices))}")
            if parts:
                return f"[{'; '.join(parts)}]"
            return ""

        try:
            return [
                (
                    cast("tuple[str, ...]", arg["options"]),
                    str(arg.get("help", "")),
                    format_arg_details(arg),
                )
                for arg in ArgSpec.normalize(args_def, strict=False)  # type: ignore[arg-type]
            ]
        except ConfigValidationError as error:
            if suppress_errors:
                return []
            else:
                cls._enrich_config_error(error, args_def, task_name)
                raise

    @staticmethod
    def _enrich_config_error(
        error: ConfigValidationError, args_def: ArgsDef, task_name: str = ""
    ):
        if isinstance(error.index, int):
            if isinstance(args_def, list):
                item = args_def[error.index]
                if arg_name := (isinstance(item, dict) and item.get("name")):
                    arg_ref = arg_name
                else:
                    arg_ref = error.index
            elif arg_name := tuple(args_def.keys())[error.index]:
                arg_ref = arg_name
            else:
                arg_ref = error.index
            error.context = f"Invalid argument {arg_ref!r} declared"
        if task_name:
            error.task_name = task_name

    def build_parser(self, env: TaskEnv, program_name: str) -> ArgumentParser:
        import argparse

        parser = argparse.ArgumentParser(
            prog=f"{program_name} {self._task_name}",
            add_help=False,
            allow_abbrev=False,
        )
        for arg in self._args:
            parser.add_argument(
                *arg.options,
                **self._get_argument_params(arg, env),
            )
        return parser

    def _get_argument_params(self, arg: ArgSpec, env: TaskEnv):
        default = arg.get("default")
        if isinstance(default, str):
            default = env.fill_template(default)

        result = {
            "default": default,
            "help": arg.get("help", ""),
        }

        required = arg.get("required", False)
        multiple = arg.get("multiple", False)
        arg_type = str(arg.get("type"))

        if multiple is True:
            if required:
                result["nargs"] = "+"
            else:
                result["nargs"] = "*"
        elif multiple and isinstance(multiple, int):
            result["nargs"] = multiple

        if arg.get("positional", False):
            if not multiple and not required:
                result["nargs"] = "?"
        else:
            result["dest"] = arg.name
            result["required"] = required

        if arg.choices is not None:
            result["choices"] = arg.choices

        if arg_type == "boolean":
            if default:
                result["action"] = "store_false"
            else:
                result["action"] = "store_true"
                result["default"] = False
        else:
            result["type"] = arg_types.get(arg_type, str)

        return result

    def parse(self, args: Sequence[str], env: TaskEnv, program_name: str):
        error_stream = (
            self._io.error_output
            if self._io.verbosity > -3
            else cast("IO[str]", os.devnull)
        )
        with redirect_stderr(error_stream):
            try:
                parsed_args = vars(
                    self.build_parser(env, program_name).parse_args(args)
                )
            except SystemExit as error:
                raise ExecutionError(
                    f"Invalid arguments for task {self._task_name!r}"
                ) from error

        # Ensure positional args are still exposed by name even if they were parsed with
        # alternate identifiers
        for arg in self._args:
            if isinstance(arg.positional, str):
                parsed_args[arg.name] = parsed_args[arg.positional]
                del parsed_args[arg.positional]
        # args named with dash case are converted to snake case before being exposed
        return {name.replace("-", "_"): value for name, value in parsed_args.items()}
