"""
PodWorker | modules | download.py

Called when inputs are images or zip files.
Downloads them into a temporary directory called "input_objects".
This directory is cleaned up after the job is complete.
"""

import os
import re
import uuid
import zipfile
from concurrent.futures import ThreadPoolExecutor
from email import message_from_string
from typing import List, Union
from urllib.parse import urlparse

import backoff
from requests import RequestException

from runpod.http_client import SyncClientSession

HEADERS = {"User-Agent": "runpod-python/0.0.0 (https://runpod.io; support@runpod.io)"}


def calculate_chunk_size(file_size: int) -> int:
    """
    Calculates the chunk size based on the file size.
    """
    if file_size <= 1024 * 1024:  # 1 MB
        return 1024  # 1 KB
    if file_size <= 1024 * 1024 * 1024:  # 1 GB
        return 1024 * 1024  # 1 MB

    return 1024 * 1024 * 10  # 10 MB


def download_files_from_urls(job_id: str, urls: Union[str, List[str]]) -> List[str]:
    """
    Accepts a single URL or a list of URLs and downloads the files.
    Returns the list of downloaded file absolute paths.
    Saves the files in a directory called "downloaded_files" in the job directory.
    """
    download_directory = os.path.abspath(
        os.path.join("jobs", job_id, "downloaded_files")
    )
    os.makedirs(download_directory, exist_ok=True)

    @backoff.on_exception(backoff.expo, RequestException, max_tries=3)
    def download_file(url: str, path_to_save: str) -> str:
        with SyncClientSession().get(
            url, headers=HEADERS, stream=True, timeout=5
        ) as response:
            response.raise_for_status()
            content_disposition = response.headers.get("Content-Disposition")
            file_extension = ""
            if content_disposition:
                msg = message_from_string(f"Content-Disposition: {content_disposition}")
                params = dict(msg.items())
                file_extension = os.path.splitext(params.get("filename", ""))[1]

            # If no extension could be determined from 'Content-Disposition', get it from the URL
            if not file_extension:
                file_extension = os.path.splitext(urlparse(url).path)[1]

            file_size = int(response.headers.get("Content-Length", 0))
            chunk_size = calculate_chunk_size(file_size)

            # write the content in chunks to the file
            with open(path_to_save + file_extension, "wb") as file_path:
                for chunk in response.iter_content(chunk_size=chunk_size):
                    if chunk:  # filter out keep-alive chunks
                        file_path.write(chunk)

            return file_extension

    def download_file_to_path(url: str) -> str:
        if url is None:
            return None

        file_name = f"{uuid.uuid4()}"
        output_file_path = os.path.join(download_directory, file_name)

        try:
            file_extension = download_file(url, output_file_path)
        except RequestException as err:
            print(f"Failed to download {url}: {err}")
            return None

        return os.path.abspath(f"{output_file_path}{file_extension}")

    if isinstance(urls, str):
        urls = [urls]

    with ThreadPoolExecutor() as executor:
        downloaded_files = list(executor.map(download_file_to_path, urls))

    return downloaded_files


def file(file_url: str) -> dict:
    """
    Downloads a single file from a given URL, file is given a random name.
    First checks if the content-disposition header is set, if so, uses the file name from there.
    If the file is a zip file, it is extracted into a directory with the same name.

    Returns an object that contains:
    - The absolute path to the downloaded file
    - File type
    - Original file name
    """
    os.makedirs("job_files", exist_ok=True)

    download_response = SyncClientSession().get(file_url, headers=HEADERS, timeout=30)

    original_file_name = []
    if "Content-Disposition" in download_response.headers.keys():
        original_file_name = re.findall(
            "filename=(.+)", download_response.headers["Content-Disposition"]
        )

    if len(original_file_name) > 0:
        original_file_name = original_file_name[0]
    else:
        download_path = urlparse(file_url).path
        original_file_name = os.path.basename(download_path)

    file_type = os.path.splitext(original_file_name)[1].replace(".", "")

    file_name = f"{uuid.uuid4()}"

    output_file_path = os.path.join("job_files", f"{file_name}.{file_type}")
    with open(output_file_path, "wb") as output_file:
        output_file.write(download_response.content)

    if file_type == "zip":
        unziped_directory = os.path.join("job_files", file_name)
        os.makedirs(unziped_directory, exist_ok=True)
        with zipfile.ZipFile(output_file_path, "r") as zip_ref:
            zip_ref.extractall(unziped_directory)
        unziped_directory = os.path.abspath(unziped_directory)
    else:
        unziped_directory = None

    return {
        "file_path": os.path.abspath(output_file_path),
        "type": file_type,
        "original_name": original_file_name,
        "extracted_path": unziped_directory,
    }
