"""
Configuration for Video Collision Anticipation (Inference Only)
===============================================================
Standalone configuration - no PyTorch Lightning dependency.
"""

import os
import re
from dataclasses import asdict, dataclass, field
from typing import Any, Dict, Optional, TYPE_CHECKING

import yaml

if TYPE_CHECKING:
    from badas.inference.smoothing import SmoothingConfig
    from badas.inference.debug import DebugConfig
    from badas.inference.risk import RiskConfig


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

    model_name: str = "facebook/vjepa2-vitl-fpc16-256-ssv2"
    num_classes: int = 2
    freeze_backbone: bool = False
    use_classification_model: bool = False
    use_custom_head: bool = False
    use_frame_level: bool = False
    use_model_frame_count: bool = False

    # Future prediction (V-JEPA 2)
    use_future_prediction: bool = True
    predictor_combination_method: str = "concat"
    future_prediction_seconds: float = 1.0
    original_fps: int = 4

    # Temporal processing
    temporal_method: str = "attention"
    temporal_hidden_dim: int = 512
    temporal_num_heads: int = 8
    temporal_num_queries: int = 8

    # Classification head
    head_type: str = "mlp"
    head_hidden_dim: int = 768
    head_num_layers: int = 3
    head_num_heads: int = 8
    head_dropout: float = 0.1

    # Training enhancements
    force_num_classes: bool = False
    use_temperature_scaling: bool = False
    temperature: float = 2.0
    drop_path_rate: float = 0.1

    @property
    def is_vjepa2(self):
        return "vjepa2" in self.model_name.lower()

    @property
    def is_vjepa(self):
        return "vjepa" in self.model_name.lower()

    @property
    def is_videomae(self):
        return "videomae" in self.model_name.lower()

    @property
    def frame_count(self):
        frame_match = re.search(r"fpc(\d+)", self.model_name)
        if frame_match:
            return int(frame_match.group(1))

    @property
    def crop_size(self):
        crop_match = re.search(r"-(\d+)(?:$|[^0-9])", self.model_name)
        if crop_match:
            return int(crop_match.group(1))


@dataclass
class DataConfig:
    """Data configuration"""

    data_root: str = "."
    val_data_root: str = None
    split: list = field(default_factory=lambda: ["train"])
    frame_count: int = 16
    img_size: int = 224
    vjepa2_crop_size: int = 256
    num_workers: int = 6
    dataset_fps: int = 8
    dataset_window_size: int = 16
    frame_sampling_method: str = None
    oversample: bool = True
    use_augmentations: bool = True


@dataclass
class TrainingConfig:
    """Training parameters (kept for checkpoint compatibility)"""

    batch_size: int = 8
    val_batch_size: Optional[int] = None
    max_steps: int = 100_000
    val_check_interval: int = 1_000
    log_every_n_steps: int = 10
    trainer_kwargs: dict = field(default_factory=lambda: {})
    learning_rate: float = 1e-5
    weight_decay: float = 1e-4
    scheduler_step_size: int = 600
    scheduler_T_max: int = 2
    scheduler_eta_min: float = 1e-6
    early_stopping_patience: int = 3
    early_stopping_min_delta: float = 0.001
    early_stopping_metric: str = "val_acc"
    early_stopping_mode: str = "max"
    disable_early_stopping: bool = False
    # Step-based checkpointing (for compatibility)
    checkpoint_every_n_steps: Optional[int] = None
    checkpoint_every_n_epochs: int = 1


@dataclass
class SystemConfig:
    """System configuration (kept for checkpoint compatibility)"""

    seed: int = 42
    precision: str = "16-mixed"
    devices: int = 1
    accelerator: str = "auto"
    use_ffmpeg: bool = True


@dataclass
class LoggingConfig:
    """Logging configuration (kept for checkpoint compatibility)"""

    output_dir: str = ""
    s3_bucket: Optional[str] = None
    s3_prefix: Optional[str] = None
    use_wandb: bool = False
    wandb_project: str = "collision-anticipation"
    wandb_entity: Optional[str] = None
    run_name: Optional[str] = None
    experiment_name: str = None
    wandb_tags: list = field(default_factory=list)


