#!/usr/bin/env python3
"""
fast_downloader.py
Multi-threaded downloader with progress, speed, ETA, and Android-style headers.
Usage:
    python3 fast_downloader.py "https://example.com/file.apk" --out myfile.apk --threads 8
"""

import argparse
import os
import requests
import threading
import time
from math import ceil

# Android-like headers
ANDROID_HEADERS = {
    "User-Agent": "Dalvik/2.1.0 (Linux; U; Android 10; Pixel 3 Build/QP1A.190711.020)",
    "Accept": "*/*",
    "Accept-Encoding": "identity",   # important for range downloads so we get raw bytes
    "Connection": "keep-alive",
}

CHUNK = 8192  # 8 KB


class ProgressTracker:
    def __init__(self, total_size):
        self.total = total_size
        self.lock = threading.Lock()
        self.downloaded = 0
        self.start_time = time.time()
        self.last_print = 0

    def add(self, n):
        with self.lock:
            self.downloaded += n

    def status(self):
        with self.lock:
            dl = self.downloaded
        now = time.time()
        elapsed = max(0.0001, now - self.start_time)
        speed = dl / elapsed  # bytes/sec
        percent = (dl / self.total * 100) if self.total else 0
        remaining = (self.total - dl) if self.total else 0
        eta = (remaining / speed) if speed > 0 else float('inf')
        return dl, self.total, percent, speed, eta

    def pretty(self):
        dl, total, percent, speed, eta = self.status()
        def fmt_bytes(b):
            for u in ['B','KB','MB','GB','TB']:
                if b < 1024.0:
                    return f"{b:6.2f}{u}"
                b /= 1024.0
            return f"{b:.2f}PB"
        return f"{percent:6.2f}%  {fmt_bytes(dl)}/{fmt_bytes(total)}   {fmt_bytes(speed)}/s   ETA {format_time(eta)}"


def format_time(sec):
    if sec == float('inf'):
        return "--:--:--"
    sec = int(sec)
    h = sec // 3600
    m = (sec % 3600) // 60
    s = sec % 60
    return f"{h:02d}:{m:02d}:{s:02d}"


def get_head(url, headers):
    return requests.head(url, headers=headers, allow_redirects=True, timeout=15)


def supports_ranges(resp):
    ar = resp.headers.get("Accept-Ranges", "").lower()
    # some servers return 'bytes' or others; treat 'none' as unsupported
    return ar and ar != "none"


def download_range(url, start, end, part_file, headers, tracker):
    # Range header expects inclusive bytes: start-end
    h = headers.copy()
    h["Range"] = f"bytes={start}-{end}"
    try:
        r = requests.get(url, headers=h, stream=True, timeout=60)
        r.raise_for_status()
        with open(part_file, "wb") as f:
            for chunk in r.iter_content(chunk_size=CHUNK):
                if not chunk:
                    continue
                f.write(chunk)
                tracker.add(len(chunk))
    except Exception as e:
        # mark failure by deleting incomplete part
        if os.path.exists(part_file):
            # keep partial for resume? better to remove and let caller retry
            os.remove(part_file)
        raise


def single_stream_download(url, out_path, headers, tracker):
    # single-threaded streaming download (resumable if server supports Range)
    mode = "wb"
    local_size = 0
    if os.path.exists(out_path):
        local_size = os.path.getsize(out_path)
    h = headers.copy()
    if local_size > 0:
        h["Range"] = f"bytes={local_size}-"
        mode = "ab"
    r = requests.get(url, headers=h, stream=True, timeout=60)
    r.raise_for_status()
    # If server returned full file despite Range, maybe restarting; adjust tracker
    content_length = r.headers.get("Content-Length")
    if content_length:
        try:
            tracker.total = local_size + int(content_length)
        except:
            pass
    with open(out_path, mode) as f:
        for chunk in r.iter_content(chunk_size=CHUNK):
            if not chunk:
                continue
            f.write(chunk)
            tracker.add(len(chunk))


def merge_parts(part_files, out_path):
    with open(out_path, "wb") as outfile:
        for p in part_files:
            with open(p, "rb") as pf:
                while True:
                    data = pf.read(CHUNK)
                    if not data:
                        break
                    outfile.write(data)
    # optionally remove part files
    for p in part_files:
        os.remove(p)


