import torch
import torch.nn as nn
from transformers import Dinov2Model, AutoImageProcessor
from PIL import Image
import pandas as pd
import os
import glob  # <--- Thư viện quan trọng để tìm file bất chấp đuôi
from tqdm import tqdm
from shared_resources import get_mongo_collection

# --- CẤU HÌNH IMPORT TỪ FILE CONFIG ---
# Nếu file config.py của bạn thiếu biến nào, hãy thêm vào hoặc định nghĩa trực tiếp ở đây
try:
    from core.config import CSV_FILE, BATCH_SIZE
except ImportError:
    # Giá trị mặc định phòng hờ lỗi import
    CSV_FILE = "VN_train_data.csv"
    BATCH_SIZE = 32

# --- CẤU HÌNH ĐƯỜNG DẪN LOCAL ---
# Cấu trúc thư mục: DATASET_ROOT_DIR / landmark_id / image_id.jpg (hoặc png, webp...)
# Ví dụ: dataset/train/1/3521.jpg
DATASET_ROOT_DIR = "dataset/train" 

# --- CẤU HÌNH MODEL ---
MODEL_NAME = "facebook/dinov2-base"
MODEL_PATH = "models/dinov2_hf_finetuned_ep30.pth" 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print(f"🚀 Đang sử dụng thiết bị: {DEVICE}")

# --- 1. KHAI BÁO KIẾN TRÚC MODEL ---
class FineTunedDINOv2(nn.Module):
    def __init__(self):
        super(FineTunedDINOv2, self).__init__()
        # Sử dụng lại kiến trúc bạn đã định nghĩa
        self.backbone = Dinov2Model.from_pretrained(MODEL_NAME)

    def forward(self, x):
        outputs = self.backbone(x)
        cls_token = outputs.last_hidden_state[:, 0]
        # Bắt buộc Chuẩn hóa L2 giống như khi tạo vector trong DB
        return nn.functional.normalize(cls_token, p=2, dim=1)

def load_finetuned_model():
    # Sử dụng lại logic load model đã sửa để xử lý checkpoint dictionary
    print(f"⏳ Đang tải model từ {MODEL_PATH}...")
    model = FineTunedDINOv2()
    
    if os.path.exists(MODEL_PATH):
        checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
        
        if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
            print(f"ℹ️ Đang trích xuất weights từ Checkpoint...")
            state_dict = checkpoint["model_state_dict"]
        else:
            state_dict = checkpoint

        new_state_dict = {}
        for k, v in state_dict.items():
            if k.startswith("module."):
                new_state_dict[k[7:]] = v
            else:
                new_state_dict[k] = v
        
        try:
            model.load_state_dict(new_state_dict)
        except:
            # Tắt strict để xử lý nếu có chênh lệch nhỏ về key
            model.load_state_dict(new_state_dict, strict=False)

        print("✅ Đã load weights thành công!")

    else:
        raise FileNotFoundError(f"❌ Không tìm thấy file model tại {MODEL_PATH}")
    
    model.to(DEVICE)
    model.eval() 
    
    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
    return model, processor

