import json
import logging
from datetime import datetime
from logging import Logger, handlers
from pathlib import Path
from queue import SimpleQueue
from typing import Any, Self

from pydantic import BaseModel

from test_automation_suite.benchmarks import benchmarks_by_name
from test_automation_suite.checks import checks_by_name
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
from test_automation_suite.procedures import procedures_by_name

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 TaskDispatcherState(BaseModel):
    task_name: str
    started_at: datetime
    current_task_index: int = 0
    current_task_state: CurrentTaskState = CurrentTaskState()
    results: list[TaskResult] = []


class UnknownTaskError(Exception):
    def __init__(self, task_name: str) -> None:
        super().__init__(f"Unknown task {task_name}")


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

    Survives multiple starts/reboots.
    """

    _tasks: list[Task]
    _state: TaskDispatcherState

    _task_logger: Logger
    _task_logger_queue: SimpleQueue[logging.LogRecord]

    def __init__(self, task_name: str) -> None:
        """Create a new dispatcher

        Use TaskDispatcher.from_state when loading a state file instead.
        """
        self._tasks = TaskDispatcher._load_tasks_for(task_name)
        self._state = TaskDispatcherState(task_name=task_name, started_at=datetime.now())

        self._task_logger = logging.getLogger(f"{__name__}.task")
        self._task_logger.setLevel(logging.DEBUG)
        self._task_logger_queue = SimpleQueue()
        self._task_logger.addHandler(handlers.QueueHandler(self._task_logger_queue))

    @classmethod
    def from_state(cls, state_file: Path) -> Self:
        """Re-create a dispatcher from a state file"""
        state = TaskDispatcher._load_persistent_state(state_file)
        dispatcher = cls(state.task_name)
        dispatcher._state = state
        return dispatcher

    @staticmethod
    def _load_tasks_for(task_name: str) -> list[Task]:
        if task_name.startswith("Check"):
            return [checks_by_name[task_name]()]
        if task_name.startswith("Benchmark"):
            return [benchmarks_by_name[task_name]()]
        if task_name.startswith("Test"):
            return procedures_by_name[task_name]

        raise UnknownTaskError(task_name)

    @staticmethod
    def _load_persistent_state(state_file: Path) -> TaskDispatcherState:
        assert state_file.is_file(), "cannot load non-existing state file"
        with state_file.open(encoding="utf-8") as f:
            state = json.load(f)
        return TaskDispatcherState.model_validate(state)

    def _drain_task_logger_queue(self) -> list[str]:
        log_formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
        items = []
        while not self._task_logger_queue.empty():
            items.append(log_formatter.format(self._task_logger_queue.get_nowait()))
        return items

    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 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._state.current_task_index < len(self._tasks):
            current_task = self._tasks[self._state.current_task_index]

            logger.info(
                "Running task %s of type %s",
                self._state.current_task_index,
                current_task.__class__.__name__,
            )
            result: TaskResult = self.run_task(current_task, context)

            self._state.results.append(result)
            assert len(self._state.results) == self._state.current_task_index + 1

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

            self._state.current_task_index += 1
            self._state.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._state.current_task_state

        self._task_logger.info("Starting %s", task.__class__.__name__)
        while state.step < len(steps):
            step = steps[state.step]
            self._task_logger.info("Starting step %i (%s)", state.step, step.__name__)

            step_result = step(context, self._task_logger)
            state.append(step_result, self._drain_task_logger_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."""

        try:
            json_state = self._state.model_dump_json()
            with state_file.open(mode="w", encoding="utf-8") as f:
                f.write(json_state)
            logger.info("Dispatcher state saved to %s", state_file)
        except Exception:
            logger.exception("Failed to save dispatcher state.")

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

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

    def results(self) -> dict[str, TaskResult]:
        return {
            f"{i}-{self._tasks[i].__class__.__name__}": self._state.results[i]
            for i in range(len(self._tasks))
        }
