"""
    Main entry point for the auto_vicuna package.
"""
import os
import sys
from pathlib import Path

import click
import torch
from dotenv import load_dotenv

from auto_vicuna.model import load_model
from auto_vicuna.loop import main_loop
from auto_vicuna.plugins import load_plugins

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
load_dotenv()


@click.command()
@click.option(
    "--vicuna_weights",
    type=click.Path(exists=True),
    default=lambda: os.environ.get("VICUNA_WEIGHTS", ""),
)
@click.option("--num_gpus", type=int, default=1)
@click.option("--device", type=str, default=DEVICE)
@click.option("--debug", is_flag=True)
@click.option("--load_8bit", is_flag=True)
@click.option(
    "--plugins_path", type=click.Path(exists=True),
    default=lambda: os.environ.get("PLUGINS_PATH", Path(os.getcwd()) / "plugins")
)
def main(
    vicuna_weights: Path,
    num_gpus: int = 1,
    device: str = DEVICE,
    debug: bool = False,
    load_8bit: bool = False,
    plugins_path: Path = Path(os.getcwd()) / "plugins",
) -> None:
    """Auto-Vicuna: A Python package for automatically generating Vicuna randomness."""
    click.echo(f"Auto-Vicuna\n===========\nVersion: 0.1.0\nWeights: {vicuna_weights}")
    click.echo(f"Device: {device}\nTorch version: {torch.__version__}")
    if "cpu" in torch.__version__:
        click.echo("\nError: CPU not supported. Install a GPU version of PyTorch.")
        click.echo("See https://pytorch.org/get-started/locally/ for more info.\n")
        sys.exit(1)
    try:
        model, tokenizer = load_model(
            vicuna_weights,
            device=device,
            num_gpus=num_gpus,
            debug=debug,
            load_8bit=load_8bit,
        )
    except ValueError as e:  # noqa: F841
        if "Tokenizer" in str(e):
            click.echo(
                "Error: Tokenizer not found. Likely due to transformers "
                "and your weights generated by different versions..."
            )
            click.echo(
                "Replace LLaMATokenizer with LlamaTokenizer in your configs"
                " or use the same version of transformers as your weights."
            )
            sys.exit(1)
        click.echo(
            "Error: CPU not supported. Select a GPU device instead, e.g."
            " --device=cuda"
        )
        sys.exit(1)
    plugins_found = load_plugins(plugins_path)
    loaded_plugins = [plugin() for plugin in plugins_found]
    if loaded_plugins:
        click.echo(f"\nPlugins found: {len(loaded_plugins)}\n"
                   "--------------------")
    for plugin in loaded_plugins:
        click.echo(f"{plugin._name}: {plugin._version} - {plugin._description}")

    if debug:
        click.echo(f"Model: {model}")

    main_loop(model, tokenizer, "bair_v1", temperature=0.7,
              max_new_tokens=512, plugins=loaded_plugins, debug=debug,
              model_path=vicuna_weights)


if __name__ == "__main__":
    main(sys.argv[1:])
