import os
import zipfile
from huggingface_hub import login
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig

# 🔹 گرفتن توکن
hf_token = os.environ.get("HF_TOKEN")
if not hf_token:
    raise ValueError("❌ HF_TOKEN not found in environment secrets.")

# 🔹 لاگین
login(hf_token)

# 🔹 مسیرها
LORA_ZIP_PATH = "dorna-diabetes-finetuned-20250514T183411Z-1-001.zip"
LORA_PATH = "dorna-diabetes-finetuned"
BASE_MODEL = "PartAI/Dorna-Llama3-8B-Instruct"

# ✅ اکسترکت
if not os.path.exists(LORA_PATH):
    with zipfile.ZipFile(LORA_ZIP_PATH, "r") as zip_ref:
        zip_ref.extractall(LORA_PATH)
    print("✅ فایل LoRA اکسترکت شد.")

# ✅ تغییر نام safetensors به adapter_model.safetensors
for filename in os.listdir(LORA_PATH):
    if filename.endswith(".safetensors") and filename != "adapter_model.safetensors":
        os.rename(
            os.path.join(LORA_PATH, filename),
            os.path.join(LORA_PATH, "adapter_model.safetensors")
        )
        print("✅ اسم فایل تغییر کرد.")
        break

# ✅ بارگذاری مدل پایه
print("🔹 در حال بارگذاری مدل پایه...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    trust_remote_code=True,
    token=hf_token
)

# ✅ بارگذاری LoRA روی مدل پایه
print("🔹 در حال بارگذاری LoRA...")
model = PeftModel.from_pretrained(base_model, LORA_PATH)
model.eval()

# ✅ توکنایزر
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=hf_token)

print("✅ مدل و توکنایزر با موفقیت بارگذاری شدند.")

# 🧪 تست ساده
while True:
    prompt = input("📝 یک دستور وارد کن (exit برای خروج): ")
    if prompt.lower() == "exit":
        break
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=200)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print("🧠 پاسخ:", response)