import json
import os
import shutil
import requests
import spaces
import torch

import gradio as gr
from huggingface_hub import Repository
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

HF_TOKEN = os.environ.get("HF_TOKEN", None)
CHECKPOINT_URL = "Salesforce/codegen-350M-mono"

SQLMODEL_PREFIX_URL = "luna-code/sqlmodel-codegen-350M-mono-prefix"
SFEPY_PREFIX_URL = "luna-code/sfepy-codegen-350M-mono-prefix"
MEGENGINE_PREFIX_URL = "luna-code/megengine-codegen-350M-mono-prefix"
MAIN_EVO_PREFIX_URL = "luna-code/codegen-350M-mono-evo-prefix"

SQLMODEL_FFT_URL = "luna-code/sqlmodel-codegen-350M-mono-fft"
SFEPY_FFT_URL = "luna-code/sfepy-codegen-350M-mono-fft"
MEGENGINE_FFT_URL = "luna-code/megengine-codegen-350M-mono-fft"
MAIN_EVO_FFT_URL = "luna-code/codegen-350M-mono-evo-fft"
MAIN_FD_FFT_URL = "luna-code/codegen-350M-mono-fd-fft"

LANGCHAIN_PREFIX_URL = "luna-code/langchain-codegen-350M-mono-prefix"
LLAMAINDEX_PREFIX_URL = "luna-code/llamaindex-codegen-350M-mono-prefix"
DSPY_PREFIX_URL = "luna-code/dspy-codegen-350M-mono-prefix"
CS_EVO_PREFIX_URL = "luna-code/cs-codegen-350M-mono-evo-prefix"

tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_URL)
basemodel = AutoModelForCausalLM.from_pretrained(CHECKPOINT_URL)

sql_prefix = PeftModel.from_pretrained(basemodel, SQLMODEL_PREFIX_URL)
sfepy_prefix = PeftModel.from_pretrained(basemodel, SFEPY_PREFIX_URL)
megengine_prefix = PeftModel.from_pretrained(basemodel, MEGENGINE_PREFIX_URL)
main_evo_prefix = PeftModel.from_pretrained(basemodel, MAIN_EVO_PREFIX_URL)

sqlmodel_fft = AutoModelForCausalLM.from_pretrained(SQLMODEL_FFT_URL)
sfepy_fft = AutoModelForCausalLM.from_pretrained(SFEPY_FFT_URL)
megengine_fft = AutoModelForCausalLM.from_pretrained(MEGENGINE_FFT_URL)
main_evo_fft = AutoModelForCausalLM.from_pretrained(MAIN_EVO_FFT_URL)
main_fd_fft = AutoModelForCausalLM.from_pretrained(MAIN_FD_FFT_URL)

langchain_prefix = PeftModel.from_pretrained(basemodel, LANGCHAIN_PREFIX_URL)
llamaindex_prefix = PeftModel.from_pretrained(basemodel, LLAMAINDEX_PREFIX_URL)
dspy_prefix = PeftModel.from_pretrained(basemodel, DSPY_PREFIX_URL)
cs_evo_prefix = PeftModel.from_pretrained(basemodel, CS_EVO_PREFIX_URL)

# basemodel = ""
# sql_prefix = ""
# sfepy_prefix = ""
# megengine_prefix = ""
# main_evo_prefix = ""
# sqlmodel_fft = ""
# sfepy_fft = ""
# megengine_fft = ""
# main_evo_fft = ""
# main_fd_fft = ""
# langchain_prefix = ""
# llamaindex_prefix = ""
# dspy_prefix = ""
# cs_evo_prefix = ""


model_map = {
    "Base": basemodel,
    "SQLModel Prefix": sql_prefix,
    "SfePy Prefix": sfepy_prefix,
    "MegEngine Prefix": megengine_prefix,
    "Main Evo Prefix": main_evo_prefix,
    "SQLModel FFT": sqlmodel_fft,
    "SfePy FFT": sfepy_fft,
    "MegEngine FFT": megengine_fft,
    "Main Evo FFT": main_evo_fft,
    "Main FD FFT": main_fd_fft,
    "LangChain Prefix": langchain_prefix,
    "LlamaIndex Prefix": llamaindex_prefix,
    "DSpy Prefix": dspy_prefix,
    "CS Evo Prefix": cs_evo_prefix,
}


