import os
import uuid
import subprocess
import gradio as gr
from urllib.parse import urlparse

SAVE_DIR = "/data/downloads"
os.makedirs(SAVE_DIR, exist_ok=True)

def make_filename(url):
    ext = os.path.splitext(urlparse(url).path)[1]
    if not ext or len(ext) > 10:
        ext = ".bin"
    return f"{uuid.uuid4().hex}{ext}"

def download_file(url):
    if not url.startswith(("http://", "https://")):
        return "Only direct HTTP/HTTPS links allowed.", None

    filename = make_filename(url)
    output_path = os.path.join(SAVE_DIR, filename)

    cmd = [
        "aria2c",
        "-c",
        "-x", "16",
        "-s", "16",
        "-k", "8M",
        "--max-tries=0",
        "--retry-wait=10",
        "--timeout=60",
        "--connect-timeout=30",
        "--continue=true",
        "--allow-overwrite=true",
        "--auto-file-renaming=false",
        "-d", SAVE_DIR,
        "-o", filename,
        url
    ]

    try:
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True
        )

        if result.returncode != 0:
            return "aria2c error:\n" + result.stdout[-4000:], None

        if not os.path.exists(output_path):
            return "Download finished but file not found.", None

        size_gb = os.path.getsize(output_path) / 1024**3
        return f"Downloaded successfully: {filename} ({size_gb:.2f} GB)", output_path

    except Exception as e:
        return f"Error: {str(e)}", None

def list_files():
    items = []
    for name in os.listdir(SAVE_DIR):
        path = os.path.join(SAVE_DIR, name)
        if os.path.isfile(path):
            size_gb = os.path.getsize(path) / 1024**3
            items.append(f"{name} - {size_gb:.2f} GB")
    return "\n".join(items) if items else "No files found."

with gr.Blocks() as demo:
    gr.Markdown("# Large File Downloader")
    gr.Markdown("Direct HTTP/HTTPS links only. No torrent/magnet links.")

    url = gr.Textbox(label="Direct URL")
    btn = gr.Button("Download")
    status = gr.Textbox(label="Status", lines=8)
    file_out = gr.File(label="Downloaded File")

    btn.click(download_file, inputs=url, outputs=[status, file_out])

    refresh = gr.Button("Refresh Files")
    files = gr.Textbox(label="Saved Files", lines=10)
    refresh.click(list_files, outputs=files)

demo.launch()