# Import các thư viện cần thiết
import os
import json  # Thêm thư viện json để xử lý dữ liệu JSON
from threading import Thread
from typing import Iterator

import gradio as gr
import spaces
import torch
import requests  # Thêm thư viện requests để thực hiện các yêu cầu HTTP
from bs4 import BeautifulSoup  # Thêm BeautifulSoup để phân tích HTML
from functools import lru_cache  # Thêm lru_cache để cache kết quả hàm
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

# Mô tả về mô hình Llama 3.2 3B Instruct
DESCRIPTION = """\
# Llama 3.2 3B Instruct

Llama 3.2 3B là phiên bản mới nhất của LLM từ Meta.
Đây là bản demo của [`meta-llama/Llama-3.2-3B-Instruct`](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct), được tinh chỉnh để theo dõi hướng dẫn.
Để biết thêm chi tiết, hãy xem [bài đăng của chúng tôi](https://huggingface.co/blog/llama32).
"""

# Thiết lập các thông số tối đa
MAX_MAX_NEW_TOKENS = 2048  # Số token tối đa cho đầu ra mới
DEFAULT_MAX_NEW_TOKENS = 1024  # Số token mặc định cho đầu ra mới
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))  # Lấy giá trị chiều dài token đầu vào từ biến môi trường

# Kiểm tra thiết bị có hỗ trợ GPU không, nếu không thì sử dụng CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Định danh mô hình và tải mô hình cùng tokenizer
model_id = "nltpt/Llama-3.2-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",  # Tự động ánh xạ thiết bị
    torch_dtype=torch.bfloat16,  # Sử dụng kiểu dữ liệu bfloat16
)
model.eval()  # Đặt mô hình vào chế độ đánh giá (evaluation mode)

# Hàm để trích xuất văn bản từ nội dung HTML, sử dụng cache để tăng tốc độ
@lru_cache(maxsize=128)
def extract_text_from_webpage(html_content):
    """Trích xuất văn bản hiển thị từ nội dung HTML sử dụng BeautifulSoup."""
    soup = BeautifulSoup(html_content, "html.parser")
    # Loại bỏ các thẻ không hiển thị như script, style, header, footer, nav, form, svg
    for tag in soup(["script", "style", "header", "footer", "nav", "form", "svg"]):
        tag.extract()
    # Lấy văn bản hiển thị và loại bỏ khoảng trắng thừa
    visible_text = soup.get_text(strip=True)
    return visible_text

# Hàm thực hiện tìm kiếm trên Google và trả về kết quả
def search(query):
    term = query
    all_results = []
    max_chars_per_page = 8000  # Số ký tự tối đa cho mỗi trang
    with requests.Session() as session:
        # Thực hiện yêu cầu tìm kiếm trên Google
        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": 4, "udm": 14},
            timeout=5,
            verify=False,  # Không xác thực SSL
        )
        resp.raise_for_status()  # Kiểm tra phản hồi HTTP
        soup = BeautifulSoup(resp.text, "html.parser")
        result_block = soup.find_all("div", attrs={"class": "g"})  # Tìm các kết quả tìm kiếm
        for result in result_block:
            link = result.find("a", href=True)
            if link:
                link = link["href"]
                try:
                    # Lấy nội dung trang web từ liên kết
                    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)
                    # Cắt bớt văn bản nếu vượt quá giới hạn
                    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

# Danh sách metadata các hàm có thể được gọi
functions_metadata = [
    {
        "name": "web_search",
        "description": "Tìm kiếm trên Google và tìm thông tin mới nhất, thông tin về bất kỳ người, đối tượng, địa điểm nào có sẵn trên Google.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Truy vấn tìm kiếm web"
                }
            },
            "required": ["query"]
        }
    },
    {
        "name": "general_query",
        "description": "Trả lời các truy vấn chung của người dùng bằng LLM như bạn. Nhưng không trả lời các câu hỏi khó và thông tin mới nhất.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "Một đề bài chi tiết"
                }
            },
            "required": ["prompt"]
        }
    },
    {
        "name": "hard_query",
        "description": "Trả lời các truy vấn khó của người dùng, sử dụng LLM mạnh mẽ. Nhưng không trả lời các thông tin mới nhất.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "Một đề bài chi tiết"
                }
            },
            "required": ["prompt"]
        }
    },
]

# Prompt role system với các hướng dẫn về khả năng và công cụ
SYSTEM_PROMPT = f"""\
Bạn là một trợ lý thông minh và hữu ích với khả năng sử dụng các công cụ sau đây để hỗ trợ người dùng:

{json.dumps(functions_metadata, indent=4, ensure_ascii=False)}

Bạn có thể thực hiện các công việc như tìm kiếm trên web, tạo hình ảnh, tạo video, và trả lời các câu hỏi liên quan đến hình ảnh. Khi cần sử dụng một trong các công cụ này, bạn sẽ trả về một lời gọi hàm theo định dạng Python như sau:

[func1(params_name=params_value, params_name2=params_value2...), func2(params)]

Bạn chỉ được gọi một hàm tại một thời điểm. Đảm bảo rằng bạn chọn các công cụ một cách khôn ngoan để cung cấp câu trả lời tốt nhất cho người dùng.

"""

