# pisek  - Tool for developing tasks for programming competitions.
#
# Copyright (c)   2019 - 2022 Václav Volhejn <vaclav.volhejn@gmail.com>
# Copyright (c)   2019 - 2022 Jiří Beneš <mail@jiribenes.com>
# Copyright (c)   2020 - 2022 Michal Töpfer <michal.topfer@gmail.com>
# Copyright (c)   2022        Jiří Kalvoda <jirikalvoda@kam.mff.cuni.cz>
# Copyright (c)   2023        Daniel Skýpala <skipy@kam.mff.cuni.cz>

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# any later version.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from contextlib import ExitStack
from dataclasses import dataclass, field
from decimal import Decimal
from math import ceil
import os
from tempfile import NamedTemporaryFile, TemporaryDirectory
from typing import Optional, Any, Union, Callable
import signal
import subprocess
from tempfile import NamedTemporaryFile, TemporaryDirectory
from typing import assert_never, Optional, Any, Union, Callable

from pisek.config.task_config import ProgramRole, RunSection
from pisek.env.env import Env
from pisek.utils.paths import TaskPath, LogPath
from pisek.jobs.jobs import PipelineItemFailure
from pisek.utils.text import tab
from pisek.task_jobs.run_result import RunResultKind, RunResult
from pisek.task_jobs.task_job import TaskJob


@dataclass
class ProgramPoolItem:
    name: str
    executable: TaskPath
    args: list[str]
    time_limit: Decimal
    clock_limit: Decimal
    mem_limit: int
    process_limit: int
    stdin: Optional[Union[TaskPath, int]]
    stdout: Optional[Union[TaskPath, int]]
    stderr: Optional[TaskPath]
    env: dict[str, str] = field(default_factory=lambda: {})

    def to_popen(self, minibox: str, meta_file: str) -> dict[str, Any]:
        """Returns subprocess.Popen args for executing this PoolItem."""
        result: dict[str, Any] = {}

        minibox_args = []
        minibox_args.append(f"--time={self.time_limit}")
        minibox_args.append(f"--wall-time={self.clock_limit}")
        minibox_args.append(f"--mem={self.mem_limit*1024}")
        # Minibox doesn't support limits for multiple processes (#613)
        if self.process_limit != 1:
            minibox_args.append(f"--processes=0")

        for std in ("stdin", "stdout", "stderr"):
            attr = getattr(self, std)
            if isinstance(attr, TaskPath):
                minibox_args.append(f"--{std}={attr.abspath}")
            elif getattr(self, std) is None:
                minibox_args.append(f"--{std}=/dev/null")

            if isinstance(attr, int):
                result[std] = attr
            else:
                result[std] = subprocess.PIPE

        for key, val in self.env.items():
            minibox_args.append(f"--env={key}={val}")

        minibox_args.append("--silent")
        minibox_args.append(f"--meta={meta_file}")

        result["args"] = (
            [minibox]
            + minibox_args
            + ["--run", "--", self.executable.abspath]
            + self.args
        )
        return result


