from huggingface_hub import InferenceClient
import gradio as gr
import os
import re
import requests
import http.client
import typing
import urllib.request
import vertexai
from vertexai.generative_models import GenerativeModel, Image

with open(".config/application_default_credentials.json", 'w') as file:
    file.write(str(os.getenv('credentials')))

vertexai.init(project=os.getenv('project_id'))
model = GenerativeModel("gemini-1.0-pro-vision")
# client = InferenceClient("google/gemma-7b-it")

def extract_image_urls(text):
    url_regex = r"(https?:\/\/.*\.(?:png|jpg|jpeg|gif|webp|svg))"
    image_urls = re.findall(url_regex, text, flags=re.IGNORECASE)
    valid_image_url = ""
    for url in image_urls:
        try:
            response = requests.head(url)  # Use HEAD request for efficiency
            if response.status_code in range(200, 300) and 'image' in response.headers.get('content-type', ''):
                valid_image_url = url
        except requests.exceptions.RequestException:
            pass  # Ignore inaccessible URLs
    return valid_image_url

def load_image_from_url(image_url: str) -> Image:
    with urllib.request.urlopen(image_url) as response:
        response = typing.cast(http.client.HTTPResponse, response)
        image_bytes = response.read()
    return Image.from_bytes(image_bytes)

def search(url):
    image = load_image_from_url(url)
    response = model.generate_content([image,"Describe what is shown in this image."])
    return response.text

# def format_prompt(message, history, cust_p):
#     prompt = ""
#     for user_prompt, bot_response in history:
#         prompt += f"<start_of_turn>user{user_prompt}<end_of_turn>"
#         prompt += f"<start_of_turn>model{bot_response}<end_of_turn>"
#     prompt += cust_p.replace("USER_INPUT",message)
#     return prompt

# def generate(
#     prompt, history, system_prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,
# ):
#     custom_prompt="<start_of_turn>userUSER_INPUT<end_of_turn><start_of_turn>model"
#     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,
#     )

#     image = extract_image_urls(prompt)
#     if image:
#         image_description = "Image Description: " + search(image)
#         prompt = prompt.replace(image, image_description)
#         print(prompt)
#     formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history, custom_prompt)
#     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
#     output = ""

#     for response in stream:
#         output += response.token.text
#         yield output
#     return output


# additional_inputs=[
#     gr.Textbox(
#         label="System Prompt",
#         max_lines=1,
#         interactive=True,
#     ),
#     gr.Slider(
#         label="Temperature",
#         value=0.9,
#         minimum=0.0,
#         maximum=1.0,
#         step=0.05,
#         interactive=True,
#         info="Higher values produce more diverse outputs",
#     ),
#     gr.Slider(
#         label="Max new tokens",
#         value=256,
#         minimum=0,
#         maximum=1048,
#         step=64,
#         interactive=True,
#         info="The maximum numbers of new tokens",
#     ),
#     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",
#     ),
#     gr.Slider(
#         label="Repetition penalty",
#         value=1.2,
#         minimum=1.0,
#         maximum=2.0,
#         step=0.05,
#         interactive=True,
#         info="Penalize repeated tokens",
#     )
# ]

# examples=[["What are they doing here https://upload.wikimedia.org/wikipedia/commons/3/38/Two_dancers.jpg ?", None, None, None, None, None]]

# gr.ChatInterface(
#     fn=generate,
#     chatbot=gr.Chatbot(show_label=True, show_share_button=True, show_copy_button=True, likeable=True, layout="bubble", bubble_full_width=False),
#     additional_inputs=additional_inputs,
#     title="Gemma Gemini Multimodal Chatbot",
#     description="Gemini Sprint submission by Rishiraj Acharya. Uses Google's Gemini 1.0 Pro Vision multimodal model from Vertex AI with Google's Gemma 7B Instruct model from Hugging Face. Google Cloud credits are provided for this project.",
#     theme="Soft",
#     examples=examples,
#     concurrency_limit=20,
# ).launch(show_api=False)




import random

models=[
    "google/gemma-7b",
    "google/gemma-7b-it",
    "google/gemma-2b",
    "google/gemma-2b-it"
]
clients=[
InferenceClient(models[0]),
InferenceClient(models[1]),
InferenceClient(models[2]),
InferenceClient(models[3]),
]

def load_models(inp):
    return gr.update(label=models[inp])

def format_prompt(message, history, cust_p):
    prompt = ""
    if history:
        for user_prompt, bot_response in history:
            prompt += f"<start_of_turn>user{user_prompt}<end_of_turn>"
            prompt += f"<start_of_turn>model{bot_response}<end_of_turn>"
    prompt+=cust_p.replace("USER_INPUT",message)
    return prompt

