import gradio as gr
import os
import time
import urllib.request
from pathlib import Path
from urllib.parse import urlparse, parse_qs, unquote

CHUNK_SIZE = 1638400
TOKEN_FILE = Path.home() / '.civitai' / 'config'
USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'

def get_token():
    try:
        with open(TOKEN_FILE, 'r') as file:
            token = file.read()
            return token
    except Exception:
        return None

def store_token(token: str):
    TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(TOKEN_FILE, 'w') as file:
        file.write(token)

def download_file(url: str, output_path: str, token: str):
    headers = {
        'Authorization': f'Bearer {token}',
        'User-Agent': USER_AGENT,
    }

    class NoRedirection(urllib.request.HTTPErrorProcessor):
        def http_response(self, request, response):
            return response
        https_response = http_response

    request = urllib.request.Request(url, headers=headers)
    opener = urllib.request.build_opener(NoRedirection)
    response = opener.open(request)

    if response.status in [301, 302, 303, 307, 308]:
        redirect_url = response.getheader('Location')
        parsed_url = urlparse(redirect_url)
        query_params = parse_qs(parsed_url.query)
        content_disposition = query_params.get('response-content-disposition', [None])[0]

        if content_disposition:
            filename = unquote(content_disposition.split('filename=')[1].strip('"'))
        else:
            return None, f"ERROR: Unable to determine filename"

        response = urllib.request.urlopen(redirect_url)
    elif response.status == 404:
        return None, f"ERROR: File not found"
    else:
        return None, f"ERROR: No redirect found, something went wrong"

    total_size = response.getheader('Content-Length')
    if total_size is not None:
        total_size = int(total_size)

    output_file = os.path.join(output_path, filename)

    with open(output_file, 'wb') as f:
        downloaded = 0
        start_time = time.time()

        while True:
            chunk_start_time = time.time()
            buffer = response.read(CHUNK_SIZE)
            chunk_end_time = time.time()

            if not buffer:
                break

            downloaded += len(buffer)
            f.write(buffer)
            chunk_time = chunk_end_time - chunk_start_time

            if chunk_time > 0:
                speed = len(buffer) / chunk_time / (1024 ** 2)  # Speed in MB/s

            if total_size is not None:
                progress = downloaded / total_size
                gr.Progress.update(f'Downloading: {filename} [{progress*100:.2f}%] - {speed:.2f} MB/s')

    end_time = time.time()
    time_taken = end_time - start_time
    hours, remainder = divmod(time_taken, 3600)
    minutes, seconds = divmod(remainder, 60)

    if hours > 0:
        time_str = f'{int(hours)}h {int(minutes)}m {int(seconds)}s'
    elif minutes > 0:
        time_str = f'{int(minutes)}m {int(seconds)}s'
    else:
        time_str = f'{int(seconds)}s'

    return output_file, f"Download completed. File saved as: {filename}. Downloaded in {time_str}"

def run_downloader(url, output_path, token_input):
    token = get_token() or token_input
    if not token:
        return None, "ERROR: Token not provided."

    store_token(token)
    file_path, message = download_file(url, output_path, token)
    return file_path, message

with gr.Blocks() as demo:
    gr.Markdown("# CivitAI Downloader")
    with gr.Row():
        url_input = gr.Textbox(label="CivitAI Download URL")
        output_path_input = gr.Textbox(label="Output Path", value="./")  # Default value set to current directory
        token_input = gr.Textbox(label="CivitAI API Token (Optional)")
    download_button = gr.Button("Download")
    output_file = gr.File(label="Downloaded Model")
    message_output = gr.Textbox(label="Status Message", interactive=False)
    
    download_button.click(
        run_downloader, 
        inputs=[url_input, output_path_input, token_input],
        outputs=[output_file, message_output]
    )

demo.launch()