from pathlib import Path
import functools
import os
import re
import tempfile
from urllib.parse import unquote, urlparse

import fal
from fal.toolkit.utils.download_utils import (
    FAL_MODEL_WEIGHTS_DIR,
    DownloadError,
    _hash_url,
)

FAL_VERSION = getattr(fal, "__version__", "<1.0.0")
_REQUEST_HEADERS = {"User-Agent": f"fal-client ({FAL_VERSION}/python)"}

TEMP_FILE_SUFFIX = ".tmp"


@functools.cache
def get_boto3_client():
    import boto3

    # boto reads AWS credentials from environment variables
    # which should be set using `fal_secrets`
    s3_client = boto3.client("s3")
    return s3_client


@functools.cache
def get_civitai_headers() -> dict[str, str]:
    headers: dict[str, str] = {}

    civitai_token = os.getenv("CIVITAI_TOKEN", None)

    if not civitai_token:
        print("CIVITAI_TOKEN is not set in the environment variables.")
        return headers

    headers["Authorization"] = f"Bearer {civitai_token}"

    return headers


@functools.cache
def get_huggingface_headers() -> dict[str, str]:
    headers: dict[str, str] = {}

    hf_token = os.getenv("HF_TOKEN", None)

    if not hf_token:
        print("HF_TOKEN is not set in the environment variables.")
        return headers

    headers["Authorization"] = f"Bearer {hf_token}"

    return headers


def get_local_file_content_length(file_path: Path) -> int:
    return file_path.stat().st_size


def download_url_to_file(
    url: str,
    dst: str | Path,
    progress: bool = True,
    headers: dict[str, str] = None,
    chunk_size_in_mb=16,
) -> Path:
    """Download object at the given URL to a local path.

    Args:
        url (str): URL of the object to download
        dst (str): Full path where object will be saved, e.g. ``/tmp/temporary_file``
        progress (bool, optional): whether or not to display a progress bar to stderr
            Default: True
        headers (dict, optional): HTTP headers to include with the request
            Default: None

    """
    from tqdm import tqdm

    file_size = None

    request_headers = {
        **_REQUEST_HEADERS,
        **(headers or {}),
    }

    import requests

    response = requests.get(
        url, headers=request_headers, stream=True, allow_redirects=True
    )
    response.raise_for_status()

    content_length = response.headers.get("Content-Length", None)  # type: ignore
    if content_length is not None and len(content_length) > 0:
        file_size = int(content_length[0])

    basename = os.path.basename(dst)

    with tempfile.NamedTemporaryFile(
        delete=False,
        dir=os.path.dirname(dst),
        prefix=f"{basename}",
        suffix=TEMP_FILE_SUFFIX,
    ) as temp_file:
        file_path = temp_file.name
        try:
            with tqdm(
                total=file_size,
                disable=not progress,
                unit="B",
                unit_scale=True,
                unit_divisor=1024,
            ) as pbar, open(file_path, "wb") as f:
                for chunk in response.iter_content(
                    chunk_size=chunk_size_in_mb * 1024 * 1024
                ):
                    if chunk:
                        f.write(chunk)
                        pbar.update(len(chunk))

            # NOTE: Atomically renaming the file into place when the file is downloaded
            # completely.
            #
            # Since the file used is temporary, in a case of an interruption, the
            # downloaded content will be lost. So, it is safe to redownload the file in
            # such cases.
            os.rename(file_path, dst)
        except Exception as error:
            raise error
        finally:
            Path(temp_file.name).unlink(missing_ok=True)

    return Path(dst)


