import os
import zipfile
from huggingface_hub import login
import torch
from transformers import AutoTokenizer
from peft import AutoPeftModelForCausalLM

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

# 🔹 لاگین به HuggingFace
login(hf_token)

# 🔹 تنظیم مسیرها
LORA_ZIP_PATH = "dorna-diabetes-finetuned-20250514T183411Z-1-001.zip"  # نام فایل زیپ که آپلود کردی
LORA_PATH = "./dorna-diabetes-finetuned"

# ✅ اکسترکت کردن فایل LoRA
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("🔹 در حال بارگذاری مدل پایه + LoRA...")

model = AutoPeftModelForCausalLM.from_pretrained(
    LORA_PATH,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    device_map="auto",
    token=hf_token,
    trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained(
    model.base_model.config._name_or_path,
    token=hf_token,
    trust_remote_code=True
)

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