"""
Lightning Training Script for Video Collision Anticipation
=========================================================

Main training script using PyTorch Lightning framework with YAML configuration.
Provides cleaner code structure, built-in callbacks, and better reproducibility.

Uses shared preprocessing from core.preprocessing - no HuggingFace dependency.
"""

import os

import torch
import pytorch_lightning as pl
from pytorch_lightning.callbacks import (
    EarlyStopping,
    LearningRateMonitor,
    ModelCheckpoint,
)
from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger

from badas.training.config import (
    Config,
    get_config,
    get_or_set_run_name,
    log_config_to_s3,
    wandb_log_config,
)
from badas.training.inference_callback import InferenceProbeCallback
from badas.training.lightning_data_module import VideoDataModule
from badas.training.lightning_module import VideoLightningModule


def setup_callbacks(root_dir: str, config: Config) -> list:
    """Setup Lightning callbacks."""
    callbacks = []

    # Model checkpointing - save all checkpoints
    checkpoint_callback = ModelCheckpoint(
        monitor="val_ap",
        mode="max",
        save_top_k=-1,
        save_last=True,
        filename=f"{config.logging.run_name}-{{step}}-{{val_ap:.4f}}",
        auto_insert_metric_name=False,
        dirpath=root_dir,
    )
    callbacks.append(checkpoint_callback)

    # Early stopping
    if not config.training.disable_early_stopping:
        early_stop_callback = EarlyStopping(
            monitor=config.training.early_stopping_metric,
            patience=config.training.early_stopping_patience,
            min_delta=config.training.early_stopping_min_delta,
            mode=config.training.early_stopping_mode,
            verbose=True,
        )
        callbacks.append(early_stop_callback)

    # Learning rate monitor
    lr_monitor = LearningRateMonitor(logging_interval="step")
    callbacks.append(lr_monitor)

    # Optional: inference probe on a fixed video after each checkpoint
    probe = config.inference_probe
    if probe.video_path:
        probe_output_dir = os.path.join(root_dir, "inference_probes")
        inference_probe_callback = InferenceProbeCallback(
            video_path=probe.video_path,
            output_dir=probe_output_dir,
            stride=probe.stride,
        )
        callbacks.append(inference_probe_callback)
        print(f"Inference probe enabled: {probe.video_path} → {probe_output_dir}")

    return callbacks


def main():
    config = get_config()

    run_name = get_or_set_run_name()
    config.logging.run_name = run_name
    if not config.logging.output_dir:
        base_dir = os.path.join(os.path.dirname(__file__), "logs")
    else:
        base_dir = config.logging.output_dir
    unique_output_dir = os.path.join(base_dir, run_name)
    config.logging.output_dir = unique_output_dir
    os.makedirs(config.logging.output_dir, exist_ok=True)

    print(f"Starting Lightning training run: {run_name}")
    print(f"Output directory: {config.logging.output_dir}")
    if config.logging.s3_bucket is not None:
        default_root_dir = f"s3://{config.logging.s3_bucket.strip('/')}"
        if config.logging.s3_prefix is not None:
            default_root_dir += f"/{config.logging.s3_prefix.strip('/')}"
        default_root_dir += f"/{run_name}"
        log_config_to_s3(default_root_dir + "/config.yml", config)
        print(f"Using S3 output directory: {default_root_dir}")
    else:
        default_root_dir = config.logging.output_dir

    # Save config
    config_file = os.path.join(config.logging.output_dir, "config.yml")
    config.save(config_file)
    print(f"Config saved to: {config_file}")

    # Override parameters from model name
    if config.model.frame_count:
        config.data.frame_count = config.model.frame_count
        print(f"Frame count set to {config.data.frame_count} from model name")

    if config.model.crop_size and config.model.is_vjepa2:
        config.data.vjepa2_crop_size = config.model.crop_size
        print(f"Crop size set to {config.data.vjepa2_crop_size} from model name")

    pl.seed_everything(config.system.seed)

    data_module = VideoDataModule(config)
    model = VideoLightningModule(config)

    logger = None
    if config.logging.use_wandb:
        try:
            import wandb

            logger = WandbLogger(
                project=config.logging.wandb_project,
                entity=config.logging.wandb_entity,
                name=config.logging.run_name,
                save_dir=config.logging.output_dir,
                log_model=False,
                tags=config.logging.wandb_tags,
            )

            wandb_log_config(logger, config)

            print(
                f"Weights & Biases logging enabled: {config.logging.wandb_project}/{config.logging.run_name}"
            )

        except ImportError:
            print("Weights & Biases not installed. Install with: pip install wandb")
            print("Falling back to TensorBoard logging")
            logger = None
        except Exception as e:
            print(f"Failed to setup Weights & Biases: {e}")
            print("Falling back to TensorBoard logging")
            logger = None

    # Use TensorBoard as fallback or if W&B not requested
    if logger is None:
        logger = TensorBoardLogger(
            save_dir=config.logging.output_dir,
            name="lightning_logs",
            version=None,
        )
        print("TensorBoard logging enabled")

    callbacks = setup_callbacks(default_root_dir, config)

    trainer_kwargs = {
        "accelerator": config.system.accelerator,
        "devices": config.system.devices,
        "precision": config.system.precision,
        "enable_checkpointing": True,
        "enable_progress_bar": True,
        "enable_model_summary": True,
        "default_root_dir": default_root_dir,
        "log_every_n_steps": config.training.log_every_n_steps,
        "val_check_interval": config.training.val_check_interval,
        "check_val_every_n_epoch": None,
        "max_steps": config.training.max_steps,
        **config.training.trainer_kwargs,
    }

    trainer = pl.Trainer(
        callbacks=callbacks,
        logger=logger,
        **trainer_kwargs,
    )

    # Optional warm-start: load weights from a previous checkpoint without restoring
    # optimizer state or epoch counters.  Set CKPT_PATH env var to enable.
    ckpt_path = os.environ.get('CKPT_PATH', None)
    if ckpt_path and os.path.exists(ckpt_path):
        print(f"[Warm-Start] Loading model weights from: {ckpt_path}")
        ckpt = torch.load(ckpt_path, map_location='cpu')
        state_dict = ckpt.get('state_dict', ckpt)
        result = model.load_state_dict(state_dict, strict=False)
        n_ok = len(state_dict) - len(result.missing_keys)
        print(f"[Warm-Start] Loaded {n_ok}/{len(state_dict)} param tensors "
              f"({len(result.missing_keys)} missing, {len(result.unexpected_keys)} unexpected)")
    elif ckpt_path:
        raise FileNotFoundError(f"[Warm-Start] CKPT_PATH not found: {ckpt_path}")

    print("Starting training...")
    resume_ckpt = os.environ.get('RESUME_CKPT_PATH', None)
    if resume_ckpt:
        print(f"[Resume] Resuming from checkpoint: {resume_ckpt}")
    trainer.fit(model, data_module, ckpt_path=resume_ckpt)

    best_model_path = trainer.checkpoint_callback.best_model_path
    print(f"Training complete!")
    print(f"Best model: {best_model_path}")
    print(f"Best val accuracy: {trainer.checkpoint_callback.best_model_score:.4f}")


if __name__ == "__main__":
    main()
