"""Module for managing pipeline optimization.

This module handles the tracking and logging of optimization artifacts,
trials, and modules during the pipeline's execution.
"""

import json
import logging
import shutil
import tempfile
from collections.abc import Generator
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np
from numpy.typing import NDArray

from autointent._dump_tools import Dumper
from autointent.configs import EmbedderConfig, InferenceNodeConfig
from autointent.custom_types import NodeType

from ._data_models import Artifact, Artifacts, EmbeddingArtifact, ScorerArtifact, Trial, Trials

if TYPE_CHECKING:
    from autointent.modules.base import BaseModule


logger = logging.getLogger(__name__)


@dataclass
class ModulesList:
    """Container for managing the best module for each node type.

    Attributes:
        regex: Best module for the regex node.
        embedding: Best module for the embedding node.
        scoring: Best module for the scoring node.
        decision: Best module for the decision node.
    """

    regex: "BaseModule | None" = None
    embedding: "BaseModule | None" = None
    scoring: "BaseModule | None" = None
    decision: "BaseModule | None" = None

    def get(self, node_type: str) -> "BaseModule | None":
        """Retrieve the module for a specific node type.

        Args:
            node_type: The type of node (e.g., "regex", "embedding").

        Returns:
            Module for the specified node type, or None if not set.
        """
        return getattr(self, node_type)  # type: ignore[no-any-return]

    def add_module(self, node_type: str, module: "BaseModule") -> None:
        """Set the module for a specific node type.

        Args:
            node_type: The type of node.
            module: The module to set.
        """
        setattr(self, node_type, module)


