from typing import IO, List, Optional, Tuple, Union

import json
import os
import re
import shutil
import sys
import tempfile
import textwrap
import time
import warnings
from urllib.parse import parse_qs, unquote, urlparse

import requests

from tqdm import tqdm

CHUNK_SIZE = 512 * 1024
HOME = os.path.expanduser("~")


def indent(text: str, prefix: str):
    """Indent each non-empty line of text with the given prefix."""
    return "".join(
        (prefix + line if line.strip() else line) for line in text.splitlines(True)
    )


class FileURLRetrievalError(Exception):
    """Custom exception for issues retrieving file URLs."""


def _extract_download_url_from_confirmation(contents: str, url_origin: str):
    """Extract the download URL from a Google Drive confirmation page."""
    patterns = [
        r'href="(\/uc\?export=download[^"]+)',
        r'href="/open\?id=([^"]+)"',
        r'"downloadUrl":"([^"]+)',
    ]
    for pattern in patterns:
        match = re.search(pattern, contents)
        if match:
            url = match.group(1)
            if pattern == r'href="/open\?id=([^"]+)"':
                uuid_match = re.search(
                    r'<input\s+type="hidden"\s+name="uuid"\s+value="([^"]+)"',
                    contents,
                )
                if uuid_match:
                    uuid = uuid_match.group(1)
                    return (
                        "https://drive.usercontent.google.com/download?id="
                        + url
                        + "&confirm=t&uuid="
                        + uuid
                    )
                raise FileURLRetrievalError(
                    f"Could not find UUID for download from {url_origin}",
                )
            if pattern == r'"downloadUrl":"([^"]+)':
                return url.replace("\\u003d", "=").replace("\\u0026", "&")
            return "https://docs.google.com" + url.replace("&", "&")

    error_match = re.search(r'<p class="uc-error-subcaption">(.*)</p>', contents)
    if error_match:
        error = error_match.group(1)
        raise FileURLRetrievalError(error)

    raise FileURLRetrievalError(
        "Cannot retrieve the public link of the file. "
        "You may need to change the permission to "
        "'Anyone with the link', or have had many accesses.",
    )


def _create_session(
    proxy: str | None = None,
    use_cookies: bool = True,
    return_cookies_file: bool = False,
):
    """Create a requests session with optional proxy and cookie handling."""
    sess = requests.session()
    sess.headers.update(
        {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6)"},
    )

    if proxy:
        sess.proxies = {"http": proxy, "https": proxy}

    cookies_file = os.path.join(HOME, ".cache/gdown/cookies.json")
    if os.path.exists(cookies_file) and use_cookies:
        try:
            with open(cookies_file) as f:
                cookies = json.load(f)
            for k, v in cookies:
                sess.cookies[k] = v
        except json.JSONDecodeError:
            warnings.warn("Corrupted Cookies file")

    return (sess, cookies_file) if return_cookies_file else sess


