import asyncio
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import gc

# 1. ระบุตำแหน่งที่อยู่ไฟล์ GGUF บน Hugging Face Hub (ดึงเวอร์ชัน Q4_K_M ตามที่เราเลือก)
MODEL_REPOS = {
    "qwen_main": {"repo": "Qwen/Qwen2.5-3B-Instruct-GGUF", "file": "qwen2.5-3b-instruct-q4_k_m.gguf"},
    "qwen_coder": {"repo": "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF", "file": "qwen2.5-coder-3b-instruct-q4_k_m.gguf"},
    "gemma2": {"repo": "bartowski/gemma-2-2b-it-GGUF", "file": "gemma-2-2b-it-Q4_K_M.gguf"},
    "llama3": {"repo": "bartowski/Llama-3.2-3B-Instruct-GGUF", "file": "Llama-3.2-3B-Instruct-Q4_K_M.gguf"},
    "smollm2": {"repo": "HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF", "file": "smollm2-1.7b-instruct-q4_k_m.gguf"}
}

async def run_model_isolated(model_key: str, prompt: str) -> str:
    """
    ฟังก์ชันโหลดโมเดล GGUF เข้าแรมประมวลผลเสร็จแล้วเคลียร์ทิ้งทันที (Offloading)
    """
    print(f"\n[System] 📥 กำลังตรวจสอบและดาวน์โหลดโมเดล: {model_key}...")
    
    # ดาวน์โหลดหรือดึงไฟล์ GGUF จากแคชระบบ
    model_path = hf_hub_download(
        repo_id=MODEL_REPOS[model_key]["repo"],
        filename=MODEL_REPOS[model_key]["file"]
    )
    
    print(f"[System] 🧠 โหลด {model_key} เข้าสู่หน่วยความจำ...")
    # เรียกโมเดลเข้า RAM จำกัด Context ไว้ที่ 2048 เพื่อประหยัดพื้นที่แรม
    llm = Llama(model_path=model_path, n_ctx=2048, n_threads=2)
    
    print(f"[System] ⚡ กำลังประมวลผลคำสั่ง...")
    # สั่งให้โมเดลเจนคำตอบ
    output = llm(f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n", max_tokens=512)
    response_text = output["choices"][0]["text"]
    
    # --- ขั้นตอนเด็ด: ทำลายทิ้งสลัดออกจากแรม ---
    del llm
    gc.collect() # บังคับล้างขยะในหน่วยความจำ RAM ของเครื่องทันที
    print(f"[System] 🧼 เคลียร์แรมของ {model_key} เรียบร้อยแล้ว!")
    
    return response_text

# ทดสอบรันฟังก์ชันสายพานเบื้องต้น
async def main():
    print("--- เริ่มต้นการทดสอบระบบสลับโหลดโมเดลบน Spaces ---")
    
    # ทดสอบเรียกตัวที่ 1 (Qwen Main)
    blueprint = await run_model_isolated("qwen_main", "วางแผนสร้างระบบบล็อกตัวต่อในเกมสร้างฐานสั้นๆ")
    print(f"\n[ผลลัพธ์จากตัวที่ 1]:\n{blueprint}")

if __name__ == "__main__":
    asyncio.run(main())