"""

automl_lib.cli

~~~~~~~~~~~~~~

Command-line interface for no-code model training and inference.
"""

from __future__ import annotations

import argparse

import sys



def main(argv: list[str] | None = None):

    """Entry-point for the ``automl`` CLI."""

    parser = argparse.ArgumentParser(

        prog="automl",

        description="Universal AutoML Library – train, evaluate, and deploy ML/DL models from the command line.",
        epilog="""

Examples:

  automl train --data dataset.csv --target price --task regression

  automl train --data images/ --data-type image --tune --epochs 100

  automl predict --model-path automl_output/checkpoints/best.pth --data test.csv
  automl info


For detailed help on a command: automl <command> --help

        """,

        formatter_class=argparse.RawDescriptionHelpFormatter

    )

    subparsers = parser.add_subparsers(dest="command", help="Available commands")


    # ── train ──────────────────────────────────────────────────────────

    train_p = subparsers.add_parser(

        "train", 

        help="Train a model on a dataset",
        description="""

Train a machine learning or deep learning model on your dataset.

The library automatically detects task type, data modality, and selects

appropriate models. Supports tabular, image, text, audio, and timeseries data.


WORKFLOW:

  1. Data is loaded and preprocessed automatically

  2. Task type (classification/regression) is detected from data

  3. Model architecture is selected based on data type

  4. Training runs with early stopping and checkpointing

  5. Best model is saved to output directory


EXAMPLES:

  # Basic tabular classification

  automl train --data data.csv --target label
  

  # Regression with hyperparameter tuning

  automl train --data housing.csv --target price --task regression --tune
  

  # Image classification with custom settings

  automl train --data images/ --data-type image --epochs 100 --batch-size 64
  

  # Export model after training

  automl train --data data.csv --target class --export onnx

        """,

        formatter_class=argparse.RawDescriptionHelpFormatter

    )
    

    # Data arguments

    data_group = train_p.add_argument_group('Data Configuration')

    data_group.add_argument("--data", required=True, 

                           help="Path to dataset (CSV file for tabular, directory for images/audio)")

    data_group.add_argument("--task", default=None, choices=["classification", "regression", "forecasting"],

                           help="ML task type (default: auto-detected from data)")

    data_group.add_argument("--data-type", default=None,

                           choices=["tabular", "image", "text", "audio", "timeseries"],

                           help="Data modality (default: auto-detected from file/folder structure)")

    data_group.add_argument("--target", default=None, 

                           help="Target column name for tabular data (required for CSV files)")
    

    # Model arguments

    model_group = train_p.add_argument_group('Model Configuration')

    model_group.add_argument("--model", default=None, 

                            help="Specific model name (e.g., 'resnet18', 'xgboost') or family (e.g., 'cnn', 'transformer')")

    model_group.add_argument("--no-pretrained", action="store_true", 

                            help="Train from scratch without pretrained weights (default: use pretrained)")
    

    # Training arguments

    training_group = train_p.add_argument_group('Training Parameters')

    training_group.add_argument("--epochs", type=int, default=50, 

                               help="Number of training epochs (default: 50)")

    training_group.add_argument("--batch-size", type=int, default=32, 

                               help="Batch size for training (default: 32)")

    training_group.add_argument("--lr", type=float, default=1e-3, 

                               help="Learning rate (default: 0.001)")

    training_group.add_argument("--optimizer", default="adam", choices=["adam", "adamw", "sgd"],

                               help="Optimizer algorithm (default: adam)")
    

    # Hyperparameter tuning

    tune_group = train_p.add_argument_group('Hyperparameter Tuning')

    tune_group.add_argument("--tune", action="store_true", 

                           help="Enable automatic hyperparameter optimization")

    tune_group.add_argument("--tune-trials", type=int, default=30, 

                           help="Number of HPO trials to run (default: 30)")
    

    # Output and system

    system_group = train_p.add_argument_group('System & Output')

    system_group.add_argument("--output", default="automl_output", 

                             help="Output directory for models and logs (default: automl_output)")

    system_group.add_argument("--device", default=None, 

                             help="Device to use: 'cpu', 'cuda', 'mps' (default: auto-detected)")

    system_group.add_argument("--seed", type=int, default=42, 

                             help="Random seed for reproducibility (default: 42)")
    

    # Export

    export_group = train_p.add_argument_group('Model Export')

    export_group.add_argument("--export", default=None, choices=["torchscript", "onnx", "both"],

                             help="Export trained model to specified format(s)")


    # ── predict ───────────────────────────────────────────────────────

    pred_p = subparsers.add_parser(

        "predict", 

        help="Run inference with a trained model",
        description="""

Run predictions on new data using a trained model checkpoint.


WORKFLOW:

  1. Load the trained model from checkpoint

  2. Preprocess input data using saved preprocessing pipeline

  3. Generate predictions

  4. Save results to output file


EXAMPLES:

  # Predict on new data

  automl predict --model-path automl_output/checkpoints/best.pth --data test.csv
  

  # Custom output location

  automl predict --model-path model.pth --data new_data.csv --output results.csv


NOTE: For full functionality, use the Python API:

  from automl_lib import AutoML

  predictor = AutoML.load('model.pth')

  predictions = predictor.predict('test.csv')

        """,

        formatter_class=argparse.RawDescriptionHelpFormatter

    )

    pred_p.add_argument("--model-path", required=True, 

                       help="Path to saved model checkpoint (.pth file)")

    pred_p.add_argument("--data", required=True, 

                       help="Path to input data (same format as training data)")

    pred_p.add_argument("--output", default="predictions.csv", 

                       help="Output file for predictions (default: predictions.csv)")


    # ── export ────────────────────────────────────────────────────────

    export_p = subparsers.add_parser(

        "export", 

        help="Export a trained model to production formats",
        description="""

Export trained models to deployment-ready formats for production use.


SUPPORTED FORMATS:

  - torchscript: Optimized PyTorch format for C++ deployment

  - onnx: Open Neural Network Exchange format (cross-framework)

  - both: Export to both formats


WORKFLOW:

  1. Load trained model checkpoint

  2. Convert to specified format(s)

  3. Save exported model(s) to output directory

  4. Validate exported model can run inference


EXAMPLES:

  # Export to ONNX for cross-platform deployment

  automl export --model-path automl_output/checkpoints/best.pth --format onnx
  

  # Export to both formats

  automl export --model-path model.pth --format both --output exports/


NOTE: For full functionality, use the Python API:

  from automl_lib import AutoML

  automl = AutoML.load('model.pth')

  paths = automl.export(format='onnx')

        """,

        formatter_class=argparse.RawDescriptionHelpFormatter

    )

    export_p.add_argument("--model-path", required=True, 

                         help="Path to trained model checkpoint (.pth file)")

    export_p.add_argument("--format", default="both", choices=["torchscript", "onnx", "both"],

                         help="Export format (default: both)")

    export_p.add_argument("--output", default="automl_output/exported",

                         help="Output directory for exported models (default: automl_output/exported)")


    # ── deploy ────────────────────────────────────────────────────────

    deploy_p = subparsers.add_parser(

        "deploy", 

        help="Generate a FastAPI deployment scaffold",
        description="""

Generate a production-ready FastAPI application for model serving.


WORKFLOW:

  1. Create FastAPI application structure

  2. Generate endpoint for model inference

  3. Add health check and documentation endpoints

  4. Create requirements.txt and README


OUTPUT STRUCTURE:

  api/

  ├── app.py           # FastAPI application

  ├── requirements.txt # Python dependencies

  └── README.md        # Deployment instructions


EXAMPLES:

  # Generate API scaffold

  automl deploy --model-path automl_output/exported/model.onnx
  

  # Custom output directory

  automl deploy --model-path model.onnx --output my_api/


RUNNING THE API:

  cd automl_output/api

  pip install -r requirements.txt

  uvicorn app:app --reload
  

  # API will be available at http://localhost:8000

  # Docs at http://localhost:8000/docs

        """,

        formatter_class=argparse.RawDescriptionHelpFormatter

    )

    deploy_p.add_argument("--model-path", required=True, 

                         help="Path to exported model file (.onnx or .pt)")

    deploy_p.add_argument("--output", default="automl_output/api",

                         help="Output directory for API scaffold (default: automl_output/api)")


    # ── info ──────────────────────────────────────────────────────────

    subparsers.add_parser(

        "info", 

        help="Show hardware and library information",
        description="""

Display system information including available hardware accelerators,

library version, and detected compute capabilities.


DISPLAYS:

  - Best available device (CPU, CUDA, MPS)

  - GPU information (if available)

  - Memory information

  - Library version and configuration


EXAMPLE:
  automl info

        """,

        formatter_class=argparse.RawDescriptionHelpFormatter

    )

    # ── research ──────────────────────────────────────────────────────

    research_p = subparsers.add_parser(

        "research", 

        help="Run the WebResearch Engine",
        description="""

Run the WebResearch Engine for querying information.

EXAMPLES:

  # Basic research query
  automl research "What is Retrieval-Augmented Generation?"

  # Research with specific mode and max sources
  automl research "Attention is All You Need" --mode academic --max-sources 3

        """,

        formatter_class=argparse.RawDescriptionHelpFormatter

    )
    
    research_p.add_argument("query", help="Natural language prompt to research")
    research_p.add_argument("--mode", default="general", choices=["general", "academic", "documentation", "github"],
                            help="Research mode (default: general)")
    research_p.add_argument("--max-sources", type=int, default=5, 
                            help="Maximum number of search results to scrape (default: 5)")
    research_p.add_argument("--style", default="concise", 
                            help="Summarization style (default: concise)")
    research_p.add_argument("--llm-provider", default="ollama", choices=["ollama", "huggingface"],
                            help="LLM provider: ollama or huggingface (default: ollama)")
    research_p.add_argument("--llm", default="llama3", 
                            help="LLM model name (default: llama3 for Ollama, or HF model like 'meta-llama/Llama-2-7b-chat-hf')")
    research_p.add_argument("--hf-token", default=None,
                            help="HuggingFace API token (required for gated models)")
    research_p.add_argument("--embed-model", default="all-MiniLM-L6-v2", 
                            help="Sentence Transformers embedding model (default: all-MiniLM-L6-v2)")


    args = parser.parse_args(argv)


    if args.command is None:

        parser.print_help()
        return


    if args.command == "info":

        _cmd_info()

    elif args.command == "train":

        _cmd_train(args)

    elif args.command == "predict":

        _cmd_predict(args)

    elif args.command == "export":

        _cmd_export(args)

    elif args.command == "research":

        _cmd_research(args)

    elif args.command == "deploy":

        _cmd_deploy(args)