def download(
    output: str | None = None,
    quiet: bool = False,
    proxy: str | None = None,
    speed: float | None = None,
    use_cookies: bool = True,
    verify: bool | str = True,
    id: str | None = None,
    fuzzy: bool = True,
    resume: bool = False,
    format: str | None = None,
    url: str | None = None,
):
    """
    Download a file from a URL, supporting Google Drive links.

    Args:
        output: Output filepath. Default is basename of URL.
        quiet: Suppress terminal output.
        proxy: HTTP/HTTPS proxy.
        speed: Download speed limit (bytes per second).
        use_cookies: Flag to use cookies.
        verify: Verify TLS certificates.
        id: Google Drive's file ID.
        fuzzy: Fuzzy Google Drive ID extraction.
        resume: Resume download from a tmp file.
        format: Format for Google Docs/Sheets/Slides.
        url: URL to download from.

    Returns:
        Output filename, or None on error.

    """
    if not (id is None) ^ (url is None):
        raise ValueError("Either url or id has to be specified")

    if id is not None:
        url = f"https://drive.google.com/uc?id={id}"

    url_origin = url
    sess, cookies_file = _create_session(
        proxy=proxy,
        use_cookies=use_cookies,
        return_cookies_file=True,
    )

    while True:
        res = sess.get(url, stream=True, verify=verify)
        res.raise_for_status()

        if url == url_origin and res.status_code == 500:
            url = f"https://drive.google.com/open?id={id}"
            continue

        if res.headers.get("Content-Type", "").startswith("text/html"):
            title_match = re.search(r"<title>(.+)</title>", res.text)
            if title_match:
                title = title_match.group(1)
                if title.endswith(" - Google Docs"):
                    url = (
                        f"https://docs.google.com/document/d/{id}/export?format={'docx' if format is None else format}"
                    )
                    continue
                if title.endswith(" - Google Sheets"):
                    url = (
                        f"https://docs.google.com/spreadsheets/d/{id}/export?format={'xlsx' if format is None else format}"
                    )
                    continue
                if title.endswith(" - Google Slides"):
                    url = (
                        f"https://docs.google.com/presentation/d/{id}/export?format={'pptx' if format is None else format}"
                    )
                    continue
            if (
                "Content-Disposition" in res.headers
                and res.headers["Content-Disposition"].endswith("pptx")
                and format not in (None, "pptx")
            ):
                url = (
                    f"https://docs.google.com/presentation/d/{id}/export?format={'pptx' if format is None else format}"
                )
                continue

        if use_cookies:
            os.makedirs(os.path.dirname(cookies_file), exist_ok=True)
            cookies = [
                (k, v)
                for k, v in sess.cookies.items()
                if not k.startswith("download_warning_")
            ]
            with open(cookies_file, "w") as f:
                json.dump(cookies, f, indent=2)

        if "Content-Disposition" in res.headers:
            break

        parsed_url = urlparse(url)
        is_gdrive = parsed_url.hostname in ("drive.google.com", "docs.google.com")
        is_download_link = parsed_url.path.endswith("/uc")

        if not (is_gdrive and is_download_link and fuzzy):
            break

        try:
            url = _extract_download_url_from_confirmation(res.text, url_origin)
        except FileURLRetrievalError as e:
            raise FileURLRetrievalError(e)

    content_disposition = res.headers.get("Content-Disposition", "")
    filename_match = re.search(
        r"filename\*=UTF-8''(.*)",
        content_disposition,
    ) or re.search(r'filename=["\']?(.*?)["\']?$', content_disposition)
    filename_from_url = (
        unquote(filename_match.group(1)) if filename_match else os.path.basename(url)
    )
    download_path = output or filename_from_url

    if isinstance(download_path, str) and download_path.endswith(os.path.sep):
        os.makedirs(download_path, exist_ok=True)
        download_path = os.path.join(download_path, filename_from_url)

    temp_dir = os.path.dirname(download_path) or "."
    prefix = os.path.basename(download_path)

    if isinstance(download_path, str):
        existing_tmp_files = [
            os.path.join(temp_dir, file)
            for file in os.listdir(temp_dir)
            if file.startswith(prefix)
        ]
        if resume and existing_tmp_files:
            if len(existing_tmp_files) > 1:
                print(
                    "There are multiple temporary files to resume:",
                    file=sys.stderr,
                )
                for file in existing_tmp_files:
                    print(f"\t{file}", file=sys.stderr)
                print(
                    "Please remove them except one to resume downloading.",
                    file=sys.stderr,
                )
                return None
            temp_file_path = existing_tmp_files[0]
        else:
            resume = False
            temp_file_path = tempfile.mktemp(
                suffix=tempfile.template,
                prefix=prefix,
                dir=temp_dir,
            )

        try:
            file_obj: IO = open(temp_file_path, "ab")
        except Exception as e:
            print(
                f"Could not open the temporary file {temp_file_path}: {e}",
                file=sys.stderr,
            )
            return None
    else:
        temp_file_path = None
        file_obj = download_path

    if temp_file_path is not None and file_obj.tell() != 0:
        headers = {"Range": f"bytes={file_obj.tell()}-"}
        res = sess.get(url, headers=headers, stream=True, verify=verify)
        res.raise_for_status()

    try:
        total = int(res.headers.get("Content-Length", 0))
        if total > 0:
            if not quiet:
                pbar = tqdm(
                    total=total,
                    unit="B",
                    unit_scale=True,
                    desc=filename_from_url,
                )
        elif not quiet:
            pbar = tqdm(unit="B", unit_scale=True, desc=filename_from_url)

        t_start = time.time()
        for chunk in res.iter_content(chunk_size=CHUNK_SIZE):
            file_obj.write(chunk)
            if not quiet:
                pbar.update(len(chunk))
            if speed is not None:
                elapsed_time_expected = 1.0 * pbar.n / speed
                elapsed_time = time.time() - t_start
                if elapsed_time < elapsed_time_expected:
                    time.sleep(elapsed_time_expected - elapsed_time)
        if not quiet:
            pbar.close()

        if temp_file_path:
            file_obj.close()
            shutil.move(temp_file_path, download_path)
    finally:
        sess.close()

    return download_path
