import os
import sys
import time
import requests

def download_file(url, dest_path):
    print(f"Downloading from: {url}")
    print(f"Saving to: {dest_path}")
    
    # Ensure destination directory exists
    dest_dir = os.path.dirname(dest_path)
    os.makedirs(dest_dir, exist_ok=True)
    
    # We will support resuming if the file exists and is smaller
    temp_path = dest_path + ".download"
    initial_bytes = 0
    headers = {}
    
    if os.path.exists(temp_path):
        initial_bytes = os.path.getsize(temp_path)
        headers["Range"] = f"bytes={initial_bytes}-"
        print(f"Resuming download from byte offset: {initial_bytes:,}")
    
    response = requests.get(url, stream=True, headers=headers, timeout=30)
    
    # Check if range request was accepted
    if response.status_code == 206:
        total_size = int(response.headers.get('content-range').split('/')[-1])
        mode = "ab"
    elif response.status_code == 200:
        total_size = int(response.headers.get('content-length', 0))
        initial_bytes = 0
        mode = "wb"
    else:
        print(f"Failed to start download. HTTP Status Code: {response.status_code}")
        if response.status_code == 416:
            print("Range Not Satisfiable. The file may already be fully downloaded or server has issues.")
        return False
        
    print(f"Total target size: {total_size:,} bytes ({total_size / (1024**3):.2f} GB)")
    
    downloaded = initial_bytes
    start_time = time.time()
    last_print = time.time()
    last_downloaded = downloaded
    
    with open(temp_path, mode) as f:
        for chunk in response.iter_content(chunk_size=1024 * 1024): # 1MB chunks
            if chunk:
                f.write(chunk)
                downloaded += len(chunk)
                
                # Print progress every 2 seconds
                current_time = time.time()
                if current_time - last_print >= 2.0:
                    elapsed = current_time - start_time
                    speed = (downloaded - last_downloaded) / (current_time - last_print) if current_time > last_print else 0
                    percent = (downloaded / total_size) * 100 if total_size > 0 else 0
                    eta = (total_size - downloaded) / speed if speed > 0 else 0
                    
                    print(f"Progress: {percent:.2f}% | "
                          f"{downloaded / (1024**2):.1f}/{total_size / (1024**2):.1f} MB | "
                          f"Speed: {speed / (1024**2):.2f} MB/s | "
                          f"ETA: {int(eta // 60)}m {int(eta % 60)}s")
                          
                    last_print = current_time
                    last_downloaded = downloaded
                    
    # Rename temp to target
    if os.path.exists(dest_path):
        os.remove(dest_path)
    os.rename(temp_path, dest_path)
    print("Download completed successfully!")
    return True

if __name__ == "__main__":
    url = "https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/resolve/main/model-00001-of-00002.safetensors"
    dest = r"C:\Users\BHUVAN GIRISH DESAI\.cache\huggingface\hub\models--microsoft--Phi-3-mini-4k-instruct\snapshots\f39ac1d28e925b323eae81227eaba4464caced4e\model-00001-of-00002.safetensors"
    
    success = download_file(url, dest)
    if success:
        # Clean up stale locks and incomplete files
        blobs_dir = r"C:\Users\BHUVAN GIRISH DESAI\.cache\huggingface\hub\models--microsoft--Phi-3-mini-4k-instruct\blobs"
        if os.path.exists(blobs_dir):
            for file in os.listdir(blobs_dir):
                if file.endswith(".incomplete"):
                    try:
                        os.remove(os.path.join(blobs_dir, file))
                        print(f"Cleaned up incomplete blob file: {file}")
                    except Exception as e:
                        print(f"Could not clean up incomplete file: {e}")
                        
        print("\nModel files are fully complete! You can now load the model instantly via the Web Dashboard.")
    else:
        sys.exit(1)