def _cmd_info():

    from .hardware import device_summary, detect_device
    from .utils import get_logger

    logger = get_logger()

    dev = detect_device()

    summary = device_summary()

    logger.info(f"Best device: {dev}")

    logger.info(f"Hardware summary: {summary}")

    print(f"automl_lib — Universal AutoML Library")

    print(f"Best device  : {dev}")

    print(f"Hardware info: {summary}")



def _cmd_train(args):

    from .config import trainiqConfig
    from . import trainiq


    config = trainiqConfig(

        data_path=args.data,

        task=args.task,

        data_type=args.data_type,

        target_column=args.target,

        model_name=args.model,

        epochs=args.epochs,

        batch_size=args.batch_size,

        learning_rate=args.lr,

        optimizer=args.optimizer,

        tune=args.tune,

        tune_trials=args.tune_trials,

        output_dir=args.output,

        device=args.device,

        seed=args.seed,

        pretrained=not args.no_pretrained,

    )

    if args.export:

        config.export_format = args.export

    automl = trainiq(config)

    results = automl.train()


    if args.export:

        automl.export()


    print("\nTraining complete!")

    if "best_val_acc" in results and results["best_val_acc"] is not None:

        print(f"   Best validation accuracy: {results['best_val_acc']:.4f}")

    print(f"   Best validation loss:     {results['best_val_loss']:.4f}")

    print(f"   Model saved:              {results['best_model_path']}")