def chat_inf(system_prompt,prompt,history,memory,client_choice,seed,temp,tokens,top_p,rep_p,chat_mem,cust_p):
    print(client_choice)
    hist_len=0
    client=clients[int(client_choice)-1]
    if not history:
        history = []
        hist_len=0
    if not memory:
        memory = []
        mem_len=0        
    if memory:
        for ea in memory[0-chat_mem:]:
            hist_len+=len(str(ea))
    in_len=len(system_prompt+prompt)+hist_len

    if (in_len+tokens) > 8000:
        history.append((prompt,"Wait, that's too many tokens, please reduce the 'Chat Memory' value, or reduce the 'Max new tokens' value"))
        yield history,memory
    else:
        generate_kwargs = dict(
            temperature=temp,
            max_new_tokens=tokens,
            top_p=top_p,
            repetition_penalty=rep_p,
            do_sample=True,
            seed=seed,
        )
        
        image = extract_image_urls(prompt)
        if image:
            image_description = "Image Description: " + search(image)
            prompt = prompt.replace(image, image_description)
            print(prompt)
            
        if system_prompt:
            formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", memory[0-chat_mem:],cust_p)
        else:
            formatted_prompt = format_prompt(prompt, memory[0-chat_mem:],cust_p)
        stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=True)
        output = ""
        for response in stream:
            output += response.token.text
            yield [(prompt,output)],memory
        history.append((prompt,output))
        memory.append((prompt,output))
        yield history,memory

def clear_fn():
    return None,None,None,None
rand_val=random.randint(1,1111111111111111)

def check_rand(inp,val):
    if inp==True:
        return gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=random.randint(1,1111111111111111))
    else:
        return gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=int(val))
    
with gr.Blocks() as app:
    memory=gr.State()
    gr.HTML("""<center><h1 style='font-size:xx-large;'>Gemma Gemini Multimodal Chatbot</h1><br><h2>Gemini Sprint submission by Rishiraj Acharya. Uses Google's Gemini 1.0 Pro Vision multimodal model from Vertex AI with Google's Gemma 7B Instruct model from Hugging Face. Google Cloud credits are provided for this project.</h2>""")
    chat_b = gr.Chatbot(show_label=True, show_share_button=True, show_copy_button=True, likeable=True, layout="bubble", bubble_full_width=False)
    with gr.Group():
        with gr.Row():
            with gr.Column(scale=3):
                inp = gr.Textbox(label="Prompt")
                sys_inp = gr.Textbox(label="System Prompt (optional)")
                with gr.Accordion("Prompt Format",open=False):
                    custom_prompt=gr.Textbox(label="Modify Prompt Format", info="For testing purposes. 'USER_INPUT' is where 'SYSTEM_PROMPT, PROMPT' will be placed", lines=3,value="<start_of_turn>userUSER_INPUT<end_of_turn><start_of_turn>model")                
                with gr.Row():
                    with gr.Column(scale=2):
                        btn = gr.Button("Chat")
                    with gr.Column(scale=1):
                        with gr.Group():
                            stop_btn=gr.Button("Stop")
                            clear_btn=gr.Button("Clear")                
                client_choice=gr.Dropdown(label="Models",type='index',choices=[c for c in models],value=models[0],interactive=True)
            with gr.Column(scale=1):
                with gr.Group():
                    rand = gr.Checkbox(label="Random Seed", value=True)
                    seed=gr.Slider(label="Seed", minimum=1, maximum=1111111111111111,step=1, value=rand_val)
                    tokens = gr.Slider(label="Max new tokens",value=1600,minimum=0,maximum=8000,step=64,interactive=True, visible=True,info="The maximum number of tokens")
                    temp=gr.Slider(label="Temperature",step=0.01, minimum=0.01, maximum=1.0, value=0.49)
                    top_p=gr.Slider(label="Top-P",step=0.01, minimum=0.01, maximum=1.0, value=0.49)
                    rep_p=gr.Slider(label="Repetition Penalty",step=0.01, minimum=0.1, maximum=2.0, value=0.99)
                    chat_mem=gr.Number(label="Chat Memory", info="Number of previous chats to retain",value=4)

    
    client_choice.change(load_models,client_choice,[chat_b])
    app.load(load_models,client_choice,[chat_b])
    
    chat_sub=inp.submit(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,memory,client_choice,seed,temp,tokens,top_p,rep_p,chat_mem,custom_prompt],[chat_b,memory])
    go=btn.click(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,memory,client_choice,seed,temp,tokens,top_p,rep_p,chat_mem,custom_prompt],[chat_b,memory])
    
    stop_btn.click(None,None,None,cancels=[go,im_go,chat_sub])
    clear_btn.click(clear_fn,None,[inp,sys_inp,chat_b,memory])
app.queue(default_concurrency_limit=10).launch()