FIM_PREFIX = "<fim_prefix>"
FIM_MIDDLE = "<fim_middle>"
FIM_SUFFIX = "<fim_suffix>"

FIM_INDICATOR = "<FILL_HERE>"

FORMATS = """## Model Formats

The model is pretrained on code and is formatted with special tokens in addition to the pure code data,\
such as prefixes specifying the source of the file or tokens separating code from a commit message.\
Use these templates to explore the model's capacities:

### 1. Prefixes 🏷️
For pure code files, use any combination of the following prefixes:

```
<reponame>REPONAME<filename>FILENAME<gh_stars>STARS\ncode<|endoftext|>
```
STARS can be one of: 0, 1-10, 10-100, 100-1000, 1000+

### 2. Commits 💾
The commits data is formatted as follows:

```
<commit_before>code<commit_msg>text<commit_after>code<|endoftext|>
```

### 3. Jupyter Notebooks 📓
The model is trained on Jupyter notebooks as Python scripts and structured formats like:

```
<start_jupyter><jupyter_text>text<jupyter_code>code<jupyter_output>output<jupyter_text>
```

### 4. Issues 🐛
We also trained on GitHub issues using the following formatting:
```
<issue_start><issue_comment>text<issue_comment>...<issue_closed>
```

### 5. Fill-in-the-middle 🧩
Fill in the middle requires rearranging the model inputs. The playground handles this for you - all you need is to specify where to fill:
```
code before<FILL_HERE>code after
```
"""

theme = gr.themes.Monochrome(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    radius_size=gr.themes.sizes.radius_sm,
    font=[
        gr.themes.GoogleFont("Open Sans"),
        "ui-sans-serif",
        "system-ui",
        "sans-serif",
    ],
)

