from typing import Any

from nbconvert import HTMLExporter

from utils.create_info_files import create_hf_card
from utils.notebook_generator import *
from utils.components_creator import *

finetuning_notebook = "Finetuning_NoteBook"

notebook = None

css = """
.container {
    align-items: center;
    justify-content: center;
}
.center_text {
    text-align: center;
}

.a_custom {
    border-radius: var(--button-large-radius);
    padding: var(--button-large-padding);
    font-weight: var(--button-large-text-weight);
    font-size: var(--button-large-text-size);
    border: var(--button-border-width) solid var(--button-primary-border-color);
    background: var(--button-primary-background-fill);
    color: var(--button-primary-text-color);
    justify-content: center;
    align-items: center;
    transition: var(--button-transition);
    box-shadow: var(--button-shadow);
    text-align: center
}
.a_custom a {
    text-decoration: none;
    color: white;
}
"""


def centered_column():
    return gr.Column(elem_classes=["container"])


def should_login_to_hf_model(model_id: str):
    return model_id == gemma.name or model_id == llama.name


def change_model_selection(model_id):
    if model_id == gemma.name:
        gr.Warning("""
        Access Gemma:
        
        To load Gemma from Hugging Face, you’re required to review and agree to Google’s usage license.
        """)
    if model_id == llama.name:
        gr.Warning("""
        Access Llama 2:
        
        To load Llama 2 from Hugging Face, you’re required to review and agree to Meta’s usage license.
        """)

    for m in models:
        if m.name == model_id:
            return gr.Dropdown(choices=m.versions, interactive=True,
                               visible=True, info=f"Select the version of the model {m.name} you wish to use.")
    return None


def check_valid_input(value):
    if isinstance(value, str):
        return value and value.strip()
    if isinstance(value, list):
        return len(value) > 0
    return not None


def get_dataset(dataset_path):
    for d in ft_datasets:
        if d.path == dataset_path:
            return d
    return None


def get_value(components: dict[Component, Any], elem_id: str) -> Any:
    for component, val in components.items():
        if component.elem_id == elem_id:
            return val
    return None


def preview_notebook():
    html_exporter = HTMLExporter()
    global notebook
    (body, resources) = html_exporter.from_notebook_node(notebook)

    html_path = f"{finetuning_notebook}.html"
    with open(html_path, 'w') as f:
        f.write(body)

    return f'<iframe src="file={html_path}" width="100%" height="250px"></iframe>'


