#!/usr/bin/env python3
"""Training script for alignment faking research."""

import argparse
import logging
import random
import sys
from pathlib import Path
from typing import Any

import numpy as np
import torch
import yaml
from .data.system_prompts import BASELINE_SYSTEM_PROMPT, HELPFUL_ONLY_SYSTEM_PROMPT
from .src.data import create_datasets
from .src.grpo_training import GRPOTrainer


def setup_logging(log_level: str = "INFO") -> None:
    """Setup clean logging configuration."""
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format="%(asctime)s | %(levelname)-7s | %(message)s",
        datefmt="%H:%M:%S",
    )


def set_seed(seed: int) -> None:
    """Set random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def load_config(config_path: str) -> dict[str, Any]:
    """Load configuration from YAML file."""
    try:
        with open(config_path) as f:
            config = yaml.safe_load(f)
        return config  # type: ignore[no-any-return]
    except Exception as e:
        logging.error(f"Failed to load config from {config_path}: {e}")
        sys.exit(1)


def validate_config(config: dict[str, Any]) -> None:
    """Validate configuration parameters."""
    required_keys = ["model", "data", "training", "logging"]

    for key in required_keys:
        if key not in config:
            raise ValueError(f"Missing required config section: {key}")

    # Validate paths
    data_path = Path(config["data"]["harmful_queries_path"])
    if not data_path.exists():
        raise FileNotFoundError(f"Data file not found: {data_path}")

    # Validate numeric ranges
    if not 0 < config["data"]["train_eval_split"] < 1:
        raise ValueError("train_eval_split must be between 0 and 1")

    if config["training"]["num_epochs"] <= 0:
        raise ValueError("num_epochs must be positive")


def main() -> None:
    parser = argparse.ArgumentParser(description="Train alignment faking model")
    parser.add_argument("--config", default="config.yaml", help="Config file path")
    parser.add_argument("--output", default="results", help="Output directory")
    parser.add_argument("--no-cuda", action="store_true", help="Disable CUDA")
    args = parser.parse_args()

    # Load and validate config
    config = load_config(args.config)

    # Override device if CUDA disabled
    if args.no_cuda:
        config["model"]["device"] = "cpu"

    # Setup logging
    setup_logging(config["logging"].get("log_level", "INFO"))
    set_seed(config["research"].get("seed", 42))

    try:
        validate_config(config)

        # Create output directory
        Path(args.output).mkdir(parents=True, exist_ok=True)

        print("\n" + "=" * 50)
        print("ALIGNMENT FAKING TRAINING")
        print("=" * 50)

        logging.info("Creating datasets...")
        train_dataset, eval_dataset = create_datasets(
            config, HELPFUL_ONLY_SYSTEM_PROMPT, BASELINE_SYSTEM_PROMPT
        )

        logging.info("Initializing GRPO trainer with pipeline parallelism...")
        trainer = GRPOTrainer(config)

        print("\n" + "-" * 50)
        logging.info("Starting training...")
        trainer.train(
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            num_epochs=config["training"]["num_epochs"],
        )

        # Save model
        save_path = Path(args.output) / "final_model"
        trainer.save_model(str(save_path))
        print("\n" + "=" * 50)
        logging.info(f"Training completed! Model saved to {save_path}")
        print("=" * 50 + "\n")

    except Exception as e:
        logging.error(f"Training failed: {e}")
        sys.exit(1)
    finally:
        # Cleanup
        if "trainer" in locals():
            trainer.cleanup()


if __name__ == "__main__":
    main()
