import json
import logging
import re
import sys
from collections.abc import Iterator
from datetime import datetime
from enum import StrEnum
from logging import handlers
from pathlib import Path
from queue import SimpleQueue
from typing import Any, Self

import yaml
from pydantic import BaseModel, ValidationError

from test_automation_suite.benchmarks import benchmarks_by_name
from test_automation_suite.benchmarks.benchmark import (
    Benchmark,
    BenchmarkConfig,
)
from test_automation_suite.checks import checks_by_name
from test_automation_suite.checks.check import Check
from test_automation_suite.core.config import Config
from test_automation_suite.core.context import Context
from test_automation_suite.core.reboot import reboot
from test_automation_suite.core.task import StepResult, Task, TaskResult
from test_automation_suite.dependencies.dependency import Dependency

logger = logging.getLogger(__name__)


class CurrentTaskState(BaseModel):
    """Accumulated state of a task."""

    step: int = 0
    """The current step within the task."""

    failed: bool = False
    """Whether any of the steps have failed."""

    log: list[str] = []
    """Accumulated log of all past steps."""

    harvest_data: dict[str, Any] = {}
    """Accumulated harvest data of all past steps."""

    def append(self, step: StepResult, log: list[str]) -> None:
        """Merges a new step result."""
        self.step += 1
        self.failed |= not step.success
        self.log += log
        self.harvest_data |= step.harvest_data


class TaskType(StrEnum):
    check = "check"
    benchmark = "benchmark"


class UnknownTaskError(Exception):
    task_type: TaskType
    task_name: str

    def __init__(self, task_type: TaskType, task_name: str) -> None:
        super().__init__(f"unknown {task_type} {task_name}")
        self.task_type = task_type
        self.task_name = task_name


class DuplicatedTaskKeyError(Exception):
    key: str

    def __init__(self, key: str) -> None:
        super().__init__(f"duplicated task key: {key}")
        self.key = key