def download_s3_to_file(
    s3_uri: str,
    dst: str | Path,
    progress: bool = True,
    chunk_size_in_mb=16,
) -> Path:
    """Download object from S3 to a local path using boto3.

    Args:
        s3_url (str): S3 URL of the object to download (s3://bucket/key format)
        dst (str): Full path where object will be saved, e.g. ``/tmp/temporary_file``
        progress (bool, optional): whether or not to display a progress bar to stderr
            Default: True
    """
    from boto3.s3.transfer import TransferConfig, MB
    from botocore.exceptions import ClientError, NoCredentialsError
    from tqdm import tqdm

    # Parse S3 URL
    parsed_url = urlparse(s3_uri)

    bucket_name = parsed_url.netloc
    object_key = parsed_url.path.lstrip("/")

    try:
        s3_client = get_boto3_client()

        # Get object metadata to determine file size
        try:
            response = s3_client.head_object(Bucket=bucket_name, Key=object_key)
            file_size = response["ContentLength"]
        except ClientError as e:
            if e.response["Error"]["Code"] == "404":
                raise FileNotFoundError(f"S3 object not found: {s3_uri}")
            raise
    except NoCredentialsError as exc:
        raise ValueError(
            "AWS credentials not found. Please provide credentials using `fal secrets set` for related AWS keys. See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html#environment-variables"
        ) from exc

    basename = os.path.basename(dst)

    with tempfile.NamedTemporaryFile(
        delete=False,
        dir=os.path.dirname(dst),
        prefix=f"{basename}",
        suffix=TEMP_FILE_SUFFIX,
    ) as temp_file:
        file_path = temp_file.name
        try:
            with tqdm(
                total=file_size,
                disable=not progress,
                unit="B",
                unit_scale=True,
                unit_divisor=1024,
                desc=f"Downloading {os.path.basename(object_key)}",
            ) as pbar:

                def progress_callback(bytes_transferred):
                    pbar.update(bytes_transferred)

                # Download the file with progress tracking
                s3_client.download_file(
                    bucket_name,
                    object_key,
                    file_path,
                    Callback=progress_callback,
                    Config=TransferConfig(
                        multipart_chunksize=chunk_size_in_mb * MB,
                        multipart_threshold=chunk_size_in_mb * MB,
                    ),
                )

            # NOTE: Atomically renaming the file into place when the file is downloaded
            # completely.
            #
            # Since the file used is temporary, in a case of an interruption, the
            # downloaded content will be lost. So, it is safe to redownload the file in
            # such cases.
            os.rename(file_path, dst)
        except Exception as error:
            raise error
        finally:
            Path(temp_file.name).unlink(missing_ok=True)

    return Path(dst)


def download_model_weights(url: str, force: bool = False) -> Path:
    parsed_url = urlparse(url)

    headers = {}
    download_url = url

    if parsed_url.netloc == "civitai.com":
        headers.update(get_civitai_headers())
        # Also add token to URL query params for better compatibility
        civitai_token = os.getenv("CIVITAI_TOKEN", None)
        if civitai_token and "token=" not in url:
            separator = "&" if "?" in url else "?"
            download_url = f"{url}{separator}token={civitai_token}"
    elif parsed_url.netloc == "huggingface.co":
        headers.update(get_huggingface_headers())

    if parsed_url.scheme == "s3":
        return download_model_weights_fal_s3(url, force=force)
    else:
        return download_model_weights_fal(
            download_url, request_headers=headers, force=force
        )


def download_model_weights_fal_s3(url: str, force: bool = False) -> Path:
    weights_dir = Path(FAL_MODEL_WEIGHTS_DIR / _hash_url(url))

    if weights_dir.exists() and not force:
        try:
            weights_path = next(
                wp for wp in weights_dir.glob("*") if wp.suffix != TEMP_FILE_SUFFIX
            )
            return weights_path

        # The model weights directory is empty, so we need to download the weights
        except StopIteration:
            pass

    try:
        parsed_url = urlparse(url)

        s3 = get_boto3_client()
        bucket_name = parsed_url.netloc
        object_path = parsed_url.path.lstrip("/")

        object_info = s3.head_object(Bucket=bucket_name, Key=object_path)
        file_name = _truncate_filename(object_path.split("/")[-1])
        file_content_length = object_info["ContentLength"]
    except Exception as e:
        print(e)
        raise DownloadError(f"Failed to get remote file properties for {url}")

    target_path = weights_dir / file_name

    if (
        target_path.exists()
        and get_local_file_content_length(target_path) == file_content_length
        and not force
    ):
        return target_path

    # Make sure the parent directory exists
    target_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        download_s3_to_file(url, target_path, progress=True)
    except Exception:
        raise DownloadError(f"Failed to download {url}")

    return target_path