def generate_code(components: dict[Component, Any]):
    global notebook
    notebook = nbf.v4.new_notebook()
    create_install_libraries_cells(notebook['cells'])
    flash_attention_value = get_value(components, FLASH_ATTENTION_ID)
    if flash_attention_value:
        create_install_flash_attention(notebook['cells'])

    dataset_value = get_value(components, DATASET_SELECTION_ID)
    seed_value = get_value(components, DATASET_SHUFFLING_SEED)
    if not check_valid_input(dataset_value):
        gr.Warning("No dataset is selected")
    else:
        create_datasets_cells(notebook['cells'], get_dataset(dataset_value), seed_value)

    model_value = get_value(components, MODEL_SELECTION_ID)
    should_login = should_login_to_hf_model(model_value)

    version_value = ""
    if not check_valid_input(model_value):
        gr.Warning("No model is selected!")
    else:
        version_value = get_value(components, MODEL_VERSION_SELECTION_ID)
        if not check_valid_input(version_value):
            gr.Warning("No version of the model is selected")
        else:
            if should_login:
                create_login_hf_cells(notebook['cells'], should_login=True, model_name=model_value)

            load_in_4bit = get_value(components, LOAD_IN_4_BIT_ID)
            bnb_4bit_use_double_quant = get_value(components, BNB_4BIT_USE_DOUBLE_QUANT)
            bnb_4bit_quant_type = get_value(components, BNB_4BIT_QUANT_TYPE)
            bnb_4bit_compute_dtype = get_value(components, BNB_4BIT_COMPUTE_DTYPE)
            pad_side = get_value(components, PAD_SIDE_ID)
            pad_value = get_value(components, PAD_VALUE_ID)
            create_model_cells(notebook['cells'], model_id=model_value, version=version_value,
                               flash_attention=flash_attention_value, pad_value=pad_value,
                               pad_side=pad_side, load_in_4bit=load_in_4bit,
                               bnb_4bit_use_double_quant=bnb_4bit_use_double_quant,
                               bnb_4bit_quant_type=bnb_4bit_quant_type, bnb_4bit_compute_dtype=bnb_4bit_compute_dtype)

    r_value = get_value(components, LORA_R_ID)
    alpha_value = get_value(components, LORA_ALPHA_ID)
    dropout_value = get_value(components, LORA_DROPOUT_ID)
    bias_value = get_value(components, LORA_BIAS_ID)
    create_lora_config_cells(notebook['cells'], r_value, alpha_value, dropout_value, bias_value)

    epochs = get_value(components, NUM_TRAIN_EPOCHS_ID)
    max_steps = get_value(components, MAX_STEPS_ID)
    logging_steps = get_value(components, LOGGING_STEPS_ID)
    per_device_train_batch_size = get_value(components, PER_DEVICE_TRAIN_BATCH_SIZE)
    save_strategy = get_value(components, SAVE_STRATEGY_ID)
    gradient_accumulation_steps = get_value(components, GRADIENT_ACCUMULATION_STEPS_ID)
    gradient_checkpointing = get_value(components, GRADIENT_CHECKPOINTING_ID)
    learning_rate = get_value(components, LEARNING_RATE_ID)
    max_grad_norm = get_value(components, MAX_GRAD_NORM_ID)
    warmup_ratio = get_value(components, WARMUP_RATIO_ID)
    lr_scheduler_type = get_value(components, LR_SCHEDULER_TYPE_ID)
    output_dir = get_value(components, OUTPUT_DIR_ID)

    report_to = get_value(components, REPORT_TO_ID)

    if not check_valid_input(output_dir):
        gr.Warning("No output_dir is given")

    create_training_args_cells(notebook['cells'], epochs=epochs, max_steps=max_steps, logging_steps=logging_steps,
                               per_device_train_batch_size=per_device_train_batch_size, save_strategy=save_strategy,
                               gradient_accumulation_steps=gradient_accumulation_steps,
                               gradient_checkpointing=gradient_checkpointing, learning_rate=learning_rate,
                               max_grad_norm=max_grad_norm, warmup_ratio=warmup_ratio,
                               lr_scheduler_type=lr_scheduler_type, output_dir=output_dir, report_to=report_to,
                               seed=seed_value)

    max_seq_length = get_value(components, MAX_SEQ_LENGTH_ID)
    packing = get_value(components, PACKING_ID)
    create_sft_trainer_cells(notebook['cells'], max_seq_length, packing)

    push_to_hub = get_value(components, PUSH_TO_HUB_ID)

    create_start_training_cells(notebook['cells'], epochs, max_steps, push_to_hub, output_dir)

    create_free_gpu_cells(notebook['cells'])

    create_merge_lora_cells(notebook['cells'], output_dir)

    merge_model_cells(notebook['cells'], output_dir)

    create_readme = get_value(components, README_ID)
    if create_readme:
        create_hf_card(notebook['cells'], name=output_dir, base_model_name=model_value,
                       base_model_version=version_value,
                       dataset_name=dataset_value, output_dir=output_dir, report_to=report_to)

    if push_to_hub:
        if not should_login:
            create_login_hf_cells(notebook['cells'])
        push_to_hub_cells(notebook['cells'], output_dir)

    file_name = f"{finetuning_notebook}.ipynb"

    with open(file_name, 'w') as f:
        nbf.write(notebook, f)

    return gr.Button(
        visible=True), f'''<div class="a_custom"><a href="file={file_name}" download={file_name}>
        💾️ Download {finetuning_notebook}.ipynb</a> </div> ''', "<div></div>"


