"""
Configuration for DummyTraining.

This module defines the configuration class for the DummyTraining step,
which is responsible for taking a pretrained model.tar.gz file and
hyperparameters.json, adding the hyperparameters to the model, and making
it available for downstream packaging and registration steps.
"""

from pydantic import Field, model_validator
from typing import TYPE_CHECKING, Optional, Dict, Any
from pathlib import Path

from .config_processing_step_base import ProcessingStepConfigBase
from ...core.base.hyperparameters_base import ModelHyperparameters

# Import the script contract
from ..contracts.dummy_training_contract import DUMMY_TRAINING_CONTRACT

# Import for type hints only
if TYPE_CHECKING:
    from ...core.base.contract_base import ScriptContract


class DummyTrainingConfig(ProcessingStepConfigBase):
    """
    Configuration for DummyTraining step.

    This step takes a pretrained model.tar.gz file, adds hyperparameters.json
    to it, and makes it available for downstream packaging and registration steps,
    bypassing the actual training process.
    """

    # Override with specific default for this step
    processing_entry_point: str = Field(
        default="dummy_training.py",
        description="Entry point script for dummy training.",
    )

    # Unique to this step
    pretrained_model_path: str = Field(
        default="", description="Local path to pretrained model.tar.gz file."
    )

    # Add hyperparameters - compatible with both XGBoost and PyTorch models
    hyperparameters: ModelHyperparameters = Field(
        default_factory=ModelHyperparameters,
        description="Model hyperparameters to be included in the model package.",
    )

    # Add hyperparameters_s3_uri for specifying the S3 location
    hyperparameters_s3_uri: Optional[str] = Field(
        default=None, description="S3 URI where hyperparameters.json will be saved."
    )

    model_config = ProcessingStepConfigBase.model_config

    @model_validator(mode="after")
    def validate_config(self) -> "DummyTrainingConfig":
        """
        Validate configuration and ensure defaults are set.

        Raises:
            ValueError: If configuration is invalid
            FileNotFoundError: If pretrained_model_path is provided but file doesn't exist

        Returns:
            Self with validated configuration
        """
        # Basic validation
        if not self.processing_entry_point:
            raise ValueError("DummyTraining step requires a processing_entry_point")

        # Check for pretrained model path
        if not self.pretrained_model_path:
            raise ValueError("pretrained_model_path is required in DummyTrainingConfig")

        # Check if file exists (if path is concrete and not a variable)
        if not hasattr(self.pretrained_model_path, "expr"):
            model_path = Path(self.pretrained_model_path)
            if not model_path.exists():
                raise FileNotFoundError(f"Pretrained model not found at {model_path}")

            # Check file extension
            if not model_path.suffix == ".tar.gz" and not str(model_path).endswith(
                ".tar.gz"
            ):
                self._logger.warning(
                    f"Model file {model_path} does not have .tar.gz extension"
                )

        # Validate hyperparameters
        if not self.hyperparameters:
            raise ValueError(
                "Model hyperparameters are required in DummyTrainingConfig"
            )

        # If hyperparameters_s3_uri is not provided, it will be generated in the builder

        # Validate script contract - this will be the source of truth
        contract = self.get_script_contract()
        if not contract:
            raise ValueError("Failed to load script contract")

        if "pretrained_model_path" not in contract.expected_input_paths:
            raise ValueError(
                "Script contract missing required input path: pretrained_model_path"
            )

        if "hyperparameters_s3_uri" not in contract.expected_input_paths:
            raise ValueError(
                "Script contract missing required input path: hyperparameters_s3_uri"
            )

        if "model_input" not in contract.expected_output_paths:
            raise ValueError(
                "Script contract missing required output path: model_input"
            )

        return self

    def get_script_contract(self) -> "ScriptContract":
        """
        Get script contract for this configuration.

        Returns:
            The DummyTraining script contract
        """
        return DUMMY_TRAINING_CONTRACT

    def get_script_path(self) -> str:
        """
        Get script path with priority order:
        1. Use processing_entry_point if provided
        2. Fall back to script_contract.entry_point if available

        Always combines with effective source directory.

        Returns:
            Script path or None if no entry point can be determined
        """
        # Determine which entry point to use
        entry_point = None

        # First priority: Use processing_entry_point if provided
        if self.processing_entry_point:
            entry_point = self.processing_entry_point
        # Second priority: Use contract entry point
        else:
            contract = self.get_script_contract()
            if contract and hasattr(contract, "entry_point"):
                entry_point = contract.entry_point

        if not entry_point:
            return None

        # Get the effective source directory
        effective_source_dir = self.get_effective_source_dir()
        if not effective_source_dir:
            return entry_point  # No source dir, just return entry point

        # Combine source dir with entry point
        if effective_source_dir.startswith("s3://"):
            full_path = f"{effective_source_dir.rstrip('/')}/{entry_point}"
        else:
            full_path = str(Path(effective_source_dir) / entry_point)

        return full_path