def download_model_weights_fal(
    url: str, force: bool = False, request_headers: dict[str, str] | None = None
) -> Path:
    weights_dir = Path(FAL_MODEL_WEIGHTS_DIR / _hash_url(url))

    if weights_dir.exists() and not force:
        try:
            weights_path = next(
                wp for wp in weights_dir.glob("*") if wp.suffix != TEMP_FILE_SUFFIX
            )
            return weights_path

        # The model weights directory is empty, so we need to download the weights
        except StopIteration:
            pass

    try:
        file_name, file_content_length = _get_remote_file_properties(
            url, request_headers=request_headers
        )
    except Exception as e:
        print(e)
        raise DownloadError(f"Failed to get remote file properties for {url}")

    target_path = weights_dir / file_name

    if (
        target_path.exists()
        and get_local_file_content_length(target_path) == file_content_length
        and not force
    ):
        return target_path

    # Make sure the parent directory exists
    target_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        download_url_to_file(url, target_path, progress=True, headers=request_headers)
    except Exception as e:
        print(e)
        raise DownloadError(f"Failed to download {url}")

    return target_path


def _get_filename_from_content_disposition(cd: str | None) -> str | None:
    if not cd:
        return None

    filenames = re.findall('filename="(.+)"', cd)

    if len(filenames) == 0:
        filenames = re.findall("filename=(.+)", cd)

    if len(filenames) == 0:
        return None

    return unquote(filenames[0])


def _truncate_filename(file_name: str, max_length: int = 200) -> str:
    """Truncate a filename to fit within filesystem limits.

    Preserves the file extension and replaces the excess stem with a hash
    to avoid collisions. The default max_length of 200 leaves room for
    temp file prefixes/suffixes added by NamedTemporaryFile.
    """
    if len(file_name) <= max_length:
        return file_name

    name, ext = os.path.splitext(file_name)
    hash_suffix = _hash_url(file_name)[:16]
    max_name_length = max_length - len(ext) - len(hash_suffix) - 1  # 1 for separator

    if max_name_length <= 0:
        # Extension alone consumes the budget; trim the hash instead
        return hash_suffix[: max_length - len(ext)] + ext

    truncated_name = name[:max_name_length] + "_" + hash_suffix
    return truncated_name + ext


def _parse_filename(url: str, cd: str | None) -> str:
    file_name = _get_filename_from_content_disposition(cd)

    if not file_name:
        parsed_url = urlparse(url)

        if parsed_url.scheme == "data":
            file_name = _hash_url(url)
        else:
            url_path = parsed_url.path
            file_name = Path(url_path).name or _hash_url(url)

    return _truncate_filename(file_name)


def _get_remote_file_properties(
    url: str, request_headers: dict[str, str] = None
) -> tuple[str, int]:
    import requests

    headers = {
        **_REQUEST_HEADERS,
        **(request_headers or {}),
    }

    req = requests.get(
        url, headers=headers, stream=True, allow_redirects=True, verify=False
    )
    req.raise_for_status()

    headers = req.headers  # type: ignore
    content_disposition = headers.get("Content-Disposition", None)
    file_name = _parse_filename(url, content_disposition)
    content_length = int(headers.get("Content-Length", -1))

    return file_name, content_length


def download_model_weights_fal_fast(
    url: str,
    request_headers: dict[str, str] | None = None,
    force: bool = False,
) -> Path:
    """Download model weights using hf_transfer with specified parameters.

    Args:
        url: The URL to download the model weights from.
        request_headers: Headers to use for the download request. Defaults to None.
        force: Whether to force the download if the model weights already exist. Defaults to False.

    Returns:
        Path: The path to the downloaded model weights.
    """
    from hf_transfer import download

    weights_dir = Path(FAL_MODEL_WEIGHTS_DIR / _hash_url(url))

    if weights_dir.exists() and not force:
        try:
            weights_path = next(weights_dir.glob("*"))
            return weights_path

        # The model weights directory is empty, so we need to download the weights
        except StopIteration:
            pass

    try:
        file_name, file_content_length = _get_remote_file_properties(
            url, request_headers=request_headers
        )
    except Exception as e:
        print(e)
        raise DownloadError(f"Failed to get remote file properties for {url}")

    target_path = weights_dir / file_name

    if (
        target_path.exists()
        and get_local_file_content_length(target_path) == file_content_length
        and not force
    ):
        return target_path

    def callback(*args):
        print("Args", args)

    download(
        url,
        filename=str(target_path),
        max_files=100,
        chunk_size=10 * 1024 * 1024,
        max_retries=5,
        parallel_failures=3,
        headers=request_headers,
        callback=callback,
    )

    return target_path
