"""
Configuration Management for Video Collision Anticipation
========================================================

Clean configuration system using YAML files with command-line overrides.
"""

import argparse
import io
import os
import re
from copy import deepcopy
from dataclasses import asdict, dataclass, field
from itertools import product
from typing import Any, Dict, Optional

import pytorch_lightning as pl
import yaml

from badas.training.utils import generate_unique_run_name


@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

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

    # Temporal processing
    temporal_method: str = "attention"  # mean, max, lstm, attention, probe
    temporal_hidden_dim: int = 512
    temporal_num_heads: int = 8
    temporal_num_queries: int = 8

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

    # Partial freeze: when freeze_backbone=True, unfreeze this many last transformer blocks.
    # 0 = fully frozen backbone (head-only); N > 0 = last N blocks + head are trainable.
    unfreeze_last_n_blocks: int = 0

    # Layer extraction: which hidden state index to use for classification.
    # None  = use last_hidden_state (default, current behaviour).
    # 0     = embedding layer output.
    # 1-23  = output of transformer block N (pre-final layernorm).
    # 24    = post-final layernorm = last_hidden_state (same as None).
    # When set, future-prediction augmentation is automatically disabled because
    # the predictor was not trained on intermediate-layer representations.
    extraction_layer: Optional[int] = None

    # Training enhancements
    force_num_classes: bool = False
    use_temperature_scaling: bool = False
    temperature: float = 2.0
    drop_path_rate: float = 0.1
    gradient_checkpointing: bool = False
    attn_implementation: Optional[str] = None  # sdpa, flash_attention_2, or None (eager)

    @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))

    @property
    def processor_type(self):
        try:
            from transformers import VideoMAEImageProcessor

            HAS_VIDEOMAE_PROCESSOR = True
        except ImportError:
            HAS_VIDEOMAE_PROCESSOR = False

        try:
            from transformers import AutoImageProcessor

            HAS_AUTO_PROCESSOR = True
        except ImportError:
            HAS_AUTO_PROCESSOR = False

        try:
            from transformers import AutoVideoProcessor

            HAS_AUTO_VIDEO_PROCESSOR = True
        except ImportError:
            HAS_AUTO_VIDEO_PROCESSOR = False

        if self.is_vjepa2 and HAS_AUTO_VIDEO_PROCESSOR:
            return "auto_video"
        elif self.is_vjepa and HAS_AUTO_PROCESSOR:
            return "auto_image"
        elif self.is_videomae and HAS_VIDEOMAE_PROCESSOR:
            return "videomae"
        return "manual"


@dataclass
class DataConfig:
    """Data and preprocessing configuration.

    Uses flat directory structure with CSV metadata.
    All clips are stored in 'clips/' folder with labels managed via CSV file.

    CSV file must have columns:
        - path: relative path to clip (e.g., 'clips/video_win001_l1.mp4')
        - label: numeric (0/1) for binary, or string for multi-class
        - split: 'train', 'val', or 'test'
    """

    # Required: root directory containing the clips folder
    data_root: str = "."
    # Required: path to CSV with columns: path, label, split
    csv_path: str = None

    # Frame configuration
    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

    # Data augmentation
    oversample: bool = True
    oversample_ratio: float = (
        1.0  # Binary: neg:pos ratio. Multi-class: ignored (all balanced)
    )
    use_augmentations: bool = True

    # Group boosting: oversample a subgroup (both pos and neg) to improve
    # within-group discrimination. Both sides of the group are boosted equally
    # so the model learns to discriminate within the group, not just correlate
    # group membership with the label.
    # group_boost_column: boolean column in the CSV (e.g. 'has_person', 'has_animal',
    #                     or a precomputed 'has_pedestrian_or_animal' column).
    #                     None = disabled (default behavior, no change).
    # group_boost_factor: how many times more likely a boosted sample is drawn.
    #                     1.0 = no effect. 3.0 = 3x more draws from that group.
    group_boost_column: Optional[str] = None
    group_boost_factor: float = 1.0

    # Class configuration (for multi-class)
    class_names: list = (
        None  # Explicit class ordering. If None, inferred from CSV labels (sorted).
    )

    # Deprecated fields (kept for backwards compatibility with old configs)
    val_data_root: str = None  # Deprecated: validation data comes from CSV split
    split: list = field(
        default_factory=lambda: ["train"]
    )  # Deprecated: splits are managed via CSV
    max_val_samples: int = None  # Deprecated: no longer used


@dataclass
class TrainingConfig:
    """Training parameters"""

    batch_size: int = 8
    val_batch_size: Optional[int] = None

    # Loss function
    loss_type: str = "bce"  # bce, focal, negative_margin
    focal_gamma: float = 2.0  # Focal loss gamma parameter
    focal_alpha: float = 0.25  # Focal loss alpha parameter
    negative_margin: float = 0.3  # Target margin for negatives (push toward 0)

    # Trainer kwargs
    max_steps: int = 100_000
    val_check_interval: int = 1_000
    log_every_n_steps: int = 10
    trainer_kwargs: dict = field(default_factory=lambda: {})

    # Optimizer and scheduler
    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

    # Label smoothing (0.0 = disabled)
    label_smoothing: float = 0.0

    # Early stopping
    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


