import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ใช้ Hugging Face Cache Directory
MODEL_NAME = "deepseek-ai/deepseek-llm-7b-chat"
MODEL_DIR = os.getenv("HF_HOME", "/tmp")

# ✅ โหลด Tokenizer และ Model ครั้งเดียวเท่านั้น
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=MODEL_DIR)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto",
    cache_dir=MODEL_DIR
)

def classify_esi(symptoms):
    """
    ใช้โมเดล AI ประเมินระดับ ESI จากข้อความอาการของผู้ป่วย
    """
    prompt = (
        "คุณเป็นแพทย์เวชศาสตร์ฉุกเฉินที่ต้องประเมิน ESI 1-5 ตามอาการที่ได้รับ..."
        f"\nอาการของผู้ป่วย: {symptoms}\n"
        "กรุณาประเมิน ESI เป็นตัวเลข 1-5 เท่านั้น ถ้าข้อมูลไม่พอให้ตอบว่า 'ไม่สามารถประเมินได้'"
    )

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=5,  # ✅ จำกัดคำตอบให้สั้น
            do_sample=False  # ✅ ปิดการสุ่มเพื่อลดเวลาประมวลผล
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    for esi in ["1", "2", "3", "4", "5"]:
        if esi in response:
            return int(esi)

    return "ไม่สามารถประเมินได้"