# Copyright 2026 Arturo Roberti
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import errno
import json
import os
import pty
import subprocess
import sys
import termios
from dataclasses import dataclass, field
from io import TextIOWrapper
from pathlib import Path
from queue import Queue
from tempfile import NamedTemporaryFile, TemporaryDirectory
from textwrap import dedent
from threading import Event, Lock, Thread

from rich.markup import escape

from gurk.lib.context import Logger, TaskTerminationType, get_logger
from gurk.lib.core.plugins import get_venv_dir
from gurk.lib.shared.scripts import (
    Command,
    CommandKind,
    SchedulerTask,
    ScriptBlock,
    ScriptBlockTypes,
    get_block_spans,
    run_script_function,
)
from gurk.lib.shared.system_info import get_system_info
from gurk.lib.utils import PatternCollection, generate_random_path, typecheck


@dataclass
class Scheduler:
    """Schedules and runs tasks with dependencies, handling logging and progress tracking."""

    # fmt: off
    tasks:        list[SchedulerTask] = field(repr=False)
    askpass_file: Path | None         = field(repr=False)
    logger:       Logger              = field(init=False, repr=False)

    results:   dict[SchedulerTask, TaskTerminationType] = field(init=False, repr=False, default_factory=dict)
    scheduled: set[SchedulerTask]                       = field(init=False, repr=False, default_factory=set)

    lock:      Lock  = field(init=False, repr=False, default_factory=Lock)
    queue:     Queue = field(init=False, repr=False, default_factory=Queue)
    # fmt:

    def __post_init__(self):
        # Validate that if the askpass_file is not provided, pytests are being run
        if self.askpass_file is None and "pytest" not in sys.modules:
            raise ValueError(
                "UNEXPECTED: Scheduler initialized without "
                "'askpass_file', but pytests don't seem to be running."
            )

    @staticmethod
    @typecheck
    def _prepare_script(command: Command) -> tuple[Path, int]:
        """
        Prepare a copy of the desired script that
        - Uses STEP statements only if in the desired function (or entrypoint)
        - Converts all STEP comments into equivalent print statements.

        :param command: Command to prepare
        :type command: Command
        :return: Tuple of (path to modified script, number of steps)
        :rtype: tuple[Path, int]
        """
        # Create temporary file to run later
        original_path = Path(command.script)
        tmp_path = generate_random_path(
            suffix=f"_{original_path.name}", prefix="modified_"
        )

        # Analyze script blocks
        script_blocks = get_block_spans(original_path)

        # Main processing loop
        n_steps = 0
        with original_path.open(
            "r", encoding="utf-8", errors="replace"
        ) as src, tmp_path.open("w", encoding="utf-8") as dst:
            for idx, line in enumerate(src, start=1):
                # Detect current block
                curr_block = [
                    block
                    for block in script_blocks
                    if block["lines"][0] <= idx <= block["lines"][1]
                ]
                if not curr_block:
                    curr_block = ScriptBlock(
                        type=ScriptBlockTypes.OTHER,
                        name=None,
                        lines=(0, 0),
                    )
                else:
                    curr_block = curr_block[0]

                # Resolve bash imports
                if (
                    command.kind == CommandKind.BASH
                    and curr_block["type"] == ScriptBlockTypes.IMPORT
                ):
                    # Resolve source expression
                    bash_script = dedent(
                        f"""
                    BASH_SOURCE[0]="{original_path.as_posix()}"
                    echo {curr_block['name'].strip()}
                    """
                    )
                    result = subprocess.run(
                        ["bash", "-c", bash_script],
                        capture_output=True,
                        text=True,
                        check=True,
                    )

                    # Write source statement with absolute path
                    abs_path = Path(result.stdout.strip())
                    if not abs_path.is_absolute():
                        abs_path = (
                            Path(original_path.as_posix()).parent / abs_path
                        ).resolve()
                    dst.write(f"source {abs_path.as_posix()}\n")
                    continue

                # Look for STEP instances
                step_patterns = PatternCollection.STEP.patterns
                m_comment = step_patterns["comment"](progress=True).match(line)
                m_any = step_patterns["any"](progress=True).match(line)
                if m_comment:
                    # STEP comment is found
                    step = m_comment.group(1).strip()
                elif m_any:
                    # Assumed (unwanted, manual) STEP print statement is found
                    step = m_any.group(1).strip()
                else:
                    # Not a STEP instance, write as is
                    dst.write(line)
                    continue

                # Handle STEP replacement/removal
                indent = len(line.rstrip("\n")) - len(
                    line.lstrip().rstrip("\n")
                )
                if (
                    (
                        command.function
                        and curr_block["type"] == ScriptBlockTypes.FUNCTION
                        and curr_block["name"] == command.function
                    )
                    or (
                        not command.function
                        and curr_block["type"] == ScriptBlockTypes.ENTRYPOINT
                    )
                ) and m_comment:
                    # Replace STEP comments with print statements
                    step_msg = f"__STEP__: {step}"
                    if command.kind == CommandKind.PYTHON:
                        step_msg = "\n" + step_msg
                        msg = f"print(f{step_msg!r})"
                    else:
                        msg = step_msg.replace("\\", "\\\\").replace(
                            '"', '\\"'
                        )
                        msg = f'printf "\\n%s\\n" "{msg}"'

                    # Write replaced STEP print statement
                    dst.write(f"{' ' * indent}{msg}\n")

                    # Add to step count
                    n_steps += 1
                else:
                    # Write removed (assumed) STEP print statement
                    dst.write(
                        f"{' ' * indent}{'pass' if command.kind == CommandKind.PYTHON else ':'}\n"
                    )

        return tmp_path, n_steps

    @typecheck
    def _spawn_and_stream(
        self, proc_cmd: list[str], flog: TextIOWrapper, task_id: int
    ) -> TaskTerminationType:
        """
        Spawn a subprocess and stream its output to the logfile and progress tracker.

        :param proc_cmd: Command to run
        :type proc_cmd: list[str]
        :param flog: Log file to write output to
        :type flog: TextIOWrapper
        :param task_id: ID of the task for progress tracking
        :type task_id: int
        :return: Task termination type (SUCCESS, FAILURE, PARTIAL)
        :rtype: TaskTerminationType
        """
        # 1. Initialize Event for PARTIAL status tracking and log output
        warning_event = Event()

        # 2. Create the PTY master and slave file descriptors
        master_fd, slave_fd = pty.openpty()

        # 3. Define the single reader function for the PTY master
        def pty_reader(master_fd: int) -> None:
            """Reads lines from the single PTY master stream."""
            # Use os.fdopen in a 'with' block to guarantee closing the master_fd upon exit.
            with os.fdopen(master_fd, "rb", 0) as master_file:
                try:
                    # Save partial lines (not yet terminated with newline) for clean logging
                    _partial_line = ""

                    # Main reading loop
                    while True:
                        # Read data and check end of file (EOF)
                        try:
                            raw_data = os.read(master_file.fileno(), 4096)
                        except OSError as e:
                            if e.errno == errno.EIO:
                                # EOF shows up as EIO (Input/Output error)
                                self.logger.debug(
                                    "(This is normal) PTY reader encountered "
                                    "an EIO error - treating as EOF."
                                )
                                break
                            else:
                                raise e

                        # Read and normalize data
                        data = raw_data.decode(
                            encoding="utf-8", errors="replace"
                        )
                        data = PatternCollection.ANSI.patterns.sub("", data)

                        # Split data into lines, handling partial lines
                        buffer = _partial_line + data.rstrip("\n")
                        lines = buffer.split("\n")
                        _partial_line = (
                            "" if data.endswith("\n") else lines.pop()
                        )

                        # Process each line for STEP statements
                        for line_raw in lines:
                            # Handle carriage return - only keep last part
                            if "\r" in line_raw.rstrip("\r"):
                                parts = line_raw.rstrip("\r").split("\r")
                                line = parts[-1]
                            else:
                                line = line_raw

                            # Write to logfile
                            flog.write(line + "\n")
                            flog.flush()

                            # Extract STEP statements with progress
                            m_progress = PatternCollection.STEP.patterns[
                                "output"
                            ](progress=True).match(line)
                            if m_progress:
                                self.logger.update_task(
                                    task_id, m_progress.group(1).strip()
                                )

                            # Extract STEP statements without progress
                            m_no_progress = PatternCollection.STEP.patterns[
                                "output"
                            ](progress=False).match(line)
                            m_no_progress_warning = (
                                PatternCollection.STEP.patterns["output"](
                                    progress=False, warning=True
                                ).match(line)
                            )
                            if m_no_progress or m_no_progress_warning:
                                if m_no_progress_warning:
                                    match = m_no_progress_warning
                                    warning_event.set()
                                else:
                                    match = m_no_progress

                                self.logger.update_task(
                                    task_id,
                                    match.group(1).strip(),
                                    advance=False,
                                )

                    # Log any remaining partial line
                    if _partial_line:
                        flog.write(_partial_line + "\n")
                        flog.flush()

                except Exception as e:
                    self.logger.debug(f"PTY reader encountered an error: {e}")

        # 4. Define preexec function for session isolation and FD cleanup in child process
        def preexec_setup():
            """Creates a new session and closes the PTY master FD in child."""
            # Use os.setsid to become session leader and claim TTY control
            os.setsid()
            # Child closes the PTY master FD
            os.close(master_fd)

            # Set terminal parameters for unbuffered output
            try:
                attrs = termios.tcgetattr(slave_fd)
                attrs[1] = attrs[1] & ~termios.ECHO
                termios.tcsetattr(slave_fd, termios.TCSANOW, attrs)
            except termios.error:
                pass

        # 5. Define environment for usage with SUDO_ASKPASS
        def create_sudo_wrapper() -> str:
            """Create a temporary sudo wrapper script, to avoid having to use 'sudo -A' everywhere."""
            # Temporary directory
            wrapper_dir = Path(TemporaryDirectory().name)
            if not wrapper_dir.is_dir():
                wrapper_dir.mkdir(parents=True, exist_ok=True)

            # Temporary sudo wrapper script
            sudo_wrapper = wrapper_dir / "sudo"
            with open(sudo_wrapper, "w") as f:
                f.write(
                    """
                    #!/bin/sh
                    exec /usr/bin/sudo -A "$@"
                """
                )
            os.chmod(sudo_wrapper, 0o700)

            return wrapper_dir.as_posix()

        env = os.environ.copy()
        if self.askpass_file is not None:
            env["SUDO_ASKPASS"] = str(self.askpass_file)
            env["PATH"] = f"{create_sudo_wrapper()}:{env.get('PATH', '')}"

        # 6. Set non-interactive environment variables
        env["DEBIAN_FRONTEND"] = "noninteractive"

        # 7. Spawn the process with PTY connections
        process = subprocess.Popen(
            proc_cmd,
            bufsize=0,
            stdin=slave_fd,
            stdout=slave_fd,
            stderr=slave_fd,
            preexec_fn=preexec_setup,
            env=env,
            text=False,
        )

        # 8. Parent closes its reference to the PTY slave
        os.close(slave_fd)

        # 9. Start the single reader thread on the PTY master
        t_out = Thread(target=pty_reader, args=(master_fd,), daemon=True)
        t_out.start()

        # 10. Wait for process exit and clean up
        exit_code = process.wait()
        t_out.join()

        # 11. Final Termination Logic: Check status and PARTIAL event
        if exit_code != 0:
            self.logger.debug("Task failed with non-zero exit code.")
            return TaskTerminationType.FAILURE
        elif warning_event.is_set():
            return TaskTerminationType.PARTIAL
        else:
            return TaskTerminationType.SUCCESS

    @typecheck
    def run_task(
        self,
        task: SchedulerTask,
        task_id: int,
    ) -> TaskTerminationType:
        """
        Run a single task, logging its output and tracking progress.

        :param task: The task to run
        :type task: SchedulerTask
        :param task_id: ID of the task for progress tracking
        :type task_id: int
        :return: Task termination type indicating success, failure, or partial completion
        :rtype: TaskTerminationType
        """
        # Prepare script with modified step statements
        modified_script, n_steps = self._prepare_script(task.command)
        self.logger.log_script(
            modified_script, task.name, ext=task.command.kind.ext
        )

        def safe_unlink(path: Path | str) -> None:
            """
            Unlink files that may or may not have been unlinked yet

            :param path: Path to the file to unlink
            :type path: Path | str
            """
            if not isinstance(path, (str, Path)):
                return
            elif Path(path).exists():
                try:
                    path.unlink()
                except Exception:
                    # Already unlinked
                    pass

        # Create args
        args = task.args + (
            "--system-info",
            json.dumps(get_system_info()),
        )
        if task.config_file:
            args += ("--config-file", task.config_file)
        args = (task.name,) + args

        # Create temporary file that will run script/call function
        try:
            # Create
            tmpwrap = NamedTemporaryFile(
                delete=False,
                suffix=Path(task.command.script).suffix,
                prefix="wrapper_",
                mode="w",
            )
            tmpwrap_path = Path(tmpwrap.name)

            # Get venv path
            plugin_name = task.name.split("/", 1)[0]
            if plugin_name == "gurk":
                venv_path = None  # Use current venv
            else:
                venv_path = get_venv_dir(plugin_name)

            # Write
            wrapper_src = run_script_function(
                script=modified_script,
                function=task.command.function,
                args=args,
                run=False,
                venv=venv_path,
            )
            tmpwrap.write(wrapper_src)
            tmpwrap.flush()
            tmpwrap.close()
            os.chmod(tmpwrap_path, os.stat(tmpwrap_path).st_mode | 0o700)
        except Exception as e:
            safe_unlink(tmpwrap_path)
            raise e

        # Get executable to run file. Also, use unbuffered output
        if task.command.kind == CommandKind.PYTHON:
            sudo_prefix = ["sudo", "-E"] if task.privileged else []
            exe_cmd = [*sudo_prefix, task.command.kind.exe, "-u"]
        else:  # Bash
            exe_cmd = ["stdbuf", "-oL", "-eL", task.command.kind.exe]

        # Combine files and args into a runnable command
        proc_cmd = [*exe_cmd, tmpwrap.name]
        self.logger.debug(
            f"Running task '{task.name}' with command:"
            f"'{' '.join(proc_cmd)}'"
        )

        # Logging
        log_file = self.logger.generate_logfile_path(task_id)
        self.logger.set_total(task_id, n_steps + 1)  # +1 for finishing step
        self.logger.info(
            f"{escape(f'[{task.name}]')} Logging to {log_file}",
            syntax_highlight=False,
        )
        flog = log_file.open("w", encoding="utf-8", errors="replace")

        # Run and stream
        try:
            success = self._spawn_and_stream(proc_cmd, flog, task_id)
        except Exception as e:
            self.logger.debug(
                f"Task '{task.name}' failed, as an exception occurred during '_spawn_and_stream': {e}"
            )
            success = TaskTerminationType.FAILURE
        finally:
            safe_unlink(modified_script)
            safe_unlink(tmpwrap_path)
            flog.close()
            return success

    @typecheck
    def _worker(self, task: SchedulerTask) -> None:
        """
        Run a task in a worker thread.

        :param task: The task to run
        :type task: SchedulerTask
        """
        task_id = self.logger.add_task(task.name, total=1)
        try:
            success = self.run_task(task, task_id)
        except Exception as e:
            self.logger.debug(
                f"Task '{task.name}' failed, as an exception occurred during 'run_task': {e}"
            )
            success = TaskTerminationType.FAILURE
        finally:
            self.logger.finish_task(task_id, success)
            self.logger.debug(
                f"Task '{task.name}' completed {'successfully' if success == TaskTerminationType.SUCCESS else 'with errors'}"
            )
            with self.lock:
                self.results[task] = success
                self.queue.put(task)

    def run(self) -> None:
        """Run all scheduled tasks, respecting dependencies."""
        # Get logger
        self.logger = get_logger()

        running = {}
        while True:
            with self.lock:
                for task in self.tasks:
                    # Skip already running or completed tasks
                    if task in self.results or task in self.scheduled:
                        continue

                    results_to_name = {
                        t.name: res for t, res in self.results.items()
                    }

                    # Skip tasks whose dependencies have failed
                    if any(
                        results_to_name.get(dep, None)
                        in {
                            TaskTerminationType.FAILURE,
                            TaskTerminationType.SKIPPED,
                        }
                        for dep in task.depends_on
                    ):
                        self.results[task] = TaskTerminationType.SKIPPED
                        self.logger.warning(
                            f"Skipping task '{task.name}' because a dependency failed or was skipped"
                        )

                        task_id = self.logger.add_task(task.name, total=1)
                        self.logger.finish_task(
                            task_id, TaskTerminationType.SKIPPED
                        )
                        continue

                    # Start tasks whose dependencies are all met
                    if all(
                        results_to_name.get(dep, None)
                        in {
                            TaskTerminationType.SUCCESS,
                            TaskTerminationType.PARTIAL,
                        }
                        for dep in task.depends_on
                    ):
                        t = Thread(target=self._worker, args=(task,))
                        t.start()
                        running[task] = t
                        self.scheduled.add(task)

            if not running:
                break

            finished = self.queue.get()
            running[finished].join()
            del running[finished]

    def get_results(self) -> list[tuple[str, str, bool]]:
        """
        Get a list of all tasks with their results.

        :return: List of tasks in the format [task_name, task_logfile, successful]
        :rtype: list[tuple[str, str, bool]]
        """
        # Get logger
        logger = get_logger()

        all_tasks = []
        for task, result in self.results.items():
            for _, task_info in logger._task_infos.items():
                if task_info["name"] == task.name:
                    all_tasks.append(
                        (
                            task.name,
                            str(task_info["logfile"]),
                            result == TaskTerminationType.SUCCESS,
                        )
                    )
        return all_tasks