class ProgramsJob(TaskJob):
    """Job that deals with a program."""

    def __init__(self, env: Env, name: str, **kwargs) -> None:
        super().__init__(env=env, name=name, **kwargs)
        self._program_pool: list[ProgramPoolItem] = []
        self._callback: Optional[Callable[[subprocess.Popen], None]] = None

    @staticmethod
    def _env_disjoint_union(
        env1: dict[str, str], env2: dict[str, str]
    ) -> dict[str, str]:
        union = env1 | env2
        if len(union) != len(env1) + len(env2):
            duplicates = list(set(env1.keys()) & set(env2.keys()))
            raise PipelineItemFailure(
                f"Environment variable '{duplicates[0]}' is used by the protocol and cannot be overwritten."
            )
        return union

    def _load_executable(
        self,
        path: TaskPath,
        args: list[str],
        time_limit: Decimal,
        clock_limit: Decimal,
        mem_limit: int,
        process_limit: int,
        stdin: Optional[Union[TaskPath, int]] = None,
        stdout: Optional[Union[TaskPath, int]] = None,
        stderr: Optional[LogPath] = None,
        env: dict[str, str] = {},
    ):
        if self._is_file(path):
            self._access_file(path)
            executable = path
        elif self._is_dir(path):
            self._access_dir(path)
            executable = path.join("run")

        if self._is_dir(executable):
            raise PipelineItemFailure(
                f"Cannot execute '{executable:p}': Is a directory."
            )
        elif not self._exists(executable):
            raise PipelineItemFailure(f"Cannot execute '{executable:p}': File missing.")

        if isinstance(stdin, TaskPath):
            self._access_file(stdin)
        if isinstance(stdout, TaskPath):
            self.make_filedirs(stdout)
            self._access_file(stdout)
        if isinstance(stderr, TaskPath):
            self.make_filedirs(stderr)
            self._access_file(stderr)

        self._program_pool.append(
            ProgramPoolItem(
                name=path.name,
                executable=executable,
                args=args,
                time_limit=time_limit,
                clock_limit=clock_limit,
                mem_limit=mem_limit,
                process_limit=process_limit,
                stdin=stdin,
                stdout=stdout,
                stderr=stderr,
                env=env,
            )
        )

    def _load_program(
        self,
        program_role: ProgramRole,
        program: RunSection,
        args: list[str] = [],
        stdin: Optional[Union[TaskPath, int]] = None,
        stdout: Optional[Union[TaskPath, int]] = None,
        stderr: Optional[LogPath] = None,
        env={},
    ) -> None:
        """Adds program to execution pool."""
        time_limit: Decimal | None = None
        if program_role.is_solution():
            time_limit = self._env.time_limit

        self._load_executable(
            path=program.executable,
            args=program.args + args,
            time_limit=program.time_limit if time_limit is None else time_limit,
            clock_limit=program.clock_limit(time_limit),
            mem_limit=program.mem_limit,
            process_limit=program.process_limit,
            stdin=stdin,
            stdout=stdout,
            stderr=stderr,
            env=self._env_disjoint_union(env, program.env),
        )

    def _load_callback(self, callback: Callable[[subprocess.Popen], None]) -> None:
        if self._callback is not None:
            raise RuntimeError("Callback already loaded.")
        self._callback = callback

    def _read_run_result(
        self, pool_item: ProgramPoolItem, process: subprocess.Popen, meta_file: str
    ) -> RunResult:
        assert process.stderr is not None  # To make mypy happy

        if process.returncode not in (0, 1):
            raise PipelineItemFailure(
                f"Minibox error:\n{tab(process.stderr.read().decode())}"
            )

        with open(meta_file) as f:
            meta_raw = f.read().strip().split("\n")
        meta = {key: val for key, val in map(lambda x: x.split(":", 1), meta_raw)}

        time = Decimal(meta["time"])
        wall_time = Decimal(meta["time-wall"])
        memory = ceil(int(meta["max-rss"]) / 1024)

        if process.returncode == 0:
            return RunResult(
                RunResultKind.OK,
                0,
                time,
                wall_time,
                memory,
                pool_item.stdin,
                pool_item.stdout,
                pool_item.stderr,
                "Exited with return code 0",
            )

        elif process.returncode == 1:
            if meta["status"] in ("RE", "SG"):
                if meta["status"] == "RE":
                    return_code = int(meta["exitcode"])
                elif meta["status"] == "SG":
                    return_code = int(meta["exitsig"])
                    meta["message"] += f" ({signal.Signals(return_code).name})"

                return RunResult(
                    RunResultKind.RUNTIME_ERROR,
                    return_code,
                    time,
                    wall_time,
                    memory,
                    pool_item.stdin,
                    pool_item.stdout,
                    pool_item.stderr,
                    meta["message"],
                )

            elif meta["status"] == "TO":
                time_limit = (
                    f"{pool_item.time_limit}s"
                    if time > pool_item.time_limit
                    else f"{pool_item.clock_limit}ws"
                )
                return RunResult(
                    RunResultKind.TIMEOUT,
                    -1,
                    time,
                    wall_time,
                    memory,
                    pool_item.stdin,
                    pool_item.stdout,
                    pool_item.stderr,
                    f"Timeout after {time_limit}",
                )
            else:
                raise RuntimeError(f"Unknown minibox status {meta['message']}.")

        assert False

    def _run_programs(self) -> list[RunResult]:
        """Runs all programs in execution pool."""
        running_pool: list[subprocess.Popen] = []
        meta_files: list[str] = []
        tmp_dirs: list[TemporaryDirectory] = []
        minibox = TaskPath.executable_path("_minibox").abspath
        with ExitStack() as exit_stack:
            for pool_item in self._program_pool:
                meta_file = NamedTemporaryFile(prefix="pisek_", delete_on_close=False)
                exit_stack.enter_context(meta_file)
                meta_file.close()
                meta_files.append(meta_file.name)

                popen = pool_item.to_popen(minibox, meta_file.name)
                self._log(
                    "debug",
                    "Executing '" + " ".join(popen["args"]) + "'",
                    bypass_cache=True,
                )

                tmp_dir = TemporaryDirectory(prefix="pisek_")
                exit_stack.enter_context(tmp_dir)
                tmp_dirs.append(tmp_dir)

                running_pool.append(subprocess.Popen(**popen, cwd=tmp_dir.name))

            if self._callback is not None:
                callback_exec = False
                while True:
                    states = [process.poll() is not None for process in running_pool]
                    if not callback_exec and any(states):
                        callback_exec = True
                        self._callback(running_pool[states.index(True)])

                    if all(states):
                        break

            run_results = []
            for pool_item, process, tmp_dir, meta_file_path in zip(
                self._program_pool, running_pool, tmp_dirs, meta_files
            ):
                self._wait_for_subprocess(process)
                run_results.append(
                    self._read_run_result(pool_item, process, meta_file_path)
                )

                created_paths = self._listdir(TaskPath(tmp_dir.name))
                if created_paths:
                    file_list = "".join(
                        [
                            (
                                self._quote_file_with_name(p)
                                if self._is_file(p)
                                else p.col_name(self._env) + "\n"
                            )
                            for p in created_paths
                        ]
                    ).removesuffix("\n")
                    raise PipelineItemFailure(
                        f"'{pool_item.name}' created paths in cwd:\n" + tab(file_list)
                    )

        return run_results

    def _run_program(
        self,
        program_role: ProgramRole,
        program: RunSection,
        **kwargs,
    ) -> RunResult:
        """Loads one program and runs it."""
        self._load_program(program_role, program, **kwargs)
        return self._run_programs()[0]

    def _run_tool(
        self,
        program: str,
        args: list[str] = [],
        stdin: Optional[Union[TaskPath, int]] = None,
        stdout: Optional[Union[TaskPath, int]] = None,
        stderr: Optional[LogPath] = None,
        env={},
    ) -> RunResult:
        """Loads one program and runs it."""
        self._load_executable(
            TaskPath.executable_path(program),
            args=args,
            time_limit=Decimal(300),
            clock_limit=Decimal(300),
            mem_limit=0,
            process_limit=1,
            stdin=stdin,
            stdout=stdout,
            stderr=stderr,
            env=env,
        )
        return self._run_programs()[0]
