import os
import urllib.request
from ultralytics import YOLO

# 1. Định nghĩa cấu trúc đường dẫn thư mục Assets của dự án Android
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
ASSETS_MODELS_DIR = os.path.join(PROJECT_ROOT, "android-project", "app", "src", "main", "assets", "models")

# Tạo thư mục nếu chưa tồn tại
os.makedirs(ASSETS_MODELS_DIR, exist_ok=True)

print(f"🚀 Bắt đầu tiến trình tự động hóa tài nguyên AI...")
print(f"📂 Thư mục đích: {ASSETS_MODELS_DIR}\n" + "-"*50)

# =====================================================================
# PHẦN 1: TỰ ĐỘNG TẢI VÀ CHUYỂN ĐỔI YOLOV11-SEGMENTATION
# =====================================================================
try:
    print("⏳ 1. Đang tải và cấu hình YOLOv11-Segmentation (Bản Nano)...")
    # Ultralytics sẽ tự động tải file .pt từ bản phát hành chính thức nếu chưa có sẵn
    yolo_model = YOLO("yolov11n-seg.pt")
    
    print("🔄 Đang export YOLOv11-Seg sang định dạng TFLite (Float16 tối ưu GPU)...")
    # Cấu hình imgsz=640 khớp hoàn toàn với kiến trúc xử lý của DualAiEngine
    exported_path = yolo_model.export(format="tflite", imgsz=640, half=True, int8=False)
    
    # Tìm file .tflite vừa được sinh ra trong thư mục kết quả của Ultralytics
    # Mặc định cấu trúc đầu ra là: yolov11n-seg_saved_model/yolov11n-seg_float16.tflite
    yolo_generated_file = os.path.join(PROJECT_ROOT, "yolov11n-seg_saved_model", "yolov11n-seg_float16.tflite")
    
    if os.path.exists(yolo_generated_file):
        dest_yolo_path = os.path.join(ASSETS_MODELS_DIR, "yolov11n_seg_portrait.tflite")
        os.replace(yolo_generated_file, dest_yolo_path)
        print(f"✅ Đã chuyển đổi và di chuyển YOLOv11 thành công vào: {dest_yolo_path}")
    else:
        print("❌ Lỗi: Không tìm thấy file TFLite của YOLO sau khi export!")
except Exception as e:
    print(f"❌ Thất bại khi xử lý YOLOv11: {str(e)}")

print("-"*50)

# =====================================================================
# PHẦN 2: TỰ ĐỘNG TẢI GOOGLE MEDIAPIPE FACE MESH
# =====================================================================
# URL tải trực tiếp mô hình Face Landmarker (Bundle gồm Face Mesh) chính thức của Google
MEDIAPIPE_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task"
dest_mesh_path = os.path.join(ASSETS_MODELS_DIR, "face_mesh_landmark.tflite")

try:
    print("⏳ 2. Đang tải trực tiếp mô hình Google Face Mesh từ Cloud Storage...")
    
    # Thiết lập thanh tiến trình tải (Progress Bar) đơn giản
    def download_progress(count, block_size, total_size):
        percent = int(count * block_size * 100 / total_size)
        print(f"\r📥 Đang tải: {percent}%", end="")

    urllib.request.urlretrieve(MEDIAPIPE_MODEL_URL, dest_mesh_path, download_progress)
    print(f"\n✅ Đã tải và đổi đuôi mô hình Face Mesh thành công vào: {dest_mesh_path}")
    
except Exception as e:
    print(f"\n❌ Thất bại khi tải mô hình Face Mesh: {str(e)}")

print("="*50)
print("🎉 HOÀN THÀNH Quy trình tự động hóa cung cấp mô hình AI cho dự án!")