def multi_thread_download(url, out_path, headers, tracker, threads=8):
    # send a HEAD to get Content-Length and range support
    head = get_head(url, headers)
    head.raise_for_status()
    cl = head.headers.get("Content-Length")
    if cl is None:
        raise RuntimeError("Server did not provide Content-Length; cannot do multi-threaded download.")
    total = int(cl)
    tracker.total = total

    if not supports_ranges(head):
        raise RuntimeError("Server does not support range requests.")

    # split into ranges
    part_size = ceil(total / threads)
    part_files = []
    workers = []
    exceptions = []
    for i in range(threads):
        start = i * part_size
        end = min(total - 1, (i + 1) * part_size - 1)
        part_file = f"{out_path}.part{i}"
        part_files.append(part_file)
        # If part exists and has expected size, skip downloading
        if os.path.exists(part_file):
            existing = os.path.getsize(part_file)
            expected = end - start + 1
            if existing == expected:
                tracker.add(existing)
                continue
            else:
                # remove to restart
                os.remove(part_file)

        t = threading.Thread(
            target=_thread_wrapper,
            args=(url, start, end, part_file, headers, tracker, exceptions),
            daemon=True,
        )
        workers.append(t)

    # start all threads
    for w in workers:
        w.start()

    last_print = 0
    try:
        while any(w.is_alive() for w in workers):
            now = time.time()
            if now - last_print >= 0.5:
                print("\r" + tracker.pretty(), end="", flush=True)
                last_print = now
            time.sleep(0.1)
        # join
        for w in workers:
            w.join()
    except KeyboardInterrupt:
        print("\nDownload interrupted by user. Partial parts remain.")
        raise

    if exceptions:
        raise exceptions[0]

    # final print
    print("\r" + tracker.pretty())
    merge_parts(part_files, out_path)


def _thread_wrapper(url, start, end, part_file, headers, tracker, exceptions):
    try:
        download_range(url, start, end, part_file, headers, tracker)
    except Exception as e:
        exceptions.append(e)


def human_readable(n):
    for unit in ['B','KB','MB','GB','TB']:
        if n < 1024:
            return f"{n:.2f}{unit}"
        n /= 1024
    return f"{n:.2f}PB"


def main():
    parser = argparse.ArgumentParser(description="Fast downloader with progress/speed/ETA")
    parser.add_argument("url", help="URL to download")
    parser.add_argument("--out", "-o", help="Output filename (default: basename of url)", default=None)
    parser.add_argument("--threads", "-t", help="Number of threads for multi-conn download (default 8)", type=int, default=8)
    parser.add_argument("--no-android-headers", action="store_true", help="Don't use Android-style headers")
    args = parser.parse_args()

    url = args.url
    out = args.out if args.out else os.path.basename(url.split("?")[0]) or "download.bin"
    threads = max(1, args.threads)

    headers = {} if args.no_android_headers else ANDROID_HEADERS.copy()

    # Probe server
    try:
        head = get_head(url, headers)
        head.raise_for_status()
    except Exception as e:
        print(f"HEAD request failed: {e}")
        print("Trying simple GET to see if URL is reachable...")
        try:
            r = requests.get(url, headers=headers, stream=True, timeout=20)
            r.raise_for_status()
            head = r
        except Exception as e2:
            print(f"Failed to reach URL: {e2}")
            return

    content_length = head.headers.get("Content-Length")
    if content_length:
        total_size = int(content_length)
    else:
        total_size = None

    tracker = ProgressTracker(total_size or 0)

    try:
        if content_length and supports_ranges(head) and threads > 1:
            print(f"Server supports Range requests. Starting multi-threaded download using {threads} threads.")
            tracker.total = int(content_length)
            multi_thread_download(url, out, headers, tracker, threads=threads)
        else:
            if not content_length:
                print("No Content-Length; doing streaming download (single-threaded).")
                tracker.total = 0
            else:
                print("Range not supported or single-thread chosen; doing single-stream download.")
                tracker.total = int(content_length) if content_length else 0
            # show progress in background thread
            tprint = threading.Thread(target=_print_progress_loop, args=(tracker,), daemon=True)
            tprint.start()
            single_stream_download(url, out, headers, tracker)
            # final status
            print("\r" + tracker.pretty())
        print(f"Download complete: {out} ({human_readable(os.path.getsize(out))})")
    except Exception as e:
        print(f"\nError during download: {e}")


def _print_progress_loop(tracker):
    last = 0
    while True:
        print("\r" + tracker.pretty(), end="", flush=True)
        time.sleep(0.5)
        # stop when downloaded >= total (if total known)
        if tracker.total and tracker.downloaded >= tracker.total:
            break


if __name__ == "__main__":
    main()