# -*- coding: utf-8 -*-


import os
import resource
from typing import Any, Literal

import torch
from dfine.core import YAMLConfig
from dfine.misc import dist_utils
from dfine.solver import DetSolver
from pydantic.dataclasses import dataclass
from sinapsis_core.data_containers.data_packet import DataContainer
from sinapsis_generic_data_tools.helpers.file_downloader import download_file

from sinapsis_dfine.templates.dfine_base import DFINEBase, DFINEBaseAttributes


@dataclass
class TrainingArtifacts:
    """
    Dataclass to contain the best_model, last_model and log_file paths produced during the D-Fine model training.
    """

    best_model: str
    last_model: str
    log_file: str


class DFINETrainingAttributes(DFINEBaseAttributes):
    """Attributes for the D-FINE training workflow. Extend the DFINEBaseAttributes.

    Attributes:
        training_mode (Literal["scratch", "tune"]): Specifies the training mode.
            'scratch' for training the model from scratch; 'tune' for fine-tuning the model
            with provided or downloaded weights.
        seed (int | None): Seed for reproducibility. Defaults to None.
        use_amp (bool): Whether to enable Automatic Mixed Precision (AMP). Defaults to False.
        print_rank (int): Rank of the process for printing logs in distributed setups. Defaults
            to 0.
        print_method (Literal["builtin", "rich"]): Method used for printing logs. Defaults to
            "builtin".
    """

    training_mode: Literal["scratch", "tune"]
    seed: int | None = None
    use_amp: bool = False
    print_rank: int = 0
    print_method: Literal["builtin", "rich"] = "builtin"


