import os
import sys
import urllib.request
import argparse

HF_REPO_URL = "https://huggingface.co/kilanisainikhil/AerialEye/resolve/main"

def is_lfs_pointer(path):
    if not os.path.exists(path):
        return False
    if os.path.getsize(path) > 1024:
        return False
    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            first_line = f.readline()
            return first_line.startswith("version https://git-lfs.github.com/spec/v1")
    except Exception:
        return False

def download_file(filename, force=False):
    url = f"{HF_REPO_URL}/{filename}"
    target_path = filename
    
    if os.path.exists(target_path) and not is_lfs_pointer(target_path) and not force:
        print(f"File '{filename}' already exists and is not an LFS pointer. Skipping download.")
        return True
        
    print(f"Downloading {filename} from {url}...")
    try:
        # We can implement a simple progress reporting callback
        def reporthook(blocknum, blocksize, totalsize):
            readsofar = blocknum * blocksize
            if totalsize > 0:
                percent = readsofar * 1e2 / totalsize
                s = f"\rDownloading... {percent:5.1f}% [{readsofar} / {totalsize} bytes]"
                sys.stdout.write(s)
                sys.stdout.flush()
            else:
                sys.stdout.write(f"\rDownloading... {readsofar} bytes")
                sys.stdout.flush()
                
        urllib.request.urlretrieve(url, target_path, reporthook)
        print(f"\nSuccessfully downloaded {filename}.")
        return True
    except Exception as e:
        print(f"\nError downloading {filename}: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description="Download AerialEye model weights and assets from Hugging Face.")
    parser.add_argument("--files", nargs="+", default=["aerialEye.pt", "best.pt", "sutra_standard_vs_sahi.jpg"],
                        help="List of files to download.")
    parser.add_argument("--force", action="store_true", help="Force overwrite existing files even if they are not LFS pointers.")
    parser.add_argument("--all", action="store_true", help="Download all model weights and sample images.")
    
    args = parser.parse_args()
    
    files_to_download = args.files
    if args.all:
        files_to_download = [
            "aerialEye.pt", "aerialEye.onnx", "best.pt", "best.onnx", "best_full_integer_quant.tflite",
            "bus.jpg", "sample_aerial_street.jpg", "sample_drone_roundabout.jpg",
            "sample_train_1.jpg", "sample_train_2.jpg", "sample_train_3.jpg",
            "sutra_standard_vs_sahi.jpg"
        ]
        
    success = True
    for f in files_to_download:
        if not download_file(f, force=args.force):
            success = False
            
    if success:
        print("All downloads completed successfully.")
        sys.exit(0)
    else:
        print("Some downloads failed.")
        sys.exit(1)

if __name__ == "__main__":
    main()