import os
import concurrent.futures
import pandas as pd, requests
from dataclasses import dataclass
from tqdm import tqdm
from urllib3 import Retry
from requests.adapters import HTTPAdapter
from src.datasets.images.utils import verify
from src.datasets.wit.utils import get_image_name
from src.logger import Optional, get_logger

logger = get_logger(__name__)


def create_session(user_agent: Optional[str] = None) -> requests.Session:
    session = requests.Session()
    retry = Retry(
        total=3,
        read=3,
        connect=3,
        backoff_factor=1.5,
        status_forcelist=(500, 502, 504),
        allowed_methods=["HEAD", "GET", "OPTIONS"],
    )
    adapter = HTTPAdapter(max_retries=retry)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    session.headers.update({"User-Agent": user_agent} if user_agent else {})
    return session


@dataclass
class CrawlingArgs:
    images_dir: str
    remove_invalid: bool = False
    user_agent: str = "TODO: replace with your user agent"
    num_threads: int = 8


def download(session: requests.Session, url: str, args: CrawlingArgs):
    try:
        path = os.path.join(args.images_dir, get_image_name(url))
        os.makedirs(os.path.dirname(path), exist_ok=True)
        r = session.get(url, timeout=30)
        r.raise_for_status()

        if args.remove_invalid and not verify(r.content):
            logger.warning(f"Invalid image at {url}")
            return

        with open(path, "wb") as f:
            f.write(r.content)
    except Exception as e:
        logger.error(f"Failed to download {url}: {e}")


def crawl_images(wit_df: pd.DataFrame, args: CrawlingArgs) -> None:
    os.makedirs(args.images_dir, exist_ok=True)
    logger.info(f"Starting to crawl images to {args.images_dir}")

    unique_urls = wit_df["image_url"].unique().tolist()
    logger.info(f"Found {len(unique_urls)} unique image URLs")

    session = create_session(args.user_agent)
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.num_threads) as ex:
        futures = []
        for url in tqdm(unique_urls, desc="Scheduling downloads"):
            futures.append(ex.submit(download, session, url, args))

        for _ in tqdm(
            concurrent.futures.as_completed(futures),
            total=len(futures),
            desc="Downloading images",
        ):
            pass
    session.close()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_file", type=str, required=True, help="Path to the input WIT json file"
    )
    parser.add_argument(
        "--images_dir", type=str, required=True, help="Directory to save crawled images"
    )
    parser.add_argument(
        "--workers", type=int, default=8, help="Number of parallel download threads"
    )
    parser.add_argument(
        "--user_agent",
        type=str,
        default="TODO: replace with your user agent",
        help="User-Agent header for requests",
    )
    parser.add_argument(
        "--remove_invalid", action="store_true", help="Remove entries with invalid image URLs"
    )
    args = parser.parse_args()

    wit_df = pd.read_json(args.input_file, lines=args.input_file.endswith(".jsonl"))

    crawling_args = CrawlingArgs(
        images_dir=args.images_dir,
        remove_invalid=args.remove_invalid,
        user_agent=args.user_agent,
        num_threads=args.workers,
    )
    crawl_images(wit_df, crawling_args)