class TaskDispatcher:
    """Stateful dispatcher for running tasks (checks or benchmarks).

    Survives multiple starts/reboots.
    """

    _tasks: list[Task]
    current_task_index: int
    current_task_state: CurrentTaskState
    started_at: datetime
    procedure_name: str | None
    results: dict[str, TaskResult]

    def __init__(self) -> None:
        self._tasks = []
        self.current_task_index = 0
        self.current_task_state = CurrentTaskState()
        self.started_at = datetime.now()
        self.procedure_name = None
        self.results = {}

    def load_persistent_state(self, state_file: Path) -> None:
        # TODO: Move exception handling out of here
        try:
            with state_file.open(encoding="utf-8") as f:
                state = json.load(f)
        except Exception:
            logger.exception("Failed to load dispatcher state")
            sys.exit(1)

        self.current_task_index = int(state.get("current_task_index"))
        self.current_task_state = CurrentTaskState(**state.get("current_task_state"))
        self.started_at = TaskDispatcher._parse_started_at(state.get("started_at"))
        self.procedure_name = state.get("procedure_name", None)
        self.results = {key: TaskResult(**value) for key, value in state.get("results", {}).items()}

        self._tasks = TaskDispatcher._restore_tasks(state.get("tasks", []))
        assert 0 <= self.current_task_index < len(self._tasks), "current_index is out of range"

        logger.info("Loaded dispatcher state from %s", state_file)
        logger.info("%d / %d tasks remaining", self.current_task_index + 1, len(self._tasks))

    @staticmethod
    def _parse_started_at(value: str | None) -> datetime:
        if not isinstance(value, str):
            return datetime.now()

        try:
            return datetime.fromisoformat(value)
        except ValueError:
            return datetime.now()

    @staticmethod
    def _restore_tasks(raw_tasks: list) -> list[Task]:
        tasks = []

        for entry in raw_tasks:
            task_key: str = entry.get("key")
            task_name: str = entry.get("name")
            task_type: TaskType = entry.get("type")
            task_obj = None
            match task_type:
                case TaskType.check:
                    arg: str | None = entry.get("arg")
                    task_obj = checks_by_name[task_name](arg=arg, key=task_key)
                case TaskType.benchmark:
                    config: BenchmarkConfig = entry.get("config")
                    task_obj = benchmarks_by_name[task_name](benchmark_config=config, key=task_key)
                case _:
                    logger.error(
                        "Could not deserialize task '%s', as its instance has an unknown type: %s",
                        task_name,
                        task_type,
                    )
                    continue

            tasks.append(task_obj)

        return tasks

    # TODO: make sure this always runs before `run_tasks`.
    def install_dependencies(self) -> None:
        """
        Install all dependencies from all checks.
        """
        for task in self._tasks:
            logger.info("Installing dependencies for '%s'...", task.__class__.__name__)

            self.install_task_dependencies(task.dependencies())

    def load_procedure(self, config: Config, procedure_file_name: str) -> Self:
        config_file_path = config.procedure_directory / f"{procedure_file_name}.yaml"

        if not config_file_path.is_file():
            logger.fatal(
                "Config file '%s' does not exist, but config_file was submitted.",
                config_file_path,
            )
            sys.exit(1)

        try:
            with config_file_path.open(encoding="utf-8") as f:
                config_data = yaml.safe_load(f) or {}

            raw_checks = config_data.get("checks", [])
            self.append_checks(raw_checks)
            raw_benchmarks = config_data.get("benchmarks", {})
            self.append_benchmarks(raw_benchmarks)
            self.procedure_name = procedure_file_name

            logger.info(
                "Loaded %d tasks from %s",
                len(self._tasks),
                config_file_path,
            )

        except Exception:
            logger.exception("Failed to load YAML config '%s'", config_file_path)
            sys.exit(1)

        return self

    def run_tasks(self, context: Context) -> None:
        """Run all tasks from the current index to the end."""

        self.save_to_file(context.config.state_file)
        while self.current_task_index < len(self._tasks):
            current_task = self._tasks[self.current_task_index]

            logger.info("Running task %s", current_task.key)
            result: TaskResult = self.run_task(current_task, context)

            self.results[current_task.key] = result

            logger.info(
                "Finished task %s %s %s",
                current_task.key,
                "with success" if result.success else "without success",
                f"and harvest_data: {result.harvest_data}" if result.harvest_data else "",
            )

            self.current_task_index += 1
            self.current_task_state = CurrentTaskState()
            self.save_to_file(context.config.state_file)

    def run_task(self, task: Task, context: Context) -> TaskResult:
        """Runs a single task."""

        steps = task.steps()
        state = self.current_task_state

        task_logger = logging.getLogger(f"task-{task.key}")
        task_logger.setLevel(logging.DEBUG)
        log_formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
        log_queue: SimpleQueue[logging.LogRecord] = SimpleQueue()
        task_logger.addHandler(handlers.QueueHandler(log_queue))

        def drain_log_queue() -> list[str]:
            items = []
            while not log_queue.empty():
                items.append(log_formatter.format(log_queue.get_nowait()))
            return items

        task_logger.info("Starting task %s of type %s", task.key, task.__class__.__name__)
        while state.step < len(steps):
            step = steps[state.step]
            task_logger.info("Starting step %i (%s)", state.step, step.__name__)

            step_result = step(context, task_logger)
            state.append(step_result, drain_log_queue())

            if not step_result.success and step_result.abort_on_failure:
                return TaskResult(
                    success=False,
                    log=state.log,
                    harvest_data=state.harvest_data,
                )

            if step_result.needs_reboot:
                self.save_to_file(context.config.state_file)
                reboot()

        return TaskResult(
            success=not state.failed,
            log=state.log,
            harvest_data=state.harvest_data,
        )

    @staticmethod
    def install_task_dependencies(dependencies: list[Dependency]) -> bool:
        """Installs the dependencies of a single task.

        Returns `False` if any dependency fails to install.
        """

        success = True
        for dependency in dependencies:
            if not dependency.install():
                success = False

        return success

    def save_to_file(self, state_file: Path) -> None:
        """Save dispatcher state to JSON."""

        state = {
            "current_task_index": self.current_task_index,
            "current_task_state": self.current_task_state.model_dump(),
            "results": {key: value.model_dump() for key, value in self.results.items()},
            "started_at": self.started_at.isoformat(),
            "procedure_name": self.procedure_name,
            "tasks": self._serialize_tasks(),
        }
        try:
            with state_file.open(mode="w", encoding="utf-8") as f:
                json.dump(state, f, indent=4)
            logger.info("Dispatcher state saved to %s", state_file)
        except Exception:
            logger.exception("Failed to save dispatcher state.")

    def _serialize_tasks(self) -> list:
        serialized_tasks = []
        for task in self._tasks:
            match task:
                case Check():
                    serialized_tasks.append(
                        {
                            "type": TaskType.check,
                            "name": task.__class__.__name__,
                            "key": task.key,
                            "arg": task.arg,
                        },
                    )
                case Benchmark():
                    serialized_tasks.append(
                        {
                            "type": TaskType.benchmark,
                            "name": task.__class__.__name__,
                            "key": task.key,
                            "config": task.benchmark_config,
                        },
                    )
                case _:
                    logger.error(
                        "Could not serialize task '%s', as its instance has an unknown type: %s",
                        task.__class__.__name__,
                        type(task).__name__,
                    )

        return serialized_tasks

    def append_checks(self, raw_checks: list[str] | str) -> Self:
        if isinstance(raw_checks, str):
            raw_checks = [raw_checks]

        for entry in raw_checks:
            for check_name, arg in self.expand_task_entry(entry):
                if check_name not in checks_by_name:
                    logger.error("Unknown check '%s'.", check_name)
                    raise UnknownTaskError(TaskType.check, check_name)

                key = f"{check_name}({arg})" if arg else check_name
                if key in (task.key for task in self._tasks):
                    raise DuplicatedTaskKeyError(key)

                self._tasks.append(checks_by_name[check_name](arg=arg, key=key))

                if arg:
                    logger.debug("Added %s(%s) to tasks.", check_name, arg)
                else:
                    logger.debug("Added %s to tasks.", check_name)

        return self

    def append_benchmarks(self, raw_benchmarks: dict) -> Self:
        for key, details in raw_benchmarks.items():
            try:
                benchmark_config = BenchmarkConfig(**details["config"])
            except ValidationError:
                logger.exception("Invalid configuration of benchmark '%s'", key)
                continue

            benchmark_name = details["benchmark"]
            if benchmark_name not in benchmarks_by_name:
                logger.error("Unknown benchmark '%s'.", key)
                raise UnknownTaskError(TaskType.benchmark, benchmark_name)

            if key in (task.key for task in self._tasks):
                raise DuplicatedTaskKeyError(key)

            self._tasks.append(
                benchmarks_by_name[benchmark_name](
                    benchmark_config=benchmark_config,
                    key=key,
                ),
            )
            logger.debug(
                "Added benchmark to task: '%s' with %s",
                key,
                benchmark_config,
            )

        return self

    def has_remaining_tasks(self) -> bool:
        return self.current_task_index < len(self._tasks)

    def count_tasks(self) -> int:
        return len(self._tasks)

    @staticmethod
    def expand_task_entry(entry: str) -> Iterator[tuple[str, str | None]]:
        """
        Translates a YAML-Entry such as 'CheckNvidiaGaming(igpu,on-demand,dgpu)' into
        ('CheckNvidiaGaming', 'igpu'), ('CheckNvidiaGaming', 'on-demand'), ...
        If there are no arguments, returns: ('CheckNvidiaPrime', None)
        """
        match = re.match(r"^(\w+)(?:\((.*)\))?$", entry)
        if not match:
            msg = f"Invalid task entry: {entry}"
            raise ValueError(msg)

        name = match.group(1)
        args = match.group(2).split(",") if match.group(2) else []

        if args:
            for a in args:
                arg = a.strip()
                if arg:
                    yield name, arg
        else:
            yield name, None
