"""
Script để tải UW-Madison GI Tract Image Segmentation dataset từ Kaggle
Cài đặt Kaggle API trước:
    pip install kaggle
    Tạo API key từ https://www.kaggle.com/account và lưu vào ~/.kaggle/kaggle.json
"""

import os
import subprocess
import shutil
from pathlib import Path

def setup_kaggle_api():
    """Kiểm tra Kaggle API được cài đặt"""
    try:
        import kaggle
        print("✓ Kaggle API đã được cài đặt")
        return True
    except ImportError:
        print("✗ Kaggle API chưa được cài đặt")
        print("Cài đặt: pip install kaggle")
        return False

def check_kaggle_credentials():
    """Kiểm tra Kaggle credentials - hỗ trợ cả file JSON và environment variable"""
    
    # Kiểm tra KAGGLE_API_TOKEN environment variable (phương thức mới)
    kaggle_token = os.environ.get("KAGGLE_API_TOKEN")
    if kaggle_token:
        print("✓ Kaggle API Token được tìm thấy từ environment variable")
        return True
    
    # Kiểm tra KAGGLE_USERNAME và KAGGLE_KEY (phương thức cũ với env vars)
    kaggle_username = os.environ.get("KAGGLE_USERNAME")
    kaggle_key = os.environ.get("KAGGLE_KEY")
    if kaggle_username and kaggle_key:
        print("✓ Kaggle credentials được tìm thấy từ environment variables")
        return True
    
    # Kiểm tra file kaggle.json
    kaggle_dir = Path.home() / ".kaggle"
    kaggle_json = kaggle_dir / "kaggle.json"
    
    if not kaggle_json.exists():
        print("\n⚠️  Kaggle credentials không được tìm thấy!")
        print("\nCách 1: Sử dụng Environment Variable (khuyến nghị)")
        print("  PowerShell: $env:KAGGLE_API_TOKEN = 'your_token'")
        print("  Hoặc tạo token tại: https://www.kaggle.com/settings")
        print("\nCách 2: Sử dụng file kaggle.json")
        print("  1. Truy cập: https://www.kaggle.com/account")
        print("  2. Scroll xuống, click 'Create New API Token'")
        print("  3. File kaggle.json sẽ tải xuống")
        print(f"  4. Di chuyển file vào: {kaggle_dir}")
        return False
    
    # Đặt permissions (Linux/Mac)
    if not os.name == 'nt':  # Không phải Windows
        os.chmod(str(kaggle_json), 0o600)
    
    print("✓ Kaggle credentials được tìm thấy từ kaggle.json")
    return True

def download_dataset(competition_name="uw-madison-gi-tract-image-segmentation", 
                    output_dir="./data"):
    """Tải dataset từ Kaggle competition"""
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    print(f"\n📥 Đang tải dataset từ Kaggle competition: {competition_name}")
    print(f"📁 Thư mục đích: {output_path.absolute()}")
    
    try:
        cmd = [
            "kaggle", "competitions", "download",
            "-c", competition_name,
            "-p", str(output_path)
        ]
        
        print(f"\n⏳ Đang tải... (Điều này có thể mất vài phút)")
        result = subprocess.run(cmd, check=True)
        
        if result.returncode == 0:
            print("✓ Tải xuống thành công!")
            return True
        else:
            print("✗ Lỗi khi tải xuống")
            return False
    except subprocess.CalledProcessError as e:
        print(f"✗ Lỗi: {e}")
        print("Kiểm tra Kaggle credentials hoặc kết nối internet")
        return False
    except Exception as e:
        print(f"✗ Lỗi: {e}")
        return False

def extract_dataset(data_dir="./data"):
    """Giải nén các file ZIP trong thư mục"""
    data_path = Path(data_dir)
    
    if not data_path.exists():
        print(f"✗ Thư mục {data_dir} không tồn tại")
        return False
    
    zip_files = list(data_path.glob("*.zip"))
    
    if not zip_files:
        print("ℹ️  Không có file ZIP để giải nén")
        return True
    
    print(f"\n📦 Đang giải nén {len(zip_files)} file(s)...")
    
    try:
        for zip_file in zip_files:
            print(f"  → {zip_file.name}")
            shutil.unpack_archive(zip_file, data_path)
            zip_file.unlink()  # Xóa file ZIP sau khi giải nén
        
        print("✓ Giải nén thành công!")
        return True
    except Exception as e:
        print(f"✗ Lỗi: {e}")
        return False

def verify_dataset(data_dir="./data"):
    """Kiểm tra cấu trúc dataset"""
    data_path = Path(data_dir)
    
    required_dirs = ["train_images", "train_masks", "test_images"]
    existing_dirs = []
    
    print("\n🔍 Kiểm tra cấu trúc dataset:")
    
    for dir_name in required_dirs:
        dir_path = data_path / dir_name
        if dir_path.exists():
            files_count = len(list(dir_path.glob("*")))
            print(f"  ✓ {dir_name}: {files_count} files")
            existing_dirs.append(dir_name)
        else:
            print(f"  ✗ {dir_name}: không tìm thấy")
    
    return len(existing_dirs) > 0

def main():
    """Main function"""
    print("=" * 60)
    print("🎯 UW-Madison GI Tract Dataset Downloader")
    print("=" * 60)
    
    # 1. Kiểm tra Kaggle API
    if not setup_kaggle_api():
        print("\n⚠️  Vui lòng cài đặt Kaggle API trước")
        return False
    
    # 2. Kiểm tra Kaggle credentials
    if not check_kaggle_credentials():
        print("\n⚠️  Vui lòng cấu hình Kaggle credentials")
        return False
    
    # 3. Tải dataset
    if not download_dataset():
        return False
    
    # 4. Giải nén dataset
    if not extract_dataset():
        return False
    
    # 5. Kiểm tra dataset
    if not verify_dataset():
        print("\n⚠️  Dataset có thể bị lỗi, vui lòng kiểm tra thủ công")
        return False
    
    print("\n" + "=" * 60)
    print("✅ Dataset đã sẵn sàng! Tiếp theo:")
    print("   1. Chạy: python prepare_dataset.py")
    print("   2. Sau đó: python train.py")
    print("=" * 60)
    
    return True

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)