import os
import requests
import threading
import time
import argparse
import mimetypes
import sys
import urllib.parse


class Downloader:
    def __init__(self, url, dest, overide="y", num_threads=4, chunk_size=1024 * 1024, max_retries=3, retry_after=5):
        self.url = urllib.parse.unquote(url)
        self.dest = dest
        self.overide = overide
        self.num_threads = num_threads
        self.file_size = 0
        self.chunk_size = chunk_size
        self.lock = threading.Lock()
        self.running = True
        self.threads = []
        self.response = requests.head(self.url, allow_redirects=True)
        self.headers = self.response.headers
        self.max_retries = max_retries
        self.retry_after = retry_after
        self.total_bytes_downloaded = 0
        self.print_progress_thread = threading.Thread(
            target=self.print_progress)
        self.success = False
        self.failed = False

        if self.response.history:
            print('Request was redirected')
            self.url = self.response.history[-1].headers['Location']
            print(f'Last redirect URL: {self.url}')

        retry_count = 0
        while self.response.status_code != 200:
            self.response = requests.head(self.url)
            self.headers = self.response.headers
            if self.response.status_code == 200:
                break
            else:
                retry_count += 1
                if retry_count >= self.max_retries:
                    raise Exception("Max retries reached, HTTP error code: {}".format(
                        self.response.status_code))
                else:
                    print(
                        f"Retrying url in {self.retry_after} seconds...")
                    time.sleep(self.retry_after)

        if "Content-Length" in self.headers:
            self.file_size = int(self.headers["Content-Length"])
            if(self.file_size <= 0):
                raise ValueError("File size is less than 0")
        else:
            raise ValueError("No content length in headers.")

        if os.path.exists(self.dest):
            if self.overide.lower() != "y":
                raise FileExistsError()
            else:
                os.remove(self.dest)

    def start(self):
        self.init_file()
        self.print_progress_thread.start()

        for i in range(self.num_threads):
            t = threading.Thread(target=self.download, args=(i,))
            t.start()
            self.threads.append(t)

        for t in self.threads:
            t.join()

        self.running = False

        if self.success:
            print("Download successful!")
        elif self.failed:
            raise Exception("Download unsuccessful!")

    def download(self, i):
        start = i * self.chunk_size
        end = start + self.chunk_size - 1
        if end > self.file_size:
            end = self.file_size - 1
        headers = {"Range": "bytes={}-{}".format(start, end)}
        retry_count = 0
        bytes_downloaded = 0
        while self.running and start < self.file_size:
            try:
                res = requests.get(self.url, headers=headers, stream=True)
                with open(self.dest, "r+b") as f:
                    f.seek(start)
                    f.write(res.content)
                bytes_downloaded += len(res.content)
                self.total_bytes_downloaded += len(res.content)
                if not res:
                    break
            except:
                retry_count += 1
                if retry_count > self.max_retries:
                    print("Max retries exceeded, exiting...")
                    self.running = False
                    self.failed = True
                    return
                print(
                    f"Retrying download in {self.retry_after} seconds (retry {retry_count}/{self.max_retries})")
                time.sleep(self.retry_after)
            start += self.chunk_size * self.num_threads
            end += self.chunk_size * self.num_threads
            if end > self.file_size:
                end = self.file_size - 1
            headers = {"Range": "bytes={}-{}".format(start, end)}

        self.lock.acquire()
        try:
            if self.total_bytes_downloaded == self.file_size:
                self.success = True
            else:
                self.failed = True
        finally:
            self.lock.release()

    def init_file(self):
        with open(self.dest, "wb") as f:
            f.truncate(self.file_size)

    def print_progress(self):
        while self.running:
            percent = round(self.total_bytes_downloaded /
                            self.file_size * 100, 2)
            progress = "#" * int(percent // 2)
            print(
                "Downloaded: {} / {} bytes [{}] {}%".format(
                    self.total_bytes_downloaded, self.file_size, progress.ljust(
                        50, "-"), percent
                )
            )
            time.sleep(0.5)

    def get_extension_by_content_type(self, content_type):
        """
        Get file extension by content type.
        """
        ext = mimetypes.guess_extension(content_type)
        if not ext:
            ext = ""
        return ext


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--url', required=True,
                        help='URL of the file to download')
    parser.add_argument('--dest', required=True,
                        help='Destination of the file to save')
    parser.add_argument('--overide', type=str, default="y",
                        help='Overide to file download')
    parser.add_argument('--chunk-size', type=int, default=1024 *
                        1024, help='Size of each download chunk')
    parser.add_argument('--num-threads', type=int, default=4,
                        help='Number of download threads')
    parser.add_argument('--max-retries', type=int, default=3,
                        help='Number of Max retries')
    parser.add_argument('--retry-after', type=int, default=5,
                        help='Number of Retry after seconds')
    args = parser.parse_args()

    try:
        downloader = Downloader(
            args.url, args.dest, args.overide, args.num_threads, args.chunk_size, args.max_retries, args.retry_after)
        downloader.start()
        sys.exit(0)
    except Exception as e:
        print("Error:", str(e))
        sys.exit(1)