# Hàm chính để xử lý việc sinh phản hồi từ mô hình
@spaces.GPU(duration=90)  # Chỉ định hàm này chạy trên GPU trong tối đa 90 giây
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    max_new_tokens: int = 1024,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    conversation = []
    func_caller = []

    # Xây dựng cuộc hội thoại từ lịch sử trò chuyện
    for user, assistant in chat_history:
        conversation.extend(
            [
                {"role": "user", "content": user},
                {"role": "assistant", "content": assistant},
            ]
        )
    
    # Thêm tin nhắn mới của người dùng vào cuộc hội thoại
    conversation.append({"role": "user", "content": message})

    # Thêm prompt role system vào đầu cuộc hội thoại
    conversation.insert(0, {"role": "system", "content": SYSTEM_PROMPT})

    # Áp dụng mẫu hội thoại và chuyển thành tensor
    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
    
    # Kiểm tra và cắt bớt chuỗi đầu vào nếu vượt quá chiều dài tối đa
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Đã cắt bớt đầu vào từ cuộc hội thoại vì vượt quá {MAX_INPUT_TOKEN_LENGTH} tokens.")
    
    # Chuyển tensor đến thiết bị của mô hình
    input_ids = input_ids.to(model.device)

    # Khởi tạo Streamer để lấy đầu ra theo từng phần (real-time)
    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
    
    # Thiết lập các tham số cho quá trình sinh đầu ra
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        num_beams=1,
        repetition_penalty=repetition_penalty,
    )
    
    # Tạo một luồng để chạy quá trình sinh đầu ra
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        response = "".join(outputs)
        yield response

        # Kiểm tra xem phản hồi có phải là lời gọi hàm không
        if response.startswith("[") and response.endswith("]"):
            try:
                # Phân tích phản hồi JSON
                json_data = json.loads(response)
                
                # Xử lý các hàm được gọi
                for func in json_data:
                    func_name = func.get("name")
                    arguments = func.get("arguments", {})
                    
                    if func_name == "web_search":
                        query = arguments.get("query")
                        if query:
                            gr.Info("Đang tìm kiếm trên web...")
                            yield "Đang tìm kiếm trên web..."
                            web_results = search(query)
                            
                            gr.Info("Đang trích xuất thông tin liên quan...")
                            yield "Đang trích xuất thông tin liên quan..."
                            web_content = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results])
                            
                            # Xây dựng lại cuộc hội thoại với kết quả tìm kiếm
                            conversation_updated = conversation.copy()
                            conversation_updated.append({"role": "user", "content": message})
                            conversation_updated.append({"role": "system", "content": f"[WEB RESULTS] {web_content}"})
                            
                            # Gửi lại cuộc hội thoại để mô hình trả lời dựa trên kết quả tìm kiếm
                            input_ids_updated = tokenizer.apply_chat_template(conversation_updated, add_generation_prompt=True, return_tensors="pt")
                            input_ids_updated = input_ids_updated.to(model.device)
                            
                            streamer_updated = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
                            generate_kwargs_updated = dict(
                                {"input_ids": input_ids_updated},
                                streamer=streamer_updated,
                                max_new_tokens=max_new_tokens,
                                do_sample=True,
                                top_p=top_p,
                                top_k=top_k,
                                temperature=temperature,
                                num_beams=1,
                                repetition_penalty=repetition_penalty,
                            )
                            
                            t_updated = Thread(target=model.generate, kwargs=generate_kwargs_updated)
                            t_updated.start()
                            
                            for text_updated in streamer_updated:
                                yield text_updated
                    
            except json.JSONDecodeError:
                # Nếu phản hồi không phải là JSON hợp lệ, tiếp tục trả lời bình thường
                continue

# Tạo giao diện chat với Gradio
chat_interface = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.Slider(
            label="Số token mới tối đa",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        ),
        gr.Slider(
            label="Nhiệt độ (Temperature)",
            minimum=0.1,
            maximum=4.0,
            step=0.1,
            value=0.6,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.9,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="Hình phạt lặp lại (Repetition penalty)",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.2,
        ),
    ],
    stop_btn=None,  # Không có nút dừng
    examples=[
        ["Xin chào! Bạn có khỏe không?"],
        ["Bạn có thể giải thích ngắn gọn về ngôn ngữ lập trình Python không?"],
        ["Giải thích cốt truyện của Cô bé Lọ Lem trong một câu."],
        ["Mất bao nhiêu giờ để một người ăn một chiếc trực thăng?"],
        ["Viết một bài báo 100 từ về 'Lợi ích của mã nguồn mở trong nghiên cứu AI'"],
    ],
    cache_examples=False,  # Không lưu trữ các ví dụ
)

# Tạo bố cục giao diện với Gradio
with gr.Blocks(css="style.css", fill_height=True) as demo:
    gr.Markdown(DESCRIPTION)  # Hiển thị phần mô tả
    gr.DuplicateButton(value="Tạo bản sao cho sử dụng cá nhân", elem_id="duplicate-button")
    chat_interface.render()  # Hiển thị giao diện chat

# Hàm chính để khởi chạy ứng dụng
if __name__ == "__main__":
    demo.queue(max_size=20).launch()  # Khởi chạy ứng dụng với hàng đợi tối đa 20 yêu cầu