from typing import TYPE_CHECKING, Dict, Generator, List

import gradio as gr

from ...train import export_model
from ..common import get_save_dir
from ..locales import ALERTS


if TYPE_CHECKING:
    from gradio.components import Component

    from ..engine import Engine


GPTQ_BITS = ["8", "4", "3", "2"]


def save_model(
    lang: str,
    model_name: str,
    model_path: str,
    adapter_path: List[str],
    finetuning_type: str,
    template: str,
    max_shard_size: int,
    export_quantization_bit: int,
    export_quantization_dataset: str,
    export_dir: str,
) -> Generator[str, None, None]:
    error = ""
    if not model_name:
        error = ALERTS["err_no_model"][lang]
    elif not model_path:
        error = ALERTS["err_no_path"][lang]
    elif not export_dir:
        error = ALERTS["err_no_export_dir"][lang]
    elif export_quantization_bit in GPTQ_BITS and not export_quantization_dataset:
        error = ALERTS["err_no_dataset"][lang]
    elif export_quantization_bit not in GPTQ_BITS and not adapter_path:
        error = ALERTS["err_no_adapter"][lang]

    if error:
        gr.Warning(error)
        yield error
        return

    if adapter_path:
        adapter_name_or_path = ",".join(
            [get_save_dir(model_name, finetuning_type, adapter) for adapter in adapter_path]
        )
    else:
        adapter_name_or_path = None

    args = dict(
        model_name_or_path=model_path,
        adapter_name_or_path=adapter_name_or_path,
        finetuning_type=finetuning_type,
        template=template,
        export_dir=export_dir,
        export_size=max_shard_size,
        export_quantization_bit=int(export_quantization_bit) if export_quantization_bit in GPTQ_BITS else None,
        export_quantization_dataset=export_quantization_dataset,
    )

    yield ALERTS["info_exporting"][lang]
    export_model(args)
    yield ALERTS["info_exported"][lang]


def create_export_tab(engine: "Engine") -> Dict[str, "Component"]:
    with gr.Row():
        max_shard_size = gr.Slider(value=1, minimum=1, maximum=100)
        export_quantization_bit = gr.Dropdown(choices=["none", "8", "4", "3", "2"], value="none")
        export_quantization_dataset = gr.Textbox(value="data/c4_demo.json")

    export_dir = gr.Textbox()
    export_btn = gr.Button()
    info_box = gr.Textbox(show_label=False, interactive=False)

    export_btn.click(
        save_model,
        [
            engine.manager.get_elem_by_name("top.lang"),
            engine.manager.get_elem_by_name("top.model_name"),
            engine.manager.get_elem_by_name("top.model_path"),
            engine.manager.get_elem_by_name("top.adapter_path"),
            engine.manager.get_elem_by_name("top.finetuning_type"),
            engine.manager.get_elem_by_name("top.template"),
            max_shard_size,
            export_quantization_bit,
            export_quantization_dataset,
            export_dir,
        ],
        [info_box],
    )

    return dict(
        max_shard_size=max_shard_size,
        export_quantization_bit=export_quantization_bit,
        export_quantization_dataset=export_quantization_dataset,
        export_dir=export_dir,
        export_btn=export_btn,
        info_box=info_box,
    )