import gradio as gr
from huggingface_hub import InferenceClient
import json
import uuid
from PIL import Image
from bs4 import BeautifulSoup
import requests
import random
from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
from threading import Thread
import re
import time 
import torch
import cv2

model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"

processor = LlavaProcessor.from_pretrained(model_id)

model = LlavaForConditionalGeneration.from_pretrained(model_id)
model.to("cpu")


def llava(message, history):
    if message["files"]:
        image = message["files"][0]    
    else:
        for hist in history:
            if isinstance(hist[0], tuple):
                image = hist[0][0]
        
    txt = message["text"]
        
    gr.Info("Analyzing image")
    image = Image.open(image).convert("RGB")
    prompt = f"<|im_start|>user <image>\n{txt}<|im_end|><|im_start|>assistant"
    
    inputs = processor(prompt, image, return_tensors="pt")
    
    # Return the dictionary format expected by MultimodalTextbox
    return {"text": txt, "files": [image]}


def extract_text_from_webpage(html_content):
    soup = BeautifulSoup(html_content, 'html.parser')
    for tag in soup(["script", "style", "header", "footer"]):
        tag.extract()
    return soup.get_text(strip=True)

def search(query):
    term = query
    start = 0
    all_results = []
    max_chars_per_page = 8000
    with requests.Session() as session:
        resp = session.get(
            url="https://www.google.com/search",
            headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"},
            params={"q": term, "num": 3, "udm": 14},
            timeout=5,
            verify=None,
        )
        resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        result_block = soup.find_all("div", attrs={"class": "g"})
        for result in result_block:
            link = result.find("a", href=True)
            link = link["href"]
            try:
                webpage = session.get(link, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"}, timeout=5, verify=False)
                webpage.raise_for_status()
                visible_text = extract_text_from_webpage(webpage.text)
                if len(visible_text) > max_chars_per_page:
                    visible_text = visible_text[:max_chars_per_page]
                all_results.append({"link": link, "text": visible_text})
            except requests.exceptions.RequestException:
                all_results.append({"link": link, "text": None})
    return all_results

# Initialize inference clients for different models
client_gemma = InferenceClient("google/gemma-1.1-7b-it")
client_mixtral = InferenceClient("NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO")
client_llama = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")

import gradio as gr
from huggingface_hub import InferenceClient
import json
import uuid
from PIL import Image
from bs4 import BeautifulSoup
import requests
import random
from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
from threading import Thread
import re
import time 
import torch
import cv2

model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"

processor = LlavaProcessor.from_pretrained(model_id)

model = LlavaForConditionalGeneration.from_pretrained(model_id)
model.to("cpu")


def llava(message, history):
    if message["files"]:
        image = message["files"][0]    
    else:
        for hist in history:
            if isinstance(hist[0], tuple):
                image = hist[0][0]
        
    txt = message["text"]
        
    gr.Info("Analyzing image")
    image = Image.open(image).convert("RGB")
    prompt = f"<|im_start|>user <image>\n{txt}<|im_end|><|im_start|>assistant"
    
    inputs = processor(prompt, image, return_tensors="pt")
    
    # Return the dictionary format expected by MultimodalTextbox
    return {"text": txt, "files": [image]}


def extract_text_from_webpage(html_content):
    soup = BeautifulSoup(html_content, 'html.parser')
    for tag in soup(["script", "style", "header", "footer"]):
        tag.extract()
    return soup.get_text(strip=True)

def search(query):
    term = query
    start = 0
    all_results = []
    max_chars_per_page = 8000
    with requests.Session() as session:
        resp = session.get(
            url="https://www.google.com/search",
            headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"},
            params={"q": term, "num": 3, "udm": 14},
            timeout=5,
            verify=None,
        )
        resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        result_block = soup.find_all("div", attrs={"class": "g"})
        for result in result_block:
            link = result.find("a", href=True)
            link = link["href"]
            try:
                webpage = session.get(link, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"}, timeout=5, verify=False)
                webpage.raise_for_status()
                visible_text = extract_text_from_webpage(webpage.text)
                if len(visible_text) > max_chars_per_page:
                    visible_text = visible_text[:max_chars_per_page]
                all_results.append({"link": link, "text": visible_text})
            except requests.exceptions.RequestException:
                all_results.append({"link": link, "text": None})
    return all_results

# Initialize inference clients for different models
client_gemma = InferenceClient("google/gemma-1.1-7b-it")
client_mixtral = InferenceClient("NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO")
client_llama = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")

# Corrected the response format to ensure it works with Gradio's multimodal interface.
def respond(message, history):
    func_caller = []

    user_prompt = message
    # Handle image processing
    if message.get("files"):
        inputs = llava(message, history)
        streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})
        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)

        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()

        buffer = ""
        for new_text in streamer:
            buffer += new_text
            yield buffer
    else:
        functions_metadata = [
            {"type": "function", "function": {"name": "web_search", "description": "Search query on google", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "web search query"}}, "required": ["query"]}}},
            {"type": "function", "function": {"name": "general_query", "description": "Reply general query of USER", "parameters": {"type": "object", "properties": {"prompt": {"type": "string", "description": "A detailed prompt"}}, "required": ["prompt"]}}},
            {"type": "function", "function": {"name": "image_generation", "description": "Generate image for user", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "image generation prompt"}, "number_of_image": {"type": "integer", "description": "number of images to generate"}}, "required": ["query"]}}},
            {"type": "function", "function": {"name": "image_qna", "description": "Answer question asked by user related to image", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Question by user"}}, "required": ["query"]}}},
        ]
    
        message_text = message["text"]
        func_caller.append({"role": "user", "content": f'[SYSTEM]You are a helpful AI assistant for a discord server called Stars Kingdom, your job is to have fun help users and listen to what they say or want you to act. You have been created by the discord server owner named Star. You have access to the following functions: \n {str(functions_metadata)}\n\nTo use these functions respond with:\n<functioncall> {{ "name": "function_name", "arguments": {{ "arg_1": "value_1", "arg_2": "value_2", ... }} }}  </functioncall>  [USER] {message_text}'})

        response = client_gemma.chat_completion(func_caller, max_tokens=150)
        response = str(response)
        try:
            response = response[int(response.find("{")):int(response.index("</"))]
        except:
            print("An error occurred")
        response = response.replace("\\n", "")
        response = response.replace("\\'", "'")
        response = response.replace('\\"', '"')
        print(f"\n{response}")
    
        func_caller.append({"role": "assistant", "content": f"<functioncall>{response}</functioncall>"})
    
        try:
            json_data = json.loads(str(response))
            if json_data["name"] == "web_search":
                query = json_data["arguments"]["query"]
                gr.Info("Searching Web")
                web_results = search(query)
                gr.Info("Extracting relevant Info")
                web2 = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results])
                messages = f"<|im_start|>system\nYou are a helpful assistant made by Star. You are provided with WEB results from which you can find informations to answer users query in a structured and better way. Only respond with what’s important!<|im_end|>"
                for msg in history:
                    messages += f"\n<|im_start|>user\n{str(msg[0])}<|im_end|>"
                    messages += f"\n<|im_start|>assistant\n{str(msg[1])}<|im_end|>"
                messages += f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>web_result\n{web2}<|im_end|>\n<|im_start|>assistant\n"
                stream = client_mixtral.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
                output = ""
                for response in stream:
                    if not response.token.text == "<|im_end|>":
                        output += response.token.text
                        yield output
            elif json_data["name"] == "image_generation":
                query = json_data["arguments"]["query"]
                gr.Info("Generating Image, Please wait 10 sec...")
                seed = random.randint(1, 99999)
                image = f"![](https://image.pollinations.ai/prompt/{message_text}{query}?seed={seed}&nologo=True)"
                image = image.replace("\\n", "")
                image = image.replace(" ", "%20")
                yield image
                time.sleep(8)
                gr.Info("We are going to Update Our Image Generation Engine to more powerful ones in Next Update. ThankYou")
            elif json_data["name"] == "image_qna":
                inputs = llava(message, history)
                streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})
                generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)

                thread = Thread(target=model.generate, kwargs=generation_kwargs)
                thread.start()

                buffer = ""
                for new_text in streamer:
                    buffer += new_text
                    yield buffer
            else:
                # Default response from llama model if no specific function matched
                messages = f"<|im_start|>system\nYou are a helpful assistant made by Star.<|im_end|>"
                for msg in history:
                    messages += f"\n<|im_start|>user\n{str(msg[0])}<|im_end|>"
                    messages += f"\n<|im_start|>assistant\n{str(msg[1])}<|im_end|>"
                messages += f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>assistant\n"
                stream = client_llama.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
                output = ""
                for response in stream:
                    if not response.token.text == "<|eot_id|>":
                        output += response.token.text
                        yield output
        except:
            messages = f"<|im_start|>system\nYou are a helpful assistant made by Star. You answer users' queries like a human friend.<|im_end|>"
            for msg in history:
                messages += f"\n<|im_start|>user\n{str(msg[0])}<|im_end|>"
                messages += f"\n<|im_start|>assistant\n{str(msg[1])}<|im_end|>"
            messages += f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>assistant\n"
            stream = client_llama.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
            output = ""
            for response in stream:
                if not response.token.text == "<|eot_id|>":
                    output += response.token.text
                    yield output

# Gradio Chat Interface
demo = gr.ChatInterface(
    fn=respond,
    chatbot=gr.Chatbot(show_copy_button=True, likeable=True, layout="panel"),
    description="AI assistant for Stars Kingdom",
    textbox=gr.MultimodalTextbox(),
    multimodal=True,
    concurrency_limit=200,
    examples=[
        {"text": "What can I wear with a yellow Kurta?",},
        {"text": "What's the preferred shirt color for an interview?",},
        {"text": "How can I dress more smartly?",},
        {"text": "Tell about some good accessories for a traditional Indian wedding",},
        {"text": "What's the color of the frock in the given image?", "files": ["./frock.png"]},
    ],
    cache_examples=False,
)

demo.launch(share=True)