import os
from pathlib import Path
import gradio as gr
import re
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import requests

HF_TOKEN = os.getenv("HF_TOKEN")

API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
headers = {"Authorization": f"Bearer {HF_TOKEN}"}

def init_speech_to_text_model():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

    model_id = "distil-whisper/distil-large-v2"
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
    )
    model.to(device)
    processor = AutoProcessor.from_pretrained(model_id)
    return pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        max_new_tokens=128,
        torch_dtype=torch_dtype,
        device=device,
    )


whisper_pipe = init_speech_to_text_model()

code_pattern = r'```python\n(.*?)```'

starting_app_code = """import gradio as gr

def greet(name):
    return "Hello " + name + "!"

with gr.Blocks(theme="monochrome") as demo:
    name = gr.Textbox(label="Name", value="World")
    output = gr.Textbox(label="Output Box")
    greet_btn = gr.Button("Greet")
    greet_btn.click(fn=greet, inputs=name, outputs=output)
    name.submit(fn=greet, inputs=name, outputs=output)

if __name__ == "__main__":
    demo.css = "footer {visibility: hidden}"
    demo.launch()
"""

html_template = Path('gradio-lite-playground.html').read_text()
pattern = r"# APP CODE START(.*?)# APP CODE END"

load_js = f"""() => {{
const htmlString = '<iframe class="my-frame" width="100%" height="512px" src="about:blank"></iframe>';
const parser = new DOMParser();
const doc = parser.parseFromString(htmlString, 'text/html');
const iframe = doc.querySelector('.my-frame'); 
const div = document.getElementById('demoDiv');
div.appendChild(iframe);

const frame = document.querySelector('.my-frame');
frame.contentWindow.document.open('text/html', 'replace');
frame.contentWindow.document.write(`{html_template}`);
frame.contentWindow.document.close();
}}"""

#  TODO: Works but is inefficient because the iframe has to be reloaded each time
update_iframe_js = f"""(code) => {{
console.log(`UPDATING CODE`);
console.log(code)
const pattern = /# APP CODE START(.*?)# APP CODE END/gs;
const template = `{html_template}`;
const completedTemplate = template.replace(pattern, code);

console.log(completedTemplate);

const oldFrame = document.querySelector('.my-frame');
oldFrame.remove();

const htmlString = '<iframe class="my-frame" width="100%" height="512px" src="about:blank"></iframe>';
const parser = new DOMParser();
const doc = parser.parseFromString(htmlString, 'text/html');
const iframe = doc.querySelector('.my-frame'); 
const div = document.getElementById('demoDiv');
div.appendChild(iframe);

const frame = document.querySelector('.my-frame');
frame.contentWindow.document.open('text/html', 'replace');
frame.contentWindow.document.write(completedTemplate);
frame.contentWindow.document.close();
console.log(`UPDATE DONE`);
}}"""

copy_snippet_js = f"""async (code) => {{
console.log(`DOWNLOADING CODE`);
const pattern = /# APP CODE START(.*?)# APP CODE END/gs;
const template = `<div id="KiteWindApp">\n<script type="module" crossorigin src="https://cdn.jsdelivr.net/npm/@gradio/lite/dist/lite.js"></script>
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/@gradio/lite/dist/lite.css" />
<gradio-lite>\n# APP CODE START\n\n# APP CODE END\n</gradio-lite>\n</div>\n`;
// Step 1: Generate the HTML content
const completedTemplate = template.replace(pattern, code);

const snippet = completedTemplate;
console.log(snippet);

await navigator.clipboard.writeText(snippet);

console.log(`COPY DONE`);
}}"""

