import requests
import os
from google.colab import output

# ==========================================
# CONFIGURATION
# ==========================================

# 1. Paste the Civitai Download Link here
# (Get this by right-clicking the download button on the website and selecting 'Copy Link')
download_link = "" #@param {type:"string"}

# 2. (Optional) Specify a custom filename (including extension, e.g., model.safetensors)
# Leave empty to use the original filename from the server.
custom_filename = "" #@param {type:"string"}

# 3. Download destination folder
destination_folder = "/content/downloaded_models" #@param {type:"string"}

# ==========================================
# SCRIPT
# ==========================================

def download_civitai_model(url, folder, filename=None):
    if not url:
        print("Error: Please paste the download link in the 'download_link' variable.")
        return

    # Create folder if it doesn't exist
    os.makedirs(folder, exist_ok=True)
    
    print(f"Connecting to Civitai...")

    # Send a HEAD request first to get headers (allows us to find filename without downloading yet)
    # We use a User-Agent because some servers block python-requests default agent
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
    }
    
    try:
        # First, let's try to get the headers to determine filename
        response = requests.head(url, headers=headers, allow_redirects=True)
        
        # Determine filename
        if not filename:
            # Try to get filename from Content-Disposition header
            content_disp = response.headers.get('Content-Disposition')
            if content_disp and 'filename=' in content_disp:
                # Parse filename from header (handles quotes)
                fname_str = content_disp.split('filename=')[1].strip('"')
                final_filename = fname_str
            else:
                # Fallback if header missing (extract from URL or use generic name)
                print("Warning: Could not determine filename automatically. Using 'downloaded_model.safetensors'")
                final_filename = "downloaded_model.safetensors"
        else:
            final_filename = filename

        file_path = os.path.join(folder, final_filename)

        print(f"Downloading: {final_filename}")
        print(f"Saving to: {file_path}")
        
        # Stream download to handle large files
        response = requests.get(url, headers=headers, stream=True)
        
        # Check for errors
        if response.status_code != 200:
            print(f"Error: Failed to download. Status Code: {response.status_code}")
            print("Response:", response.text[:500])
            return

        # Write file
        total_size = int(response.headers.get('content-length', 0))
        bytes_downloaded = 0
        
        with open(file_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    bytes_downloaded += len(chunk)
                    # Print progress
                    if total_size > 0:
                        percent = (bytes_downloaded / total_size) * 100
                        print(f"\rProgress: {percent:.1f}% ({bytes_downloaded/1024/1024:.1f} MB / {total_size/1024/1024:.1f} MB)", end="")

        print("\n\n✅ Download Complete!")
        
        # List file info
        file_size = os.path.getsize(file_path) / (1024 * 1024)
        print(f"File Size: {file_size:.2f} MB")

    except Exception as e:
        print(f"\nAn error occurred: {e}")

# Run the function
download_civitai_model(download_link, destination_folder, custom_filename)