@dataclass
class Config:
    """Complete configuration"""

    model: ModelConfig = field(default_factory=ModelConfig)
    data: DataConfig = field(default_factory=DataConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)
    system: SystemConfig = field(default_factory=SystemConfig)
    logging: LoggingConfig = field(default_factory=LoggingConfig)

    def save(self, path: str):
        """Save configuration to YAML file"""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w") as f:
            yaml.dump(asdict(self), f, default_flow_style=False, indent=2)

    @classmethod
    def load(cls, path: str) -> "Config":
        """Load configuration from YAML file"""
        with open(path, "r") as f:
            data = yaml.safe_load(f)
        return cls.load_from_dict(data)

    @classmethod
    def load_from_dict(cls, data: dict) -> "Config":
        """Load configuration from dictionary, filtering unknown fields"""
        def filter_fields(config_cls, config_data):
            """Filter config data to only include known fields"""
            if not config_data:
                return {}
            known_fields = {f.name for f in config_cls.__dataclass_fields__.values()}
            return {k: v for k, v in config_data.items() if k in known_fields}

        return cls(
            model=ModelConfig(**filter_fields(ModelConfig, data.get("model", {}))),
            data=DataConfig(**filter_fields(DataConfig, data.get("data", {}))),
            training=TrainingConfig(**filter_fields(TrainingConfig, data.get("training", {}))),
            system=SystemConfig(**filter_fields(SystemConfig, data.get("system", {}))),
            logging=LoggingConfig(**filter_fields(LoggingConfig, data.get("logging", {}))),
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to dictionary"""
        return asdict(self)


@dataclass
class BADASConfig:
    """Configuration for BADAS inference."""

    # Output control
    quiet: bool = False
    show_progress: bool = False

    # Features (use factory functions to avoid circular imports)
    smoothing: 'SmoothingConfig' = field(default_factory=lambda: _default_smoothing_config())
    debug: 'DebugConfig' = field(default_factory=lambda: _default_debug_config())
    risk: 'RiskConfig' = field(default_factory=lambda: _default_risk_config())

    # Backend selection (None = auto-detect from file extension)
    backend: Optional[str] = None  # 'pytorch', 'onnx', 'tensorrt'

    # PyTorch backend options
    use_compile: bool = True
    use_autocast: bool = False

    # Model parameters.
    # ONNX/TRT: required (no checkpoint to read from); defaults to 16.
    # PyTorch:  None = read from checkpoint; set to override.
    num_frames: Optional[int] = None
    model_fps: int = 8
    model_img_size: int = 256

    # Override the backbone path stored in the checkpoint.  Useful when
    # deploying on a machine where the original training path is unavailable.
    # Must be the same architecture (e.g. ViT-S) as the training backbone.
    # Example: BADASConfig(model_name_override="nexar-ai/vjepa2-vits-fpc16-256-nexar")
    model_name_override: Optional[str] = None

    # Startup ramp: yield predictions before the full window is accumulated,
    # stepping every stride (rounded up to even) frames during cold start.
    startup_ramp: bool = False

    # GPU video decoder (decord GPU decode, with CPU fallback)
    # Disabled by default: GPU saves ~200ms but deadlocks on malformed EOF videos
    use_gpu_decoder: bool = False

    # VLM backend options (used when backend='vlm' or path has no file extension)
    vlm_max_pixels: int = 256 * 256          # cap vision tokens per frame
    vlm_gpu_memory_utilization: float = 0.85  # vLLM KV-cache budget

    # Gemini backend options
    # temperatures: list of sampling temps; probs are averaged across all of them.
    # Use [0.0] for greedy (fast), [0.0, 0.3, 0.7] for calibrated probability.
    gemini_temperatures: tuple = (0.0, 0.3, 0.7)
    # max_workers: thread pool size for concurrent temperature + window-level calls.
    # Set high enough to support window-parallel callers (predict_new_datasets --window-workers).
    gemini_max_workers: int = 32

    @classmethod
    def silent(cls) -> 'BADASConfig':
        """Silent mode (no console output)."""
        return cls(quiet=True)

    @classmethod
    def with_debug(cls, output_dir: str = "debug_windows") -> 'BADASConfig':
        """Enable debug mode to save inference windows as videos."""
        from badas.inference.debug import DebugConfig
        return cls(debug=DebugConfig.basic(output_dir))


def _default_smoothing_config():
    """Factory function to create default SmoothingConfig (avoids circular import)."""
    from badas.inference.smoothing import SmoothingConfig
    return SmoothingConfig()


def _default_debug_config():
    """Factory function to create default DebugConfig (avoids circular import)."""
    from badas.inference.debug import DebugConfig
    return DebugConfig()


def _default_risk_config():
    """Factory function to create default RiskConfig (avoids circular import)."""
    from badas.inference.risk import RiskConfig
    return RiskConfig()