def _cmd_predict(args):

    print(f"Prediction from CLI requires model architecture info. "

          f"Use the Python API for full predict support:")

    print(f"  from automl_lib import AutoML")

    print(f"  predictor = AutoML.load('{args.model_path}')")

    print(f"  results = predictor.predict('{args.data}')")



def _cmd_export(args):

    print(f"Export from CLI requires model architecture info. "

          f"Use the Python API for full export support:")

    print(f"  from automl_lib import AutoML")

    print(f"  automl = AutoML.load('{args.model_path}')")

    print(f"  automl.export(format='{args.format}')")



def _cmd_deploy(args):

    from .exporter import scaffold_api

    scaffold_api(args.model_path, args.output)

    print(f"\nFastAPI app created at: {args.output}")

    print(f"   Run: cd {args.output} && uvicorn app:app --reload")



def _cmd_research(args):

    from .research import research

    from .utils import get_logger
    get_logger().setLevel("INFO")

    print(f"\n[ trainiq WebResearch Engine ]")
    print(f"Query: {args.query}")
    print(f"Mode: {args.mode} | Max Sources: {args.max_sources} | Provider: {args.llm_provider} | LLM: {args.llm}\n")
    
    try:
        results = research(
            query=args.query,
            embedding_model_name=args.embed_model,
            llm_model=args.llm,
            llm_provider=args.llm_provider,
            hf_token=args.hf_token,
            max_sources=args.max_sources,
            style=args.style,
            research_mode=args.mode
        )
        
        print(f"\n{'='*60}\nANSWER:\n{'='*60}")
        print(results.get("answer", ""))
        print(f"\n{'='*60}\nREFERENCES (Used {results.get('context_used', 0)} chunks):\n{'='*60}")
        for ref in results.get("references", []):
            print(f"- {ref}")
            
    except ImportError as e:
        print(f"\n[!] Missing Dependencies: {e}")
        print("Please install them using: pip install 'trainiq[research]'")
    except Exception as e:
        print(f"\n[!] Error during research: {e}")


if __name__ == "__main__":

    main()