class DFINETraining(DFINEBase):
    """This module implements the training pipeline for the D-FINE model.

    It includes logic for initializing configurations, downloading weights,
    and setting up the training solver.

    Usage example:

        agent:
        name: my_test_agent
        templates:
        - template_name: InputTemplate
        class_name: InputTemplate
        attributes: {}
        - template_name: DFINETraining
        class_name: DFINETraining
        template_input: InputTemplate
        attributes:
            config_file: '/path/to/config/file/for/dfine'
            pretrained_model:
                size: 'n'
                variant: 'coco'
            device: 'cuda'
            weights_path: null
            output_dir: /sinapsis/cache/dir
            training_mode: 'tune'
            use_amp: false
            print_rank: 0
            print_method: builtin

    """

    AttributesBaseModel = DFINETrainingAttributes

    def initialize(self) -> None:
        """Initializes the template's common state for creation or reset.

        This method is called by both `__init__` and `reset_state` to ensure
        a consistent state. Can be overriden by subclasses for specific behaviour.
        """
        super().initialize()
        self.is_dist = self._initialize_distributed_training()
        self._validate_training_attributes()
        self._adjust_file_descriptor_limit()
        self._set_seed()
        self.cfg = self._initialize_config()
        self._prepare_hgnetv2_weights()
        self.solver = DetSolver(self.cfg)

    def _initialize_distributed_training(self) -> bool:
        """Initializes distributed training if multiple GPUs are available.

        Returns:
            bool: True if distributed training is initialized, False otherwise.
        """
        if torch.cuda.device_count() > 1:
            dist_utils.setup_distributed(
                self.attributes.print_rank,
                self.attributes.print_method,
                seed=self.attributes.seed,
            )
            return True
        return False

    @staticmethod
    def _adjust_file_descriptor_limit(target_limit: int = 65535) -> None:
        """Adjusts the file descriptor limit for the training process.

        Args:
            target_limit (int): Desired limit for file descriptors. Defaults to 65535.
        """
        soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_NOFILE)
        new_limit = min(target_limit, hard_limit)

        if soft_limit < new_limit:
            resource.setrlimit(resource.RLIMIT_NOFILE, (new_limit, hard_limit))

    def _validate_training_attributes(self) -> None:
        """Validates the attributes for training workflows."""
        self._validate_config_file()

        if self.attributes.training_mode == "scratch":
            self._log_ignored_weights_in_scratch_mode()
        elif self.attributes.training_mode == "tune":
            self._validate_tune_mode()

    def _log_ignored_weights_in_scratch_mode(self) -> None:
        """Logs a warning if pretrained_model or weights are provided in 'scratch' mode."""
        if self.attributes.weights_path or self.attributes.pretrained_model.variant != "custom":
            self.logger.warning(
                "In 'scratch' mode, any provided 'weights_path' or standard 'pretrained_model' will be ignored."
            )

    def _validate_tune_mode(self) -> None:
        """Validates attributes for 'tune' mode.

        Raises:
            ValueError: If required weights_path or pretrained_model is missing.
        """
        is_custom = self.attributes.pretrained_model.variant == "custom"
        if is_custom and not self.attributes.weights_path:
            raise ValueError("For 'tune' mode with a 'custom' variant, 'weights_path' must be provided.")
        if not is_custom and self.attributes.weights_path:
            self.logger.warning("'weights_path' is ignored when a standard pretrained model is selected for tuning.")
        if not is_custom:
            self._validate_pretrained_model()

    def _initialize_config(self) -> YAMLConfig:
        """Prepares and initializes the training configuration.

        Returns:
            YAMLConfig: The initialized training configuration object.
        """
        update_dict = self._build_solver_update_dict()
        update_dict = self._prepare_dfine_weights(update_dict)
        cfg = YAMLConfig(self.attributes.config_file, **update_dict)
        cfg.yaml_cfg[self.KEYS.HGNET_V2]["local_model_dir"] = self.attributes.output_dir
        if self.attributes.training_mode == "tune":
            cfg.yaml_cfg[self.KEYS.HGNET_V2]["pretrained"] = False
        return cfg

    def _prepare_hgnetv2_weights(self) -> None:
        """Ensures the HGNetv2 backbone weights are downloaded if necessary.

        Raises:
            ValueError: If the HGNetv2 backbone name is not found in the configuration.
        """
        if not (self.attributes.training_mode == "tune" and self.attributes.pretrained_model.variant == "custom"):
            hgnetv2_cfg = self.cfg.yaml_cfg.get(self.KEYS.HGNET_V2, {})
            hgnetv2_name = hgnetv2_cfg.get("name")
            if hgnetv2_name not in self.SUPPORTED_HGNET_BACKBONES:
                raise ValueError("HGNetv2 model name not found in configuration.")
            self._download_hgnetv2_weights(hgnetv2_name)

    def _build_solver_update_dict(self) -> dict[str, Any]:
        """Builds the dictionary to update solver configuration.

        Returns:
            dict[str, Any]: Dictionary with updated solver parameters.
        """
        return {
            self.KEYS.OUTPUT_DIR: self.attributes.output_dir,
            self.KEYS.DEVICE: self.attributes.device,
            self.KEYS.USE_AMP: self.attributes.use_amp,
            self.KEYS.SEED: self.attributes.seed,
        }

    def _prepare_dfine_weights(self, update_dict: dict[str, Any]) -> dict[str, Any]:
        """Ensures that the D-FINE weights are downloaded if necessary and updates the solver configuration.

        Args:
            update_dict (dict[str, Any]): Initial solver configuration dictionary

        Returns:
            dict[str, Any]: Updated solver configuration dictionary
        """
        if self.attributes.training_mode == "tune":
            if self.attributes.pretrained_model.variant == "custom":
                if self.attributes.weights_path:
                    update_dict[self.KEYS.TUNING] = self.attributes.weights_path
            else:
                dfine_weights_path = self._download_dfine_weights()
                update_dict[self.KEYS.TUNING] = dfine_weights_path
        return update_dict

    def _download_hgnetv2_weights(self, hgnetv2_name: Literal["B0", "B1", "B2", "B3", "B4", "B5", "B6"]) -> None:
        """Downloads HGNetv2 backbone weights.

        Args:
            hgnetv2_name (str): Name of the HGNetv2 backbone to download.
        """
        hgnetv2_weights_filename = f"PPHGNetV2_{hgnetv2_name}_stage1.pth"
        hgnetv2_weights_path = os.path.join(self.attributes.output_dir, hgnetv2_weights_filename)
        hgnetv2_weights_url = self.WEIGHTS_BASE_URL + hgnetv2_weights_filename
        download_file(hgnetv2_weights_url, hgnetv2_weights_path, f"HGNetV2_{hgnetv2_name} weights")

    def _set_seed(self) -> None:
        """Sets the random seed for reproducibility."""
        if self.attributes.seed is not None:
            torch.manual_seed(self.attributes.seed)
            torch.cuda.manual_seed_all(self.attributes.seed)
            self.logger.info(f"Seed set to: {self.attributes.seed}")

    def _cleanup(self) -> None:
        """Cleans up the training-specific stateful objects.

        This method reaches into the DetSolver to move its core PyTorch model
        components to the CPU before deleting the solver and config objects. This
        ensures all GPU resources are safely released.
        """
        if hasattr(self, "solver") and self.solver is not None:
            if hasattr(self.solver, "model") and self.solver.model is not None:
                self.solver.model.to("cpu")

            if hasattr(self.solver, "criterion") and self.solver.criterion is not None:
                self.solver.criterion.to("cpu")

            if hasattr(self.solver, "postprocessor") and self.solver.postprocessor is not None:
                self.solver.postprocessor.to("cpu")

            if hasattr(self.solver, "ema") and self.solver.ema is not None:
                self.solver.ema.to("cpu")
                
            del self.solver

        if hasattr(self, "cfg"):
            del self.cfg

    def execute(self, container: DataContainer) -> DataContainer:
        """Executes the D-fine model training.

        Args:
            container (DataContainer): Input data container.

        Returns:
            DataContainer: The data container containing the model training artifacts.
        """
        self.solver.fit()

        best_model_path = os.path.join(self.attributes.output_dir, "best_stg2.pth")
        if not os.path.exists(best_model_path):
            best_model_path = os.path.join(self.attributes.output_dir, "best_stg1.pth")

            training_artifats = TrainingArtifacts(
                best_model=best_model_path,
                last_model=os.path.join(self.attributes.output_dir, "last.pth"),
                log_file=os.path.join(self.attributes.output_dir, "log.txt"),
            )

            self._set_generic_data(container, training_artifats)

        if self.is_dist:
            dist_utils.cleanup()

        return container
