import sys
from functools import wraps
from typing import Callable, Iterable, Protocol, Sequence, TextIO, overload

import click
import z3

from problems import mark
from problems._markers import MarkerValue
from problems.problem import Problem


class AcceptingProblemSet[T, **P](Protocol):
    def __call__(
        self, problems: set[Problem], *args: P.args, **kwargs: P.kwargs
    ) -> T: ...


class AcceptingProblemAndMarkerOptions[T, **P](Protocol):
    def __call__(
        self,
        smt_files: tuple[TextIO, ...],
        names: tuple[str, ...],
        markers: tuple[str, ...],
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> T: ...


@overload
def problem_options[T, **P](
    fun: AcceptingProblemSet[T, P], /
) -> AcceptingProblemAndMarkerOptions[T, P]: ...


@overload
def problem_options[T, **P](
    **default_markers: MarkerValue,
) -> Callable[[AcceptingProblemSet[T, P]], AcceptingProblemAndMarkerOptions[T, P]]: ...


def problem_options[T, **P](
    fun: AcceptingProblemSet[T, P] | None = None, /, **default_markers: MarkerValue
) -> (
    AcceptingProblemAndMarkerOptions[T, P]
    | Callable[[AcceptingProblemSet[T, P]], AcceptingProblemAndMarkerOptions[T, P]]
):
    def wrapper_wrapper(
        non_none_fun: AcceptingProblemSet[T, P],
    ) -> AcceptingProblemAndMarkerOptions[T, P]:
        @click.argument("smt_files", type=click.File("r"), nargs=-1)
        @click.option(
            "--name",
            "-n",
            "names",
            type=str,
            multiple=True,
            help="Names of pre-defined problems. "
            "To see a list of pre-defined problems "
            "run `problems list`.",
        )
        @click.option(
            "--marker",
            "-m",
            "markers",
            type=str,
            multiple=True,
            help="Set of markers (attributes) to match against pre-defined problems. "
            "To see a list of available markers "
            "run `problems markers`.",
        )
        @wraps(non_none_fun)
        @click.pass_context
        def wrapper(
            ctx: click.Context,
            /,
            smt_files: tuple[TextIO, ...],
            names: tuple[
                str,
                ...,
            ],
            markers: tuple[str, ...],
            *args: P.args,
            **kwargs: P.kwargs,
        ) -> T:
            return ctx.invoke(
                non_none_fun,
                problems=_collect_problems(
                    smt_files, names, markers, **default_markers
                ),
                *args,
                **kwargs,
            )

        return wrapper

    if fun is None:
        return wrapper_wrapper
    else:
        return wrapper_wrapper(fun)


def _collect_problems(
    smt_files: Sequence[TextIO],
    names: Sequence[str],
    markers: Sequence[str],
    **default_markers: MarkerValue,
) -> set[Problem]:
    default_problems = mark.get(**default_markers)
    if not smt_files and not names and not markers:
        return default_problems

    input_problems = set(map(_smt_file_to_problem, smt_files))

    try:
        named_problems = {mark[problem.replace("-", "_")] for problem in names}
    except KeyError as e:
        click.secho(f"Could not find a problem named `{e.args[0]}`", fg="red")
        sys.exit(-1)

    if markers:
        marker_problems = mark.get(**(_parse_markers(markers)))
    else:
        marker_problems = set()

    return input_problems | ((named_problems | marker_problems) & default_problems)


def _parse_markers(markers: Iterable[str]) -> dict[str, MarkerValue]:
    parsed_markers = {}
    for marker in markers:
        parts = marker.split("=")
        value: MarkerValue
        match len(parts):
            case 1:
                (key,) = parts
                if key.startswith("no-"):
                    key = key[3:]
                    value = False
                else:
                    value = True
            case 2:
                key, value = parts
                try:
                    value = int(value)
                except ValueError:
                    pass
            case _:
                click.secho(f"Bad marker {marker}", fg="red")
                sys.exit(-1)

        parsed_markers[key] = value
    return parsed_markers


def _smt_file_to_problem(smt_file: TextIO) -> Problem:
    smt_source = smt_file.read()
    solver = z3.Solver()
    try:
        solver.from_string(smt_source)
    except z3.Z3Exception as e:
        click.secho(
            f"Bad SMTLIB file {smt_file.name}: {e.value.decode("utf-8")}",
            fg="red",
            nl=False,
        )
        sys.exit(-1)

    formula = z3.And(*solver.assertions())

    return Problem(smt_file.name, lambda: formula)