def main():
    # 1. Kết nối MongoDB
    try:
        collection, client = get_mongo_collection()
        print("✅ Kết nối MongoDB thành công.")
    except Exception as e:
        print(f"❌ Dừng chương trình do lỗi kết nối DB: {e}")
        return

    # 2. Xóa dữ liệu cũ (Re-indexing)
    # Bước này quan trọng vì model đã thay đổi, vector cũ không còn tác dụng
    print("⚠️ Đang xóa dữ liệu cũ trong Collection...")
    collection.delete_many({})
    print("✅ Đã xóa sạch collection. Sẵn sàng nạp mới.")

    # 3. Load Model & Processor
    try:
        model, processor = load_finetuned_model()
    except Exception as e:
        print(e)
        return

    # 4. Đọc file CSV
    if not os.path.exists(CSV_FILE):
        print(f"❌ Không tìm thấy file CSV: {CSV_FILE}")
        return
    
    df = pd.read_csv(CSV_FILE)
    print(f"📂 Đã đọc {len(df)} dòng từ file CSV.")

    # 5. Vòng lặp xử lý chính
    documents_batch = []
    missing_files_count = 0
    success_count = 0
    
    print(f"🚀 Bắt đầu Embed dữ liệu từ thư mục: {DATASET_ROOT_DIR}")
    print("👉 Logic: Đọc ID từ CSV -> Tìm file local (bất chấp đuôi) -> Embed -> Lưu MongoDB")
    
    for index, row in tqdm(df.iterrows(), total=df.shape[0], desc="Processing"):
        try:
            # Lấy thông tin metadata từ CSV
            landmark_id = str(row['landmark_id'])
            image_id = str(row['image_id'])
            original_url = row['image_url'] 

            if pd.isna(image_id) or image_id.strip() == "":
                continue

            # --- LOGIC TÌM FILE DÙNG GLOB (Xử lý đa định dạng) ---
            # Tạo đường dẫn mẫu: dataset/train/1/3521.*
            # Dấu * sẽ khớp với mọi đuôi (.jpg, .png, .webp, .JPG...)
            search_path_pattern = os.path.join(DATASET_ROOT_DIR, landmark_id, f"{image_id}.*")
            
            # Tìm tất cả các file khớp mẫu
            found_files = glob.glob(search_path_pattern)

            if not found_files:
                # Nếu list rỗng nghĩa là không tìm thấy file nào khớp ID này
                # print(f"⚠️ Missing file for ID: {image_id}") # Uncomment nếu muốn xem chi tiết
                missing_files_count += 1
                continue
            
            # Lấy file đầu tiên tìm thấy (thường chỉ có 1 file duy nhất khớp ID)
            local_image_path = found_files[0]

            # a. Mở ảnh từ ổ cứng
            try:
                # .convert("RGB") là bắt buộc để tránh lỗi kênh Alpha (trong png) hoặc Grayscale
                img = Image.open(local_image_path).convert("RGB")
            except Exception as e:
                print(f"❌ Lỗi file hỏng {local_image_path}: {e}")
                continue

            # b. Vector hóa (Embedding)
            inputs = processor(images=img, return_tensors="pt").pixel_values.to(DEVICE)
            
            with torch.no_grad():
                embedding_tensor = model(inputs)
            
            # Chuyển Tensor về List chuẩn Python
            vector = embedding_tensor.cpu().numpy().flatten().tolist()
            
            # c. Tạo Document chuẩn cấu trúc
            # QUAN TRỌNG: Metadata lấy từ CSV để đảm bảo hiển thị đúng trên Web/App
            doc = {
                "image_id": image_id,
                "landmark_id": landmark_id,
                "image_url": original_url, # Vẫn lưu URL online
                "embedding": vector
            }
            
            documents_batch.append(doc)
            success_count += 1

            # d. Nạp vào DB theo lô (Batch insert)
            if len(documents_batch) >= BATCH_SIZE:
                collection.insert_many(documents_batch)
                documents_batch = [] # Reset lô

        except Exception as e:
            print(f"⚠️ Lỗi không xác định tại dòng {index}: {e}")
            continue

    # 6. Nạp nốt số còn lại trong lô cuối cùng
    if documents_batch:
        collection.insert_many(documents_batch)

    # 7. Báo cáo kết quả
    print("\n" + "="*50)
    print(f"🎉 HOÀN TẤT QUÁ TRÌNH INGEST!")
    print(f"📊 Tổng số ảnh trong CSV: {len(df)}")
    print(f"✅ Số ảnh xử lý thành công: {success_count}")
    print(f"⚠️ Số ảnh không tìm thấy trong thư mục Local: {missing_files_count}")
    
    # Kiểm tra lại số lượng thực tế trong DB
    try:
        db_count = collection.count_documents({})
        print(f"💾 Số lượng document hiện tại trong MongoDB: {db_count}")
    except:
        pass
        
    print("-" * 50)
    print("👉 LƯU Ý CHO BƯỚC TIẾP THEO:")
    print("1. Vào MongoDB Atlas > Atlas Search.")
    print("2. Tạo hoặc Cập nhật Index.")
    print("3. Đảm bảo field 'embedding' có 'numDimensions': 768 (DINOv2 Base).")
    print("="*50)
    
    client.close()

if __name__ == "__main__":
    main()