with gr.Blocks(css=css, theme=gr.themes.Soft(text_size='lg', font=["monospace"],
                                             primary_hue=gr.themes.colors.blue)) as demo:
    gr.Label("UI-Guided LLM FineTuning Jupyter Notebook Generator 🛠️🧠", show_label=False)
    gr.Markdown(
        'Generating a **Jupyter Notebook file (.ipynb)** 📔⚙️ for **finetuning** a Large Language Model (**LLM**) '
        '🎚️🧠 on a chosen dataset and configured parameters, guided by an intuitive User Interface (UI) 👆💻.',
        elem_classes=["center_text"])

    all_components: Set[Component] = set()

    gr.HTML("<h2 style='text-align: center;'>LLM 🧠</h2>")
    with gr.Row():
        model_selection = gr.Dropdown(
            [model.name for model in models],
            elem_id=MODEL_SELECTION_ID,
            label="Select a Large Language Model (LLM)",
            info="Select a Large Language Model (LLM) to finetune using the SFTTrainer."
        )
        version_selection = gr.Dropdown(
            choices=[], label="Select a Model Version 🔄", info="", visible=False, elem_id=MODEL_VERSION_SELECTION_ID
        )
        all_components.add(model_selection)
        all_components.add(version_selection)

    gr.HTML("<h2 style='text-align: center;'>Dataset 📊</h2>")
    with gr.Row():
        all_components.update(add_dataset_components())

    gr.HTML("<h2 style='text-align: center;'>⚡ Flash Attention ⚡</h2>")
    with gr.Row():
        flash_attention = gr.Checkbox(value=True, label="Enable Flash Attention", interactive=True,
                                      elem_id=FLASH_ATTENTION_ID,
                                      info="Flash Attention is a technique that reduces the memory and runtime costs "
                                           "associated with "
                                           "the attention layer in a model. For more details, please refer to the "
                                           "Flash Attention "
                                           "repository on GitHub.")
        all_components.add(flash_attention)

    gr.HTML("<h2 style='text-align: center;'>Quantization</h2>")
    with gr.Row():
        with centered_column():
            all_components.update(add_quantization_components())
        with centered_column():
            all_components.update(add_quantization_components1())

    gr.HTML("<h2 style='text-align: center;'>Tokenizer Configuration</h2>")
    with gr.Row():
        all_components.update(add_pad_tokens())

    gr.HTML("<h2 style='text-align: center;'>Lora Configuration</h2>")
    with gr.Row():
        with centered_column():
            all_components.update(add_lora_components1())
        with centered_column():
            all_components.update(add_lora_components())

    gr.HTML("<h2 style='text-align: center;'>⚙️ Training Arguments ⚙️</h2>")
    with gr.Row():
        with centered_column():
            all_components.update(add_training_args_1())
            all_components.update(add_training_args_1_bis())
        with centered_column():
            all_components.update(add_training_args_3())

    gr.HTML("<h2 style='text-align: center;'>Optimizer Arguments</h2>")
    with gr.Row():
        with centered_column():
            optimizer1 = add_optimizer1()
            all_components.update(optimizer1)

        with centered_column():
            optimizer = add_optimizer()
            all_components.update(optimizer)

    gr.HTML("<h2 style='text-align: center;'>Outputs</h2>")
    with gr.Row():
        with centered_column():
            output_dir_cmp, push_to_hub_cmp = add_outputs()
            all_components.update({output_dir_cmp, push_to_hub_cmp})
        with centered_column():
            all_components.update(add_outputs1())

    gr.HTML("<h2 style='text-align: center;'>SFTTrainer Arguments</h2>")
    with gr.Row():
        sft_args = add_sft_trainer_args()
        all_components.update(sft_args)

    with gr.Row():
        iframe = gr.HTML(show_label=False, visible=True)

    with gr.Row():
        greet_btn = gr.Button("Generate 🛠️", variant="primary")

    with gr.Row():
        preview_btn = gr.Button(f"👀 Preview {finetuning_notebook}.ipynb", variant="primary", visible=False)
        download_btn = gr.HTML(show_label=False, visible=True)

    greet_btn.click(fn=generate_code, inputs=all_components, outputs=[preview_btn, download_btn, iframe])

    preview_btn.click(fn=preview_notebook, inputs=None, outputs=iframe)

    model_selection.change(
        fn=change_model_selection,
        inputs=model_selection,
        outputs=version_selection
    )

demo.launch(allowed_paths=["/"])

# Upload metrics to the hub....
"""
import os
from huggingface_hub import Repository

# Create a repository object
repo = Repository("Menouar/ft-phi-1")

# Push the runs directory
os.system(f"git -C {repo.local_dir} add output_dir/runs")
repo.git_commit("Adding TensorBoard logs")
repo.push_to_hub(commit_message="Adding TensorBoard logs")

"""