import os
import argparse
import yaml
import subprocess

def load_config(config_path):
    """Load configuration from YAML file."""
    with open(config_path, "r") as f:
        return yaml.safe_load(f)

def create_dataset(config):
    """Create dataset using the configuration."""
    dataset_config = config["dataset"]
    
    cmd = [
        "python", "create_dataset.py",
        "--size", str(dataset_config["size"]),
        "--ratio", str(dataset_config["predatory_ratio"]),
        "--output", dataset_config["output_path"]
    ]
    
    print(f"Creating dataset with {dataset_config['size']} examples...")
    result = subprocess.run(cmd)
    if result.returncode != 0:
        print("Dataset creation failed!")
        return False
    
    print(f"Dataset created successfully at {dataset_config['output_path']}")
    return True

def finetune_model(config):
    """Fine-tune the model using the configuration."""
    dataset_config = config["dataset"]
    model_config = config["model"]
    training_config = config["training"]
    peft_config = config["peft"]
    precision_config = config["precision"]
    logging_config = config["logging"]
    
    cmd = [
        "python", "finetune_heaven.py",
        "--model_name_or_path", model_config["name_or_path"],
        "--dataset_path", dataset_config["output_path"],
        "--output_dir", model_config["output_dir"],
        "--num_epochs", str(training_config["num_epochs"]),
        "--batch_size", str(training_config["batch_size"]),
        "--gradient_accumulation_steps", str(training_config["gradient_accumulation_steps"]),
        "--learning_rate", str(training_config["learning_rate"]),
        "--weight_decay", str(training_config["weight_decay"]),
        "--max_grad_norm", str(training_config["max_grad_norm"]),
        "--warmup_ratio", str(training_config["warmup_ratio"]),
        "--eval_ratio", str(training_config["eval_ratio"]),
        "--max_seq_length", str(training_config["max_seq_length"]),
        "--logging_steps", str(logging_config["logging_steps"]),
        "--eval_steps", str(logging_config["eval_steps"]),
        "--save_steps", str(logging_config["save_steps"]),
        "--run_name", logging_config["run_name"],
        "--compute_dtype", precision_config["compute_dtype"]
    ]
    
    # Add boolean flags
    if peft_config["use_lora"]:
        cmd.append("--use_lora")
    if peft_config["use_qlora"]:
        cmd.append("--use_qlora")
    if precision_config["fp16"]:
        cmd.append("--fp16")
    if precision_config["bf16"]:
        cmd.append("--bf16")
    if logging_config["use_wandb"]:
        cmd.append("--use_wandb")
    
    # Add LoRA parameters
    cmd.extend(["--lora_r", str(peft_config["lora_r"])])
    cmd.extend(["--lora_alpha", str(peft_config["lora_alpha"])])
    cmd.extend(["--lora_dropout", str(peft_config["lora_dropout"])])
    
    print("Starting fine-tuning process...")
    result = subprocess.run(cmd)
    if result.returncode != 0:
        print("Fine-tuning failed!")
        return False
    
    print(f"Fine-tuning completed successfully! Model saved to {model_config['output_dir']}")
    return True

def main():
    parser = argparse.ArgumentParser(description="Run the Heaven fine-tuning pipeline")
    parser.add_argument("--config", type=str, default="config.yaml", help="Path to the configuration file")
    parser.add_argument("--skip-dataset", action="store_true", help="Skip dataset creation step")
    args = parser.parse_args()
    
    print(f"Loading configuration from {args.config}...")
    config = load_config(args.config)
    
    # Create necessary directories
    os.makedirs(os.path.dirname(config["dataset"]["output_path"]), exist_ok=True)
    os.makedirs(config["model"]["output_dir"], exist_ok=True)
    
    # Create dataset if not skipped
    if not args.skip_dataset:
        success = create_dataset(config)
        if not success:
            return
    else:
        print("Skipping dataset creation...")
    
    # Fine-tune the model
    finetune_model(config)

if __name__ == "__main__":
    main()