download_code_js = f"""(code) => {{
console.log(`DOWNLOADING CODE`);
const pattern = /# APP CODE START(.*?)# APP CODE END/gs;
const template = `{html_template}`;
// Step 1: Generate the HTML content
const completedTemplate = template.replace(pattern, code);

// Step 2: Create a Blob from the HTML content
const blob = new Blob([completedTemplate], {{ type: "text/html" }});

// Step 3: Create a URL for the Blob
const url = URL.createObjectURL(blob);

// Step 4: Create a download link
const downloadLink = document.createElement("a");
downloadLink.href = url;
downloadLink.download = "gradio-lite-app.html"; // Specify the filename for the download

// Step 5: Trigger a click event on the download link
downloadLink.click();

// Clean up by revoking the URL
URL.revokeObjectURL(url);

console.log(`DOWNLOAD DONE`);
}}"""


def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()


def generate_text(code, prompt):
    print(f"Calling API with prompt:\n{prompt}")
    prompt = f"```python\n{code}```\nGiven the code above return only updated code for the following request:\n{prompt}\n<|assistant|>"
    params = {"max_new_tokens": 512}
    output = query({
        "inputs": prompt,
        "parameters": params,
    })
    print(f'API RESPONSE\n{output[0]["generated_text"]}')
    assistant_reply = output[0]["generated_text"].split('<|assistant|>')[1]
    match = re.search(code_pattern, assistant_reply, re.DOTALL)
    new_code = match.group(1)
    print(new_code)
    # TODO: error handling here
    return assistant_reply, new_code


def transcribe(audio):
    result = whisper_pipe(audio)
    return result["text"], None


def copy_notify(code):
    gr.Info("App code snippet copied!")


with gr.Blocks() as demo:
    gr.Markdown("<h1 align=\"center\">KiteWind 🪁🍃</h1>")
    gr.Markdown("<h4 align=\"center\">Chat-assisted web app creator by <a href=\"https://huggingface.co/gstaff\">@gstaff</a></h4>")
    with gr.Row():
        with gr.Column():
            gr.Markdown("## 1. Run your app in the browser!")
            html = gr.HTML(value='<div id="demoDiv"></div>')
    gr.Markdown("## 2. Customize using voice requests!")
    with gr.Row():
        with gr.Column():
            with gr.Group():
                in_audio = gr.Audio(label="Record a voice request", source='microphone', type='filepath')
                in_prompt = gr.Textbox(label="Or type a text request and press Enter",
                                       placeholder="Need an idea? Try one of these:\n- Add a button to reverse the name\n- Change the greeting to Hola\n- Put the reversed name output into a separate textbox\n- Change the theme from monochrome to soft")
            out_text = gr.TextArea(label="Chat Assistant Response")
            clear = gr.ClearButton([in_prompt, in_audio, out_text])
        with gr.Column():
            code_area = gr.Code(label="App Code - You can also edit directly and then click Update App", language='python', value=starting_app_code)
            update_btn = gr.Button("Update App", variant="primary")
            update_btn.click(None, inputs=code_area, outputs=None, _js=update_iframe_js)
            in_prompt.submit(generate_text, [code_area, in_prompt], [out_text, code_area]).then(None, inputs=code_area, outputs=None, _js=update_iframe_js)
            in_audio.stop_recording(transcribe, [in_audio], [in_prompt, in_audio]).then(generate_text, [code_area, in_prompt], [out_text, code_area]).then(None, inputs=code_area, outputs=None, _js=update_iframe_js)
    with gr.Row():
        with gr.Column():
            gr.Markdown("## 3. Export your app to share!")
            copy_snippet_btn = gr.Button("Copy app snippet to paste in another page")
            copy_snippet_btn.click(copy_notify, code_area, None, _js=copy_snippet_js)
            download_btn = gr.Button("Download app as a standalone file")
            download_btn.click(None, code_area, None, _js=download_code_js)
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Current limitations")
            with gr.Accordion("Click to view", open=False):
                gr.Markdown("- Only gradio-lite apps using the python standard libraries and gradio are supported\n- The chat hasn't been tuned on gradio library data; it may make mistakes\n- The app needs to fully reload each time it is changed")

    demo.load(None, None, None, _js=load_js)
    demo.css = "footer {visibility: hidden}"

if __name__ == "__main__":
    demo.queue().launch()