@dataclass
class SystemConfig:
    """System and hardware configuration"""

    seed: int = 42
    precision: str = "16-mixed"  # 32, 16-mixed, bf16-mixed
    devices: int = 1
    accelerator: str = "auto"  # auto, gpu, cpu
    use_ffmpeg: bool = True
    use_decord: bool = True  # Use decord for video loading (~7x faster than ffmpeg)
    decord_num_threads: int = (
        4  # Decord decode threads per worker (4 optimal on 96-core machines)
    )


@dataclass
class LoggingConfig:
    """Logging configuration"""

    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  # unused
    wandb_tags: list = field(default_factory=list)


@dataclass
class InferenceProbeConfig:
    """Optional: run inference on a fixed video after each checkpoint."""

    video_path: Optional[str] = None  # absolute path to probe video
    stride: int = 1  # frame stride (1 = every frame)


@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)
    inference_probe: InferenceProbeConfig = field(default_factory=InferenceProbeConfig)

    @pl.utilities.rank_zero_only
    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"""
        if path.startswith("s3://"):
            with io.BytesIO() as f:
                import boto3

                s3 = boto3.client("s3")
                bucket, key = path.replace("s3://", "").split("/", 1)
                s3.download_fileobj(bucket, key, f)
                f.seek(0)
                data = yaml.safe_load(f)
        else:
            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):
        return cls(
            model=ModelConfig(**data.get("model", {})),
            data=DataConfig(**data.get("data", {})),
            training=TrainingConfig(**data.get("training", {})),
            system=SystemConfig(**data.get("system", {})),
            logging=LoggingConfig(**data.get("logging", {})),
            inference_probe=InferenceProbeConfig(**data.get("inference_probe", {})),
        )

    @classmethod
    def load_matrix(cls, path: str) -> list["Config"]:
        """Load configuration matrix from YAML file"""
        with open(path, "r") as f:
            matrix_data = yaml.safe_load(f)

        base_config = cls.load(matrix_data["base_path"])

        combinations = []
        for key, values in matrix_data["matrix"].items():
            combinations.append([(key, v) for v in values])
        iterator = product(*combinations)
        configs = []
        for combo in iterator:
            config = deepcopy(base_config)
            updates = {k: v for k, v in combo}
            config.update_from_dict(updates)
            configs.append(config)

        return configs

    def update_from_dict(self, updates: Dict[str, Any]):
        """Update configuration from dictionary (for CLI overrides)"""
        for key, value in updates.items():
            if "." in key:
                sections = key.split(".")
                param = sections.pop()

                # unfold nested sections
                cur = self
                for section in sections:
                    if hasattr(cur, section):
                        cur = getattr(cur, section)
                    else:
                        raise KeyError(f"Unknown configuration section: {section}")
                if hasattr(cur, param):
                    setattr(cur, param, value)
                else:
                    raise KeyError(f"Unknown configuration parameter: {param}")
            else:
                if hasattr(self, key):
                    setattr(self, key, value)
                else:
                    raise KeyError(f"Unknown configuration key: {key}")

    def to_dict(self):
        """Convert configuration to dictionary"""
        return asdict(self)


def get_config() -> Config:
    parser = argparse.ArgumentParser(
        description="Lightning Video Classification Training",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--config", type=str, default=None, help="Path to YAML configuration file"
    )
    parser.add_argument(
        "--matrix", type=str, default=None, help="Path to YAML matrix file"
    )
    args = parser.parse_args()

    if args.config and os.path.exists(args.config):
        path = args.config
    elif args.config:
        raise FileNotFoundError(f"Configuration file not found: {args.config}")
    elif args.matrix and os.path.exists(args.matrix):
        path = args.matrix
        return Config.load_matrix(path)
    elif args.matrix:
        raise FileNotFoundError(f"Matrix file not found: {args.matrix}")
    else:
        path = os.path.join(os.path.dirname(__file__), "configs/default.yaml")
        print("📄 Using default configuration")
    config = Config.load(path)

    return config


@pl.utilities.rank_zero_only
def wandb_log_config(logger, config):
    """Store experiment config in wandb."""
    logger.experiment.config.update(config.to_dict())


@pl.utilities.rank_zero_only
def log_config_to_s3(s3_path, config):
    """Log configuration file to S3 bucket."""
    import boto3

    s3 = boto3.client("s3")
    bucket, key = s3_path.replace("s3://", "").split("/", 1)

    config_str = yaml.dump(asdict(config), default_flow_style=False, indent=2)
    config_bytes = io.BytesIO(config_str.encode("utf-8"))
    config_bytes.seek(0)
    s3.upload_fileobj(config_bytes, bucket, key)


def get_or_set_run_name():
    # DDP will rerun the generate call so use an env var to
    # fix the run name for all processes.
    run_name = os.getenv("RUN_NAME", None)
    if run_name is None:
        run_name = generate_unique_run_name()
        os.environ["RUN_NAME"] = run_name
    return run_name
