import json
import re
from decimal import Decimal
from typing import Any, Mapping, Optional, Union

import click


class ByteSizeParamType(click.ParamType):
    name = "byte"

    _rx_digits = re.compile(r"^(\d+(?:\.\d*)?)([kmgtpe]?)$", re.I)
    _scales = {
        "k": 2**10,
        "m": 2**20,
        "g": 2**30,
        "t": 2**40,
        "p": 2**50,
        "e": 2**60,
    }

    def convert(self, value, param, ctx):
        if isinstance(value, int):
            return value
        if not isinstance(value, str):
            self.fail(
                f"expected string, got {value!r} of type {type(value).__name__}",
                param,
                ctx,
            )
        m = self._rx_digits.search(value)
        if m is None:
            self.fail(f"{value!r} is not a valid byte-size expression", param, ctx)
        size = float(m.group(1))
        unit = m.group(2).lower()
        return int(size * self._scales.get(unit, 1))


class ByteSizeParamCheckType(ByteSizeParamType):
    name = "byte-check"

    def convert(self, value, param, ctx):
        if isinstance(value, int):
            return value
        if not isinstance(value, str):
            self.fail(
                f"expected string, got {value!r} of type {type(value).__name__}",
                param,
                ctx,
            )
        m = self._rx_digits.search(value)
        if m is None:
            self.fail(f"{value!r} is not a valid byte-size expression", param, ctx)
        return value


class CommaSeparatedKVListParamType(click.ParamType):
    name = "comma-seperated-KVList-check"

    def convert(self, value: Union[str, Mapping[str, str]], param, ctx) -> Mapping[str, str]:
        if isinstance(value, dict):
            return value
        if not isinstance(value, str):
            self.fail(
                f"expected string, got {value!r} of type {type(value).__name__}",
                param,
                ctx,
            )
        override_map = {}
        for assignment in value.split(","):
            try:
                k, _, v = assignment.partition("=")
                if k == "" or v == "":
                    raise ValueError(f"key or value is empty. key = {k}, value = {v}")
            except ValueError:
                self.fail(
                    f"{value!r} is not a valid mapping expression",
                    param,
                    ctx,
                )
            else:
                override_map[k] = v
        return override_map


class JSONParamType(click.ParamType):
    """
    A JSON string parameter type.
    The default value must be given as a valid JSON-parsable string,
    not the Python objects.
    """

    name = "json-string"

    def __init__(self) -> None:
        super().__init__()
        self._parsed = False

    def convert(
        self,
        value: Optional[str],
        param: Optional[click.Parameter],
        ctx: Optional[click.Context],
    ) -> Any:
        if self._parsed:
            # Click invokes this method TWICE
            # for a default value given as string.
            return value
        self._parsed = True
        if value is None:
            return None
        try:
            return json.loads(value)
        except json.JSONDecodeError:
            self.fail(f"cannot parse {value!r} as JSON", param, ctx)
        return value


def drange(start: Decimal, stop: Decimal, num: int):
    """
    A simplified version of numpy.linspace with default options
    """
    delta = stop - start
    step = delta / (num - 1)
    yield from (start + step * Decimal(tick) for tick in range(0, num))


class RangeExprOptionType(click.ParamType):
    """
    Accepts a range expression which generates a range of values for a variable.

    Linear space range: "linspace:1,2,10" (start, stop, num) as in numpy.linspace
    Pythonic range: "range:1,10,2" (start, stop[, step]) as in Python's range
    Case range: "case:a,b,c" (comma-separated strings)
    """

    _rx_range_key = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
    name = "Range Expression"

    def convert(self, arg, param, ctx):
        key, value = arg.split("=", maxsplit=1)
        assert self._rx_range_key.match(key), "The key must be a valid slug string."
        try:
            if value.startswith("case:"):
                return key, value[5:].split(",")
            elif value.startswith("linspace:"):
                start, stop, num = value[9:].split(",")
                return key, tuple(drange(Decimal(start), Decimal(stop), int(num)))
            elif value.startswith("range:"):
                range_args = map(int, value[6:].split(","))
                return key, tuple(range(*range_args))
            else:
                self.fail("Unrecognized range expression type", param, ctx)
        except ValueError as e:
            self.fail(str(e), param, ctx)


class CommaSeparatedListType(click.ParamType):

    name = "List Expression"

    def convert(self, arg, param, ctx):
        try:
            if isinstance(arg, int):
                return arg
            elif isinstance(arg, str):
                return arg.split(",")
        except ValueError as e:
            self.fail(repr(e), param, ctx)
