"""Configuration dataclasses and utilities."""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

import yaml
from dacite import from_dict, Config as DaciteConfig


@dataclass
class DataConfig:
    """Data-related configuration."""

    ud_path: str = "/Users/michal/git/ipi/ud/ud-treebanks-v2.17"
    language: str = "Polish"
    treebanks: list[str] = field(default_factory=lambda: ["LFG", "PDB"])
    max_length: int = 512
    train_batch_size: int = 16
    eval_batch_size: int = 32
    num_workers: int = 4
    max_train_samples: Optional[int] = None  # Limit training sentences
    max_eval_samples: Optional[int] = None  # Limit dev/eval sentences


@dataclass
class ModelConfig:
    """Model architecture configuration."""

    base_model: str = "xlm-roberta-large"
    use_lora: bool = False
    lora_r: int = 64
    lora_alpha: int = 128
    lora_dropout: float = 0.05
    lora_target_modules: list[str] = field(default_factory=lambda: ["query", "value"])

    # Head configurations
    arc_hidden_dim: int = 256
    rel_hidden_dim: int = 128
    char_embed_dim: int = 300
    max_char_len: int = 50
    dropout: float = 0.1

    # CNN lemmatizer
    lemma_input_proj_dim: int = 32
    lemma_input_proj_dropout: float = 0.25
    lemma_cnn_filters: list[int] = field(default_factory=lambda: [256, 256, 256])
    lemma_cnn_kernel_sizes: list[int] = field(default_factory=lambda: [3, 3, 3, 1])
    lemma_cnn_dilations: list[int] = field(default_factory=lambda: [1, 2, 4, 1])
    lemma_cnn_paddings: list[int] = field(default_factory=lambda: [1, 2, 4, 0])

    # GRU refinement
    lemma_gru_hidden_dim: int = 512

    # Loss weights
    morpho_weight: float = 1.0
    feats_weight: float = 3.0
    arc_weight: float = 1.0
    rel_weight: float = 1.0
    lemma_weight: float = 1.5
    use_uncertainty_weighting: bool = False

    # Morpho-to-lemma conditioning
    morpho_to_lemma_proj_dim: int = 64

    # Beam search for lemma decoding
    lemma_beam_width: int = 5
    use_lemma_beam_search: bool = False

    # Morphological constraint filtering
    use_morpho_constraints: bool = False

    # These are set dynamically based on data
    num_upos: int = 0
    num_xpos: int = 0
    feat_vocab_sizes: dict[str, int] = field(default_factory=dict)
    feat_categories: list[str] = field(default_factory=list)
    num_deprels: int = 0
    char_vocab_size: int = 512


@dataclass
class TrainingConfig:
    """Training configuration."""

    epochs: int = 20
    learning_rate: float = 2e-5
    encoder_lr: float = 2e-5
    head_lr: float = 1e-3
    weight_decay: float = 0.01
    warmup_ratio: float = 0.1
    gradient_clip: float = 1.0
    gradient_accumulation_steps: int = 1
    gradient_checkpointing: bool = False

    # Checkpointing
    save_every_epoch: bool = True
    save_best_only: bool = False
    metric_for_best: str = "MLAS"
    best_model_split: str = "dev"  # Which split to use for best model selection: "train", "dev", or "test"

    # Evaluation
    eval_on_train: bool = True
    train_eval_samples: Optional[int] = 2000


@dataclass
class WandbConfig:
    """Weights & Biases configuration."""

    enabled: bool = True
    project: str = "combo-nlp"
    entity: Optional[str] = None
    run_name: Optional[str] = None
    tags: list[str] = field(default_factory=list)


@dataclass
class NamingConfig:
    """Model naming configuration for HuggingFace Hub."""

    project: str = "clarin-pl"
    models_prefix: str = "combo-nlp"
    base_model_short: str = "herbert-base-cased"
    ud_version: str = "2.17"


@dataclass
class Config:
    """Root configuration."""

    data: DataConfig = field(default_factory=DataConfig)
    model: ModelConfig = field(default_factory=ModelConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)
    wandb: WandbConfig = field(default_factory=WandbConfig)
    naming: NamingConfig = field(default_factory=NamingConfig)

    output_dir: str = "data/output"
    seed: int = 42
    device: str = "auto"  # "auto", "cuda", "mps", "cpu"
    max_length: int = 512  # Also present in data.max_length; exported models store it here

    @property
    def model_name(self) -> str:
        """Get the full model name derived from naming config and data config.

        Example: combo-nlp-herbert-base-cased-polish-pdb-ud2.17
        """
        treebanks = "-".join(t.lower() for t in self.data.treebanks)
        language = self.data.language.lower().replace("_", "-")
        return f"{self.naming.models_prefix}-{self.naming.base_model_short}-{language}-{treebanks}-ud{self.naming.ud_version}"

    @property
    def repo_id(self) -> str:
        """Get the full HuggingFace repo ID.

        Example: clarin-pl/combo-nlp-herbert-base-cased-polish-pdb-ud2.17
        """
        return f"{self.naming.project}/{self.model_name}"

    @property
    def language_output_dir(self) -> Path:
        """Get model-specific output directory."""
        return Path(self.output_dir) / self.model_name


def deep_merge(base: dict, override: dict) -> dict:
    """Deep merge two dictionaries, with override taking precedence."""
    result = base.copy()
    for key, value in override.items():
        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
            result[key] = deep_merge(result[key], value)
        else:
            result[key] = value
    return result


def _load_config_file(path: Path) -> dict:
    """Load a config file (YAML or JSON)."""
    with open(path) as f:
        if path.suffix == ".json":
            import json
            return json.load(f)
        return yaml.safe_load(f)


def load_config(config_path: Path, task_config_path: Optional[Path] = None) -> Config:
    """Load configuration from YAML or JSON files.

    Args:
        config_path: Path to base configuration file.
        task_config_path: Optional path to task-specific overrides.

    Returns:
        Merged configuration object.
    """
    config_dict = _load_config_file(config_path)

    if task_config_path and task_config_path.exists():
        with open(task_config_path) as f:
            task_config = yaml.safe_load(f)
        config_dict = deep_merge(config_dict, task_config)

    dacite_config = DaciteConfig(strict=True)
    return from_dict(data_class=Config, data=config_dict, config=dacite_config)


def merge_configs(base_dict: dict[str, Any], override_dict: dict[str, Any]) -> dict[str, Any]:
    """Merge two config dictionaries."""
    return deep_merge(base_dict, override_dict)
