"""Downloaders are responsible for downloading evaluation repositories from various sources (e.g. IOM, UNHCR)."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_downloaders.ipynb.

# %% auto 0
__all__ = ['flatten_evals_id_doc', 'mk_dirs', 'extract_fname', 'download_doc', 'download_docs']

# %% ../nbs/02_downloaders.ipynb 2
from pathlib import Path
from rich import print
from fastcore.all import *
from urllib.parse import urlparse
import requests
from functools import partial

from .readers import load_evals, default_config

# %% ../nbs/02_downloaders.ipynb 3
def flatten_evals_id_doc(
    evals: L,  # List of evaluation records
    cfg: dict = default_config  # Config dict with field names
) -> L:  # Returns list of (id, url) tuples
    "Flatten evaluation records into list of (id, url) tuples"
    return L((eval_data[cfg['id_field']], doc[cfg['url_field']]) 
             for eval_data in evals 
             for doc in eval_data[cfg['docs_field']])

# %% ../nbs/02_downloaders.ipynb 5
def mk_dirs(
    docs_to_download: L, # list of (eval_id, url) tuples
    base_dir: Path # path to the base directory
    ) -> None:
    "Create directories for each unique evaluation ID"
    unique_eval_ids = set(eval_id for eval_id, url in docs_to_download)
    for eval_id in unique_eval_ids:
        (base_dir / eval_id).mkdir(parents=True, exist_ok=True)

# %% ../nbs/02_downloaders.ipynb 7
def extract_fname(url: str) -> str: 
    parsed_url = urlparse(url)
    fname = Path(parsed_url.path).name
    return fname

# %% ../nbs/02_downloaders.ipynb 8
def download_doc(
    doc_info: tuple[str, str],  # (eval_id, url) for the document to download
    base_dir: Path, # Base directory to save files to
    overwrite: bool = True # If True, overwrite existing file
) -> str:
    "Download a document from a given URL and save it to a specified directory."
    eval_id, url = doc_info    
    
    fname = extract_fname(url)
    path = base_dir / eval_id / fname
    
    if path.exists() and not overwrite: 
        return f"Skipped {fname} (already exists)"
    
    try:
        r = requests.get(url, timeout=15)
        r.raise_for_status()
        
        with open(path, 'wb') as f:
            f.write(r.content)
            
        return f"Downloaded {fname}"
    except Exception as e:
        return f"Failed to download {fname}: {e}"

# %% ../nbs/02_downloaders.ipynb 11
def download_docs(
    json_file: str, # path to the JSON file
    base_dir: Path = Path("./PDF_Library"), # path to the base directory
    n_workers: int = 5, # number of workers
    overwrite: bool = False, # if True, overwrite existing files
    cfg: dict = default_config # config for the JSON file
) -> list:
    "Download all evaluation documents from a JSON file in parallel"
    docs_to_download = flatten_evals_id_doc(load_evals(json_file))
    mk_dirs(docs_to_download, base_dir)
    download_func = partial(download_doc, base_dir=base_dir, overwrite=overwrite)
    results = parallel(download_func, 
                      docs_to_download, 
                      n_workers=n_workers,
                      total=len(docs_to_download),
                      progress=True)
    
    return results