@spaces.GPU(enable_queue=True)
def stream(model, code, generate_kwargs):
    model.to(device)
    input_ids = tokenizer(code, return_tensors="pt").to("cuda")
    generated_ids = model.generate(**input_ids, **generate_kwargs)
    return tokenizer.decode(generated_ids[0][input_ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()

@spaces.GPU(enable_queue=True)
def generate(
    prompt, temperature=0.6, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, library="LangChain", method="Prefix"
):

    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )

    if method == "Base":
        output = stream(basemodel, prompt, generate_kwargs)
    elif method == "Prefix":
        output = stream(model_map[library + " Prefix"], prompt, generate_kwargs)
    elif method == "Evo Prefix" and library in ["SQLModel", "SfePy", "MegEngine"]:
        output = stream(model_map["Main Evo Prefix"], prompt, generate_kwargs)
    elif method == "FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
        output = stream(model_map[library + " FFT"], prompt, generate_kwargs)
    elif method == "Evo FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
        output = stream(model_map["Main Evo FFT"], prompt, generate_kwargs)
    elif method == "Full Data FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
        output = stream(model_map["Main FD FFT"], prompt, generate_kwargs)
    else:
        output = ""

    return output


examples = [
    "X_train, y_train, X_test, y_test = train_test_split(X, y, test_size=0.1)\n\n# Train a logistic regression model, predict the labels on the test set and compute the accuracy score",
    "// Returns every other value in the array as a new array.\nfunction everyOther(arr) {",
    "Poor English: She no went to the market. Corrected English:",
    "def alternating(list1, list2):\n   results = []\n   for i in range(min(len(list1), len(list2))):\n       results.append(list1[i])\n       results.append(list2[i])\n   if len(list1) > len(list2):\n       <FILL_HERE>\n   else:\n       results.extend(list2[i+1:])\n   return results",
]


def process_example(args):
    for x in generate(args):
        pass
    return x


css = ".generating {visibility: hidden}"

monospace_css = """
#q-input textarea {
    font-family: monospace, 'Consolas', Courier, monospace;
}
"""


css += share_btn_css + monospace_css + ".gradio-container {color: black}"


description = """
<div style="text-align: center;">
    <h1> 🌙 LUNA Models Playground</h1>
</div>
<div style="text-align: left;">
    <p>This is a demo to generate text and code with unknown libraries. The supported based model is <a href="Salesforce/codegen-350M-mono" style='color: #e6b800;'>CodeGen-350M-mono</a></p>
    <p>The supported libraries are:</p>
    <ul>
        <li><a href="https://sqlmodel.tiangolo.com" style='color: #e6b800;'>SQLModel</a></li>
        <li><a href="https://sfepy.org" style='color: #e6b800;'>SfePy</a></li>
        <li><a href="https://megengine.org" style='color: #e6b800;'>MegEngine</a></li>
        <li><a href="https://www.langchain.com/" style='color: #e6b800;'>LangChain</a></li>
        <li><a href="https://www.llamaindex.ai/" style='color: #e6b800;'>LlamaIndex</a></li>
        <li><a href="https://dspy-docs.vercel.app/" style='color: #e6b800;'>DSpy</a></li>
    </ul>
    <p><b>Please note:</b> These models are not designed for instruction purposes.</p>
</div>
"""
disclaimer = """⚠️<b>Any use or sharing of this demo constitues your acceptance of the BigCode [OpenRAIL-M](spaces/bigcode/bigcode-model-license-agreement) License Agreement and the use restrictions included within.</b>\
 <br>**Intended Use**: this app and its [supporting model](bigcode) are provided for demonstration purposes; not to serve as replacement for human expertise. For more details on the model's limitations in terms of factuality and biases, see the [model card.](hf.co/bigcode)"""

with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
    with gr.Column():
        gr.Markdown(description)
        with gr.Row():
            library = gr.Dropdown(
                        ["SQLModel", "SfePy", "MegEngine", "LangChain", "LlamaIndex", "DSpy"],
                        value="LangChain",
                        label="Library",
                        info="Choose a library from the list",
                        )
        with gr.Row():
            method = gr.Dropdown(
                        ["Base", "Prefix", "Evo Prefix", "FFT", "Evo FFT", "Full Data FFT"],
                        value="Prefix",
                        label="Model",
                        info="Choose an expert from the list",
                        )
        with gr.Row():
            with gr.Column():
                instruction = gr.Textbox(
                    placeholder="Enter your code here",
                    lines=5,
                    label="Input",
                    elem_id="q-input",
                )
                submit = gr.Button("Generate", variant="primary")
                output = gr.Code(elem_id="q-output", lines=30, label="Output")
                with gr.Row():
                    with gr.Column():
                        with gr.Accordion("Advanced settings", open=False):
                            with gr.Row():
                                column_1, column_2 = gr.Column(), gr.Column()
                                with column_1:
                                    temperature = gr.Slider(
                                        label="Temperature",
                                        value=0.2,
                                        minimum=0.0,
                                        maximum=1.0,
                                        step=0.05,
                                        interactive=True,
                                        info="Higher values produce more diverse outputs",
                                    )
                                    max_new_tokens = gr.Slider(
                                        label="Max new tokens",
                                        value=256,
                                        minimum=0,
                                        maximum=8192,
                                        step=64,
                                        interactive=True,
                                        info="The maximum numbers of new tokens",
                                    )
                                with column_2:
                                    top_p = gr.Slider(
                                        label="Top-p (nucleus sampling)",
                                        value=0.90,
                                        minimum=0.0,
                                        maximum=1,
                                        step=0.05,
                                        interactive=True,
                                        info="Higher values sample more low-probability tokens",
                                    )
                                    repetition_penalty = gr.Slider(
                                        label="Repetition penalty",
                                        value=1.2,
                                        minimum=1.0,
                                        maximum=2.0,
                                        step=0.05,
                                        interactive=True,
                                        info="Penalize repeated tokens",
                                    )
                                    
                gr.Markdown(disclaimer)
                with gr.Group(elem_id="share-btn-container"):
                    community_icon = gr.HTML(community_icon_html, visible=True)
                    loading_icon = gr.HTML(loading_icon_html, visible=True)
                    share_button = gr.Button(
                        "Share to community", elem_id="share-btn", visible=True
                    )
                gr.Examples(
                    examples=examples,
                    inputs=[instruction],
                    cache_examples=False,
                    fn=process_example,
                    outputs=[output],
                )
                gr.Markdown(FORMATS)

    submit.click(
        generate,
        inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty, library, method],
        outputs=[output]
    )
    share_button.click(None, [], [])
demo.queue().launch(debug=True)