# SPDX-License-Identifier: Apache-2.0
# Copyright (C) 2024-2026 SYMFLUENCE Team <dev@symfluence.org>

"""
Xinanjiang Model Configuration.

Pydantic schema and configuration adapter for Xinanjiang model settings.
"""

from typing import Any, Dict, Literal, Optional

from pydantic import BaseModel, Field

from symfluence.models.base.base_config import AutoGeneratedConfigAdapter


class XinanjiangConfig(BaseModel):
    """Xinanjiang model configuration schema."""

    warmup_days: int = Field(
        default=365, ge=0,
        description="Number of warmup days for model spinup."
    )

    pet_method: Literal['input', 'hamon'] = Field(
        default='input',
        description="PET method. 'input' uses provided PET, 'hamon' computes from temperature."
    )

    latitude: Optional[float] = Field(
        default=None, ge=-90.0, le=90.0,
        description="Catchment latitude for PET calculations."
    )

    backend: Literal['jax', 'numpy'] = Field(
        default='jax',
        description="Computation backend. 'jax' enables autodiff; 'numpy' is fallback."
    )

    calibration_metric: Literal['KGE', 'NSE'] = Field(
        default='KGE',
        description="Objective function for calibration."
    )

    params_to_calibrate: str = Field(
        default='all',
        description="Comma-separated list of parameters to calibrate, or 'all' for all 15."
    )

    save_states: bool = Field(
        default=False,
        description="Save internal state variables to output."
    )

    snow_module: Literal['none', 'snow17'] = Field(
        default='none',
        description="Snow module to couple. 'none' for standalone XAJ, 'snow17' for Snow-17 coupling."
    )

    si: float = Field(
        default=100.0, ge=0.0,
        description="SWE threshold for 100% areal coverage (mm). Used when snow_module='snow17'."
    )

    model_config = {'extra': 'forbid'}


class XinanjiangConfigAdapter(AutoGeneratedConfigAdapter):
    """Configuration adapter for Xinanjiang model."""

    MODEL_NAME = "XINANJIANG"
    CONFIG_PREFIX = "XINANJIANG_"

    @classmethod
    def get_config_schema(cls) -> type:
        return XinanjiangConfig

    @classmethod
    def get_defaults(cls) -> Dict[str, Any]:
        return {
            'XINANJIANG_WARMUP_DAYS': 365,
            'XINANJIANG_PET_METHOD': 'input',
            'XINANJIANG_BACKEND': 'jax',
            'XINANJIANG_CALIBRATION_METRIC': 'KGE',
            'XINANJIANG_PARAMS_TO_CALIBRATE': 'all',
            'XINANJIANG_SAVE_STATES': False,
            'XINANJIANG_SNOW_MODULE': 'none',
            'XINANJIANG_SI': 100.0,
        }

    @classmethod
    def _get_legacy_transformers(cls) -> Dict[str, tuple]:
        return {
            'XINANJIANG_WARMUP_DAYS': ('warmup_days', int),
            'XINANJIANG_PET_METHOD': ('pet_method', str),
            'XINANJIANG_LATITUDE': ('latitude', float),
            'XINANJIANG_BACKEND': ('backend', str),
            'XINANJIANG_CALIBRATION_METRIC': ('calibration_metric', str),
            'XINANJIANG_PARAMS_TO_CALIBRATE': ('params_to_calibrate', str),
            'XINANJIANG_SAVE_STATES': ('save_states', bool),
            'XINANJIANG_SNOW_MODULE': ('snow_module', str),
            'XINANJIANG_SI': ('si', float),
        }

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> XinanjiangConfig:
        transformers = cls._get_legacy_transformers()
        defaults = cls.get_defaults()
        kwargs = {}
        for config_key, (field_name, field_type) in transformers.items():
            value = config_dict.get(config_key, defaults.get(config_key))
            if value is not None:
                try:
                    kwargs[field_name] = field_type(value)
                except (ValueError, TypeError):
                    kwargs[field_name] = value
        return XinanjiangConfig(**kwargs)

    @classmethod
    def to_dict(cls, config: XinanjiangConfig) -> Dict[str, Any]:
        transformers = cls._get_legacy_transformers()
        result = {}
        for config_key, (field_name, _) in transformers.items():
            if hasattr(config, field_name):
                result[config_key] = getattr(config, field_name)
        return result
