import os
import requests
import zipfile

MODEL_URL = "https://alphacephei.com/vosk/models/vosk-model-vn-0.4.zip"
MODEL_DIR = "model"
MODEL_PATH = os.path.join(MODEL_DIR, "vosk-model-vn-0.4")  # Đổi lại cho đúng với thư mục sau khi giải nén

if not os.path.exists(MODEL_PATH):
    os.makedirs(MODEL_DIR, exist_ok=True)

    zip_path = os.path.join(MODEL_DIR, "model.zip")

    print("📥 Đang tải model Vosk tiếng Việt...")
    response = requests.get(MODEL_URL, stream=True)

    if response.status_code != 200:
        print(f"❌ Lỗi tải model! Mã lỗi: {response.status_code}")
        exit(1)

    with open(zip_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

    print("📦 Kiểm tra file ZIP...")
    if not zipfile.is_zipfile(zip_path):
        print("❌ File tải về không phải là file ZIP hợp lệ!")
        os.remove(zip_path)
        exit(1)

    print("📦 Giải nén model...")
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(MODEL_DIR)

    # Kiểm tra xem thư mục giải nén có tồn tại không
    extracted_path = os.path.join(MODEL_DIR, "vosk-model-vn-0.4")
    if os.path.exists(extracted_path):
        print(f"✅ Model đã được giải nén tại: {extracted_path}")
    else:
        print(f"❌ Lỗi: Model không tồn tại sau khi giải nén! Kiểm tra lại.")
        exit(1)

    os.remove(zip_path)  # Xóa file zip sau khi giải nén

    print("✅ Model tiếng Việt đã tải thành công!")
else:
    print(f"✅ Model đã có sẵn tại {MODEL_PATH}, bỏ qua tải về.")