import os
import glob
import shutil
from ultralytics import YOLO

# 1. Nạp mô hình PyTorch gốc từ file bạn đã tải sẵn
model = YOLO("yolov11n-seg.pt")

print("⏳ Bước 1: Đang dịch sang định dạng trung gian ONNX...")
# Xuất sang file .onnx (Vượt qua bộ check OS của LiteRT)
onnx_path = model.export(format="onnx", imgsz=640)

print("⏳ Bước 2: Kích hoạt onnx2tf để dịch ma trận đồ họa sang TFLite...")
onnx_file = "yolov11n-seg.onnx"

if os.path.exists(onnx_file):
    # Gọi công cụ dòng lệnh chuyển dịch sang Float32
    # -nuo: Chống tối ưu hóa lỗi cấu hình phần cứng Android GPU Delegate
    os.system(f"onnx2tf -in {onnx_file} -nuo")
    
    print("🔍 Đang quét tìm vị trí file .tflite thành phẩm...")
    # 🟢 SỬA LỖI: Tìm kiếm tất cả các file .tflite được sinh ra trong các thư mục con
    tflite_matches = glob.glob("**/yolov11n-seg*.tflite", recursive=True) + glob.glob("*.tflite")
    
    # Lọc bỏ nếu tìm trúng file đích cũ để tránh trùng lặp
    tflite_matches = [f for f in tflite_matches if "yolov11n_seg_portrait.tflite" not in f]

    if tflite_matches:
        generated_tflite = tflite_matches[0]
        dest_path = "yolov11n_seg_portrait.tflite"
        
        # Di chuyển và đổi tên chuẩn xác cho dự án Android
        os.replace(generated_tflite, dest_path)
        print(f"🎉 XUẤT FILE THÀNH CÔNG! File đã nằm tại: {os.path.abspath(dest_path)}")
        
        # --- DỌN DẸP FILE RÁC TRUNG GIAN ---
        try:
            os.remove(onnx_file)
            # Xóa các thư mục tạm do onnx2tf hoặc ultralytics tạo ra
            if os.path.exists("yolov11n-seg_saved_model"):
                shutil.rmtree("yolov11n-seg_saved_model")
            if os.path.exists("saved_model"):
                shutil.rmtree("saved_model")
            # Tìm và xóa thư mục con trùng tên nếu có
            for d in os.listdir("."):
                if os.path.isdir(d) and "yolov11n-seg" in d:
                    shutil.rmtree(d)
        except Exception as clean_ex:
            print(f"⚠️ Cảnh báo dọn dẹp: {str(clean_ex)}")
            
    else:
        print("❌ Lỗi: onnx2tf không tạo ra bất kỳ file .tflite nào! Hãy kiểm tra log phía trên của onnx2tf xem có bị lỗi thiếu công cụ 'flatc' không.")
else:
    print("❌ Lỗi: Không tìm thấy file ONNX trung gian để chuyển đổi!")