from pathlib import Path
from subprocess import run
from tempfile import TemporaryDirectory

import click

from common import config
from common.cli_utils import error_reporting
from fest.heuristic import find_sym_model_by_heuristic
from problems import Problem
from problems.benchmark_utils import average_time
from problems.cli_utils import problem_options


def run_z3(path: Path, timeout: int) -> None:
    result = run(
        ["z3", str(path)],
        timeout=timeout,
        capture_output=True,
        encoding="utf-8",
    )
    output = str(result.stdout).strip()
    if output != "sat":
        raise Exception()


def run_cvc5(path: Path, timeout: int) -> None:
    result = run(
        ["cvc5", str(path)],
        timeout=timeout,
        capture_output=True,
        encoding="utf-8",
    )
    output = str(result.stdout).strip()
    if output != "sat":
        raise Exception()


def run_vampire(path: Path, timeout: int) -> None:
    result = run(
        [
            "vampire",
            "--input_syntax",
            "smtlib2",
            "-t",
            f"{timeout}s",
            str(path),
        ],
        timeout=timeout,
        capture_output=True,
        encoding="utf-8",
    )
    output = str(result.stdout).splitlines(keepends=False)[-1].strip()
    if "Success" not in output:
        raise Exception()


@click.command(
    help="Benchmark for the POPL 2024 paper "
    "`An Infinite Needle in a Finite Haystack`",
    context_settings=dict(
        show_default=True,
        help_option_names=["-h", "--help"],
    ),
)
@click.option("--with-fest/--without-fest", default=True)
@click.option("--with-z3/--without-z3", default=True)
@click.option("--with-cvc5/--without-cvc5", default=True)
@click.option("--with-vampire/--without-vampire", default=True)
@click.option(
    "-t",
    "--timeout",
    type=int,
    help="Timeout (in Seconds)",
    default=config.popl24.timeout[int],
)
@click.option(
    "-r",
    "--runs",
    type=int,
    help="Number of runs to average runtime",
    default=config.popl24.runs[int],
)
@error_reporting
@problem_options(popl24=True)
def main(
    problems: set[Problem],
    with_fest: bool,
    with_z3: bool,
    with_cvc5: bool,
    with_vampire: bool,
    timeout: int,
    runs: int,
) -> None:
    with TemporaryDirectory() as tmpdir:
        dir_path = Path(tmpdir)
        for problem in problems:
            print(f"(problem {problem.name}")
            if with_fest:
                print(
                    f"\t(fest '{average_time(lambda: find_sym_model_by_heuristic(problem, quiet=True), runs)}')"
                )

            smt_path = (dir_path / problem.name).with_suffix(".smt2")
            smt_path.write_text(problem.to_smt())

            if with_z3:
                print(
                    f"\t(z3 '{average_time(lambda: run_z3(smt_path, timeout), runs)}')"
                )

            if with_cvc5:
                print(
                    f"\t(cvc5 '{average_time(lambda: run_cvc5(smt_path, timeout), runs)}')"
                )

            if with_vampire:
                print(
                    f"\t(vampire '{average_time(lambda: run_vampire(smt_path, timeout), runs)}')"
                )

            print(")")


if __name__ == "__main__":
    main()
