import os
import torch
import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

MODEL_ID = "laurent-maille/Valuoty-industry-plc-7B-V0.9"

# Tải Tokenizer chuẩn
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)

# Cấu hình Chat Template chính xác theo định dạng ChatML của Qwen2.5
CHAT_TEMPLATE = (
    "{% for message in messages %}"
    "{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n' }}"
    "{% endfor %}"
    "{% if add_generation_prompt %}"
    "{{ '<|im_start|>assistant\\n' }}"
    "{% endif %}"
)
tokenizer.chat_template = CHAT_TEMPLATE

@spaces.GPU(duration=90)
def respond(message, chat_history):
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, 
        torch_dtype=torch.float16, 
        device_map="auto"
    )
    
    # 1. Khởi tạo mảng hội thoại chứa System Prompt định hướng AI chuyên ngành
    messages = [
        {
            "role": "system", 
            "content": "Bạn là chuyên gia lập trình PLC chuẩn quốc tế IEC 61131-3. Hãy sửa lỗi cú pháp, tối ưu thuật toán và viết code (Structured Text, Ladder ASCII) chính xác theo dòng máy yêu cầu. Không nói chuyện phiếm."
        }
    ]
    
    # 2. Đồng bộ lịch sử trò chuyện theo định dạng OpenAI-style mới chuẩn Gradio
    for turn in chat_history:
        # turn lúc này là một dict chứa {'role': '...', 'content': '...'}
        if isinstance(turn, dict):
            messages.append({"role": turn.get("role"), "content": turn.get("content")})
            
    # Thêm câu hỏi hiện tại của người dùng vào cuối mảng
    messages.append({"role": "user", "content": message})

    # 3. Biên dịch luồng hội thoại thành dữ liệu GPU tokens
    input_ids = tokenizer.apply_chat_template(
        messages, 
        tokenize=True, 
        add_generation_prompt=True, 
        return_tensors="pt"
    ).to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=20., skip_prompt=True, skip_special_tokens=True)
    
    # Tham số cấu hình logic nghiêm ngặt (Tắt tính năng sáng tạo văn học tự do)
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=2048,
        do_sample=False, 
    )
    
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        # Làm sạch chuỗi hiển thị nếu tokenizer sinh ra các ký tự điều hướng hệ thống
        clean_text = text.replace("<|im_end|>", "").replace("<|endoftext|>", "").replace("<|im_start|>", "")
        outputs.append(clean_text)
        yield "".join(outputs)

# Thiết lập giao diện hiển thị tinh gọn và cấu hình type='messages' để xóa cảnh báo
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🛠️ Trợ Lý AI Lập Trình PLC Cao Cấp (Valuoty-7B-V0.9)")
    gr.Markdown("Hệ thống đã được tối ưu hóa cấu trúc dữ liệu theo tiêu chuẩn Gradio mới nhất.")
    gr.ChatInterface(fn=respond, type="messages")

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