class OptimizationInfo:
    """Tracks optimization results, including trials, artifacts, and modules.

    This class provides methods for logging optimization results, retrieving
    the best-performing modules and artifacts, and generating configuration
    for inference nodes.

    Attributes:
        artifacts: Container for storing optimization artifacts.
        trials: Container for storing optimization trials.
        modules: Container for storing module instances.
        pipeline_metrics: Dictionary storing pipeline-level metrics.
    """

    def __init__(self) -> None:
        """Initialize optimization info."""
        self.artifacts = Artifacts()
        self.trials = Trials()
        self.modules = ModulesList()
        self.pipeline_metrics: dict[str, float] = {}

    def log_module_optimization(
        self,
        node_type: str,
        module_name: str,
        module_params: dict[str, Any],
        metric_value: float,
        metric_name: str,
        metrics: dict[str, float],
        artifact: Artifact,
        module_dump_dir: str | None,
        module: "BaseModule",
    ) -> None:
        """Log optimization results for a module.

        Args:
            node_type: Type of the node being optimized.
            module_name: Type of the module.
            module_params: Parameters of the module for the trial.
            metric_value: Metric value achieved by the module.
            metric_name: Name of the evaluation metric.
            metrics: Dictionary of metric names and their values.
            artifact: Artifact generated by the module.
            module_dump_dir: Directory where the module is dumped.
            module: The module instance, if available.
        """
        is_new_best, old_best_metric_value_idx = self._is_new_best(node_type, metric_value)
        if is_new_best:
            self.modules.add_module(node_type, module)
            if module_dump_dir is not None:
                module.dump(module_dump_dir)
            self.artifacts.add_artifact(node_type, artifact)

            if old_best_metric_value_idx is not None:
                prev_best_dump = self.trials.get_trials(node_type)[old_best_metric_value_idx].module_dump_dir
                if prev_best_dump is not None:
                    shutil.rmtree(prev_best_dump, ignore_errors=True)  # workaround for windows
        else:
            module_dump_dir = None

        trial = Trial(
            module_name=module_name,
            metric_name=metric_name,
            metric_value=metric_value,
            module_params=module_params,
            module_dump_dir=module_dump_dir,
            metrics=metrics,
        )
        self.trials.add_trial(node_type, trial)
        logger.debug("module %s fitted and saved to optimization info %s", module_name, json.dumps(trial.model_dump()))

    def _is_new_best(self, node_type: str, metric_value: float) -> tuple[bool, int | None]:
        """Check if the new trial is the best.

        Args:
            node_type: Type of the node.
            metric_value: Metric value of the trial.
        """
        best_metric_value, best_metric_value_idx = self._get_best_metric_value(node_type)
        return best_metric_value is None or metric_value > best_metric_value, best_metric_value_idx

    def _get_metrics_values(self, node_type: str) -> list[float]:
        """Retrieve all metric values for a node type.

        Args:
            node_type: Type of the node.
        """
        return [trial.metric_value for trial in self.trials.get_trials(node_type)]

    def _get_best_metric_value(self, node_type: str) -> tuple[float, int] | tuple[None, None]:
        """Retrieve the best metric value and index for a node type.

        Args:
            node_type: Type of the node.

        Returns:
            Best metric value and index, or None if no trials exist.
        """
        metric_values = self._get_metrics_values(node_type)
        if not metric_values:
            return None, None
        return max(metric_values), int(np.argmax(metric_values))

    def get_best_embedder(self) -> EmbedderConfig:
        """Retrieve the name of the best embedder from the retriever node.

        Returns:
            Configuration of the best embedder.
        """
        best_retriever_artifact: EmbeddingArtifact = self.artifacts.get_best_artifact(node_type=NodeType.embedding)  # type: ignore[assignment]
        return best_retriever_artifact.config

    def get_best_train_scores(self) -> NDArray[np.float64] | None:
        """Retrieve the train scores from the best scorer node.

        Returns:
            Train scores as a numpy array.
        """
        best_scorer_artifact: ScorerArtifact = self.artifacts.get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
        return best_scorer_artifact.train_scores

    def get_best_validation_scores(self) -> NDArray[np.float64] | None:
        """Retrieve the validation scores from the best scorer node.

        Returns:
            Validation scores as a numpy array.
        """
        best_scorer_artifact: ScorerArtifact = self.artifacts.get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
        return best_scorer_artifact.validation_scores

    def get_best_folded_scores(self) -> list[NDArray[np.float64]] | None:
        """Retrieve the validation scores from the best scorer node.

        Returns:
            Validation scores as a numpy array.
        """
        best_scorer_artifact: ScorerArtifact = self.artifacts.get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
        return best_scorer_artifact.folded_scores

    def get_best_test_scores(self) -> NDArray[np.float64] | None:
        """Retrieve the test scores from the best scorer node.

        Returns:
            Test scores as a numpy array.
        """
        best_scorer_artifact: ScorerArtifact = self.artifacts.get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
        return best_scorer_artifact.test_scores

    def dump_evaluation_results(self) -> dict[str, Any]:
        """Dump evaluation results for all nodes.

        Returns:
            Dictionary containing metrics and configurations for all nodes.
        """
        node_wise_metrics = {node_type: self._get_metrics_values(node_type) for node_type in NodeType}
        return {
            "pipeline_metrics": self.pipeline_metrics,
            "metrics": node_wise_metrics,
            "configs": self.trials.model_dump(),
        }

    def dump(self, path: Path) -> None:
        """Dump the optimization information to a file."""
        exclude = [ModulesList, dict]
        target_path = path / "optimization_info"

        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)

            # Try to save to the temporary directory first
            try:
                Dumper.dump(self, temp_path, exists_ok=False, exclude=exclude)

                # If successful, move the contents to the final destination
                if target_path.exists():
                    shutil.rmtree(target_path)
                shutil.move(temp_path, target_path)

            except Exception:
                logger.exception("Failed to save optimization info")
                raise

    def load(self, path: Path) -> None:
        """Load the optimization information from a file."""
        Dumper.load(self, path / "optimization_info")

    def get_inference_nodes_config(self, asdict: bool = False) -> list[InferenceNodeConfig]:
        """Generate configuration for inference nodes based on the best trials.

        Args:
            asdict: Whether to return the configuration as dictionaries.

        Returns:
            List of configurations for inference nodes.
        """
        res = []
        for node_type, trial, _ in self._get_best_trials():
            if trial is None:
                continue
            item = {
                "node_type": node_type.value,
                "module_name": trial.module_name,
                "module_config": trial.module_params,
                "load_path": trial.module_dump_dir,
            }
            res.append(item if asdict else InferenceNodeConfig(**item))  # type: ignore[arg-type]
        return res  # type: ignore[return-value]

    def _get_best_trials(self) -> Generator[tuple[NodeType, Trial | None, int | None], None, None]:
        """Retrieve the best trials for all node types.

        Yields:
            Tuple of node type, best trial, and best trial index.
        """
        for node_type in NodeType:
            metric_values = self._get_metrics_values(node_type)
            if not metric_values:
                yield node_type, None, None
                continue
            best_trial_idx = int(np.argmax(metric_values))
            best_trial = self.trials.get_trials(node_type)[best_trial_idx]
            yield node_type, best_trial, best_trial_idx

    def get_best_modules(self) -> dict[NodeType, "BaseModule"]:
        """Retrieve the best modules for all node types.

        Returns:
            Dictionary of the best modules for each node type.
        """
        res = {}
        for node_type in NodeType:
            module = self.modules.get(node_type)
            if module is None:
                continue
            res[node_type] = module
        return res
