import yaml
import argparse
import sys
from aiforge import Pipeline, load_plugin

def main():
    parser = argparse.ArgumentParser(description="AIForge CLI Execution Engine")
    parser.add_argument("config", help="Path to the YAML pipeline configuration file")
    
    args = parser.parse_args()
    
    try:
        with open(args.config, 'r') as file:
            config = yaml.safe_load(file)
    except FileNotFoundError:
        print(f"Error: Could not find configuration file at '{args.config}'")
        sys.exit(1)
        
    if "pipeline" not in config:
        print("Error: No 'pipeline' key found in YAML config.")
        sys.exit(1)
            
    if "plugins" in config:
        for plugin in config["plugins"]:
            load_plugin(plugin)

    pipe = Pipeline()
    for step in config["pipeline"]:
        for op_name, params in step.items():
            params = params or {}
            pipe.add(op_name, **params)
            
    input_data = config.get("input_data", None)
    
    print(f"Running pipeline with {len(pipe.steps)} steps...")
    
    try:
        result = pipe.run(input_data)
        print("\n=== Execution Complete ===")
        print(result)
    except Exception as e:
        print(f"Error during execution: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
