from urllib.parse import urlparse
import os
from tqdm import tqdm
import asyncio
from .args import parse_args
from .constant import base_headers
from .cdn import load_cdn_info, update_cdn_info
from .custom_transport import request
from .utils import generate_file_name, generate_urls_of_all_qualities, save_file


def save(result, output_path, filename):
    u = urlparse(result.url)
    filename = generate_file_name(
        u.path, result.url, result.headers.get("content-type"), filename
    )
    path = os.path.join(output_path, filename)
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    save_file(path, result.body)
    print(f"Saved {result.url} to {path}")


async def rescue_async(
    urls, ips, output_path, filename=None, workers=20, timeout=10.0
):
    total_tasks = len(urls) * (len(ips) + 1)
    result = None
    semaphore = asyncio.Semaphore(max(1, workers))

    with tqdm(total=total_tasks, desc="Rescuing") as pbar:
        for url in urls:
            stop_event = asyncio.Event()
            ip_candidates = [None] + list(ips)

            async def run_task(ip):
                if stop_event.is_set():
                    return None
                async with semaphore:
                    if stop_event.is_set():
                        return None
                    task_result = await request(url, base_headers, ip, timeout=timeout)
                    pbar.update(1)
                    if task_result and task_result.status == 200:
                        stop_event.set()
                        return task_result
                    return None

            pending_tasks = [asyncio.create_task(run_task(ip)) for ip in ip_candidates]
            try:
                for completed in asyncio.as_completed(pending_tasks):
                    task_result = await completed
                    if task_result and task_result.status == 200:
                        result = task_result
                        break
            finally:
                for task in pending_tasks:
                    if not task.done():
                        task.cancel()
                await asyncio.gather(*pending_tasks, return_exceptions=True)

            if result is not None:
                break

        pbar.n = pbar.total
        pbar.refresh()

    if not result:
        print(f"[FAILED] Unfortunately, all {total_tasks} resolves failed.")
        return None

    print(f"[SUCCESS] {result.url} | {result.ip} | {len(result.body)}")
    save(result, output_path, filename)
    return result


def rescue(urls, ips, output_path, filename=None, workers=20, timeout=10.0):
    return asyncio.run(
        rescue_async(
            urls,
            ips,
            output_path,
            filename=filename,
            workers=workers,
            timeout=timeout,
        )
    )


def main(args=None):
    if args is None:
        args = parse_args(None)
    else:
        args = parse_args(args)

    if args.update_cdn:
        update_cdn_info()
        return

    if args.f:
        update_cdn_info()
        image_url = args.f
    else:
        image_url = args.image_url
        if args.force_update:
            update_cdn_info()

    if not image_url:
        print("Image URL is required.")
        return

    ips = load_cdn_info()

    if not ips:
        print("No cached resolves found. Continue with DNS fallback only.")

    print(f"Using {len(ips)} cached resolves.")
    qualities = [item.strip() for item in args.q.split(",") if item.strip()]
    urls = generate_urls_of_all_qualities(image_url, qualities)
    rescue(
        urls,
        ips,
        args.output_dir,
        workers=args.workers,
        timeout=args.timeout,
    )


if __name__ == "__main__":
    main()
