import hashlib
import pathlib
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import as_completed

import requests


from .baai_config import Application, Progress
from .baai_downloder import check_file
from .core.baai_presign import read_presign

app = Application() # noqa
progress = Progress()

executor_spawn = ThreadPoolExecutor(max_workers=10)
executor_progress = ThreadPoolExecutor(max_workers=app.jobs)
executor_sign = ThreadPoolExecutor(max_workers=app.jobs)
executor_down = ThreadPoolExecutor(max_workers=app.jobs)
executor_watch = ThreadPoolExecutor(max_workers=50)
executor_merge = ThreadPoolExecutor(max_workers=app.jobs/2)


def download_sign_file(raw_data: str):
    from .presign import extract


    config = Application()

    meta_file: dict[str, str] = extract.jwtsign_read_extract(raw_data)
    hash_prefix = hashlib.md5(meta_file.get("path").encode("utf-8")).hexdigest()
    hash_ext = pathlib.Path(meta_file.get("path")).suffix

    sign_addr: str = read_presign(raw_data)
    resp_meta = requests.head(sign_addr, headers={"accept-encoding": None}) # noqa
    size = int(resp_meta.headers.get('Content-Length'))

    progress.add_require_size(size)
    # 文件存在不进行下载
    if check_file(f"{config.save_path}/{hash_prefix}{hash_ext}", size):
        progress.add_download_size(size)
        progress.add_download_count(1)
        return

    chunk_size = config.chunk_size
    total_parts = int((size + chunk_size - 1) / chunk_size)

    save_path = pathlib.Path(config.save_path)
    save_path.mkdir(parents=True, exist_ok=True)

    download_futures = []
    for part_index in range(total_parts):
        start = part_index * chunk_size
        end = min(start + chunk_size, size)
        task_down = executor_down.submit(download_sign_part_write, raw_data, start, end, hash_prefix)
        download_futures.append(task_down)

    # 等待完成
    for future in as_completed(download_futures):
        future.result()

    executor_merge.submit(download_sign_part_merge,  hash_prefix, hash_ext, total_parts)



def download_sign_part_write(raw_data, start: int, end: int, part_hash):
    config = Application()

    part_idx = int(start / config.chunk_size)
    part_name =  pathlib.Path(f"{config.save_path}/{part_hash}_{part_idx}.bin")

    if check_file(part_name.__str__(),  end-start):
        pass
    else:
        range_header = f"bytes={start}-{end - 1}"
        sign_addr = read_presign(raw_data)
        resp_file = requests.get(sign_addr, headers={"Range": range_header})

        # md5加密
        with part_name.open("wb") as fwriter:
            fwriter.write(resp_file.content)

    progress.add_download_size(end-start)


def download_sign_part_merge(name_prefix, name_ext, total_parts): # noqa

    config = Application()
    target_path = pathlib.Path(f"{config.save_path}/{name_prefix}{name_ext}")
    with target_path.open("ab") as fwriter:
        for part_idx in range(total_parts):
            part_name = pathlib.Path(f"{config.save_path}/{name_prefix}_{part_idx}.bin")
            with part_name.open("rb") as f:
                data = f.read()
                fwriter.write(data)

    for part_idx in range(total_parts):
        part_name = pathlib.Path(f"{config.save_path}/{name_prefix}_{part_idx}.bin")
        part_name.unlink()
    progress.add_download_count(1)


def download_progress():
    config = Application()
    with open(config.mate_path, "r") as fr:
        for _line in fr:
            try:
                executor_progress.submit(download_sign_file, _line)
            except Exception as e: # noqa
                progress.add_download_count(-1)


def download_sign_dataset(dataset_id, save_path, meta_path, jobs=8):
    import time
    from tqdm import tqdm # noqa

    from .baai_config import Application, Progress # noqa
    from .baai_logger import logger_download

    progress = Progress()
    logger_download()

    from .baai_downloder import executor_spawn, download_readrows, download_watch


    p1 = executor_spawn.submit(download_readrows)
    p2 = executor_spawn.submit(download_progress)
    p3 = executor_spawn.submit(download_watch)

    with tqdm(
        desc="speed",
        # colour="green",
        ncols=100,
        ascii="░▒▓█",
        unit="B",
        unit_scale=True,
        # bar_format="{desc}: {percentage:3.1f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}] {postfix}",
        bar_format="[{elapsed}] | {percentage:3.1f}%|{bar}| {n_fmt}/{total_fmt}{postfix} [{desc}: {rate_fmt}]"
    ) as pbar:
        last_size = 0
        while True:
            time.sleep(1)
            # 更新总任务数
            pbar.total = progress.require_size
            pbar.refresh()  # 刷新以显示新的total

            # 计算增量并更新进度条
            current_size = progress.download_size
            increment = min(current_size, progress.require_size) - last_size

            if increment > 0:
                pbar.update(increment)
                last_size = current_size

            pbar.set_postfix({
                "count": f"{progress.download_count}/{progress.require_count}",
            })

            if progress.require_size == progress.download_size and progress.require_count == progress.download_count:
                break

    p1.result()
    p2.result()
    p3.result()

    logger_download()