import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import requests
import pdfkit
from docx import Document
from docx.shared import Pt
from docx.oxml.ns import qn
import os

# Thiết lập API token từ Hugging Face
api_token = os.getenv("api_token")
if api_token is None:
    raise ValueError("API token is not set. Please set it in your Hugging Face settings.")

# Tải mô hình và tokenizer từ Hugging Face
model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=api_token)
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=api_token).to("cpu")

def analyze_chat(link, extraction_request, system_message, max_tokens, temperature, top_p):
    response = requests.get(link)
    if response.status_code != 200:
        return "Không thể lấy nội dung từ link."

    chat_content = response.text
    prompt = f"{system_message}\n\nChat Content:\n{chat_content}\n\nExtraction Request: {extraction_request}"
    
    inputs = tokenizer(prompt, return_tensors="pt", max_length=8192, truncation=True).to("cpu")
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p)
    
    extracted_content = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return extracted_content

def generate_file(content, file_type):
    if file_type == "PDF":
        pdfkit.from_string(content, 'output.pdf')
        return "output.pdf"
    elif file_type == "Docx":
        doc = Document()
        doc.styles['Normal'].font.name = 'Arial'
        doc.styles['Normal']._element.rPr.rFonts.set(qn('w:eastAsia'), 'Arial')
        p = doc.add_paragraph()
        run = p.add_run(content)
        run.font.size = Pt(12)
        doc.save('output.docx')
        return "output.docx"

with gr.Blocks() as demo:
    gr.Markdown("# Chatbot Trợ Lý Cá Nhân")
    
    link = gr.Textbox(label="Đường link đoạn chat")
    extraction_request = gr.Textbox(label="Yêu cầu trích xuất nội dung")
    analyze_button = gr.Button("Phân tích")
    extracted_output = gr.Textbox(label="Nội dung đã trích xuất", interactive=True)
    generate_button = gr.Button("Tạo file")
    file_type = gr.Radio(["PDF", "Docx"], label="Loại file", value="PDF")
    file_output = gr.File(label="File kết quả")
    
    def analyze_link(link, extraction_request, system_message, max_tokens, temperature, top_p):
        extracted_content = analyze_chat(link, extraction_request, system_message, max_tokens, temperature, top_p)
        return extracted_content

    def generate_file_action(content, file_type):
        file_path = generate_file(content, file_type)
        return file_path

    system_message = gr.Textbox(value="You are a friendly Chatbot.", label="System message")
    max_tokens = gr.Slider(minimum=1, maximum=8192, value=512, step=1, label="Max new tokens")  # Adjust maximum based on model capabilities
    temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")

    analyze_button.click(analyze_link, inputs=[link, extraction_request, system_message, max_tokens, temperature, top_p], outputs=extracted_output)
    generate_button.click(generate_file_action, inputs=[extracted_output, file_type], outputs=file_output)

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