"""Functions for downloading (the source code of) articles from arXiv"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/49_helper.arxiv.ipynb.

# %% auto 0
__all__ = ['arxiv_id', 'arxiv_search', 'extract_metadata', 'ArxivMetadataEncoder', 'extract_last_names', 'folder_name_for_source',
           'create_acronym', 'file_name_for_pdf', 'analyze_arxiv_tarfile', 'read_gz_file', 'get_tex_filename_from_gz',
           'extract_tex_from_gz', 'download_from_results']

# %% ../../nbs/49_helper.arxiv.ipynb 1
import datetime
import gzip
import json
from typing import Callable, Literal, Optional, Union
import os
from os import PathLike
from pathlib import Path
import re
import tarfile

import arxiv
from arxiv import Client, Search, Result
from pathvalidate import sanitize_filename

from .files_and_folders import file_is_compressed, uncompress_file


# %% ../../nbs/49_helper.arxiv.ipynb 5
def arxiv_id(arxiv_id_or_url: str) -> str:
    """
    Return the arxiv id from a str which is either of the arxiv id itself or the url
    to the arxiv article.

    **Raises**
    - `ValueError`
        - If the input does not contain a valid arXiv ID.
    """
    new_id_pattern = r'(\d{4}\.\d{4,5}(?:v\d+)?)'
    old_id_pattern = r'([a-z-]+(?:\.[A-Z]{2})?/\d{7}(?:v\d+)?)'
    combined_pattern = f'({new_id_pattern}|{old_id_pattern})'
    
    # Check if input is a URL and extract the ID
    if 'arxiv.org' in arxiv_id_or_url:
        match = re.search(combined_pattern, arxiv_id_or_url)
        if match:
            return match.group(1)
        else:
            raise ValueError("Invalid arXiv URL provided.")
    
    # If it's not a URL, assume it's an ID and validate it
    elif re.match(combined_pattern, arxiv_id_or_url):
        return arxiv_id_or_url
    
    else:
        raise ValueError("Invalid input. Please provide a valid arXiv ID or URL.")


# %% ../../nbs/49_helper.arxiv.ipynb 8
def arxiv_search(
        arxiv_ids: Union[str, list[str]], # The ID of a single arXiv article or multiple arxiv articles
        client: Optional[Client] = None,  # an arxiv API Client. If `None`, create one on the spot.
        results: bool = True, # If `True` return a `Result` object. otherwise, return a `Search`` object`.
        ) -> Union[Result, Search]:
    if not client:
        client = Client()
    if not isinstance(arxiv_ids, list):
        arxiv_ids = [arxiv_ids]
    search = Search(id_list=arxiv_ids)
    if results:
        return client.results(search)
    return search

# %% ../../nbs/49_helper.arxiv.ipynb 14
def extract_metadata(
        results: Union[list[Result], Result],
        ) -> list[dict]: # Each dict corresponds to the metadata for each result.
    """
    Return the metadata from the arxiv search results
    """
    if not isinstance(results, list):
        results = [results]
    metadata_list = []
    for result in results:
        metadata = {
            "arxiv_id": result.get_short_id(),
            "authors": [author.name for author in result.authors],
            "title": result.title,
            "summary": result.summary,
            "primary_category": result.primary_category,
            "categories": result.categories,
            "published": result.published,
            "updated": result.updated,
            "doi": result.doi,
            "comment": result.comment,
            "journal_ref": result.journal_ref,
            "links": result.links
        }
        metadata_list.append(metadata)
    return metadata_list

# %% ../../nbs/49_helper.arxiv.ipynb 16
class ArxivMetadataEncoder(json.JSONEncoder):
    """
    `json` encoder to accomapny the `extract_metadta` function when using `json.dump`. 
    """
    def default(self, obj):
        if isinstance(obj, datetime.datetime):
            return obj.isoformat()
        elif isinstance(obj, arxiv.Result.Link):
            return obj.href
        return super().default(obj)

# %% ../../nbs/49_helper.arxiv.ipynb 19
def extract_last_names(
        authors: list[str]
        ) -> list[str]:
    last_names = []
    for author in authors:
        # Split the name into parts
        parts = author.split()
        # Handle special cases like "de Jong"
        if len(parts) > 2 and parts[-2].lower() in ['de', 'van', 'von', 'del', 'della', 'di', 'da', 'dos']:
            last_name = f"{parts[-2]} {parts[-1]}"
        else:
            last_name = parts[-1]
        
        # Remove any commas or periods
        last_name = re.sub(r'[,.]', '', last_name)
        last_names.append(last_name)
    return last_names


# %% ../../nbs/49_helper.arxiv.ipynb 22
def folder_name_for_source(
        result: Result,
        lowercase: bool = True
        ) -> str:
    family_names = extract_last_names([author.name for author in result.authors])
    if len(family_names) > 4:
        family_names_text = f'{family_names[0]}_et_al'
    else:
        family_names_text = '_'.join(family_names)
    if lowercase:
        output = f'{family_names_text.lower()}_{create_acronym(result.title)}'
    else:
        output = f'{family_names_text}_{create_acronym(result.title)}'
    return output


def create_acronym(title):
    # Words to exclude from acronym
    exclude_words = set(['a', 'an', 'the', 'on', 'and', 'of', 'to', 'over', 'in', 'for', 'with', 'by', 'at', 'from'])
    
    # Split the title into words
    words = re.findall(r'\b[\w-]+\b', title)
    
    acronym = ''
    for word in words:
        if word.lower() not in exclude_words:
            if len(word) == 1 and word.isupper():
                # Keep single uppercase letters (likely mathematical symbols) as is
                acronym += word
            elif word.lower() in ['i', 'ii', 'iii', 'iv', 'v', 'vi', 'vii', 'viii', 'ix', 'x']:
                # Handle Roman numerals
                acronym += word.lower()
            elif '-' in word:
                # Handle hyphenated words
                parts = word.split('-')
                acronym += ''.join(part[0].lower() if not (len(part) == 1 and part.isupper()) else part[0] for part in parts)
            else:
                # Take the first letter of other words, always lowercase
                acronym += word[0].lower()
    
    return acronym

# %% ../../nbs/49_helper.arxiv.ipynb 26
def file_name_for_pdf(
        result: Result
        ) -> str:
    family_names = extract_last_names([author.name for author in result.authors])
    if len(family_names) > 4:
        family_names_text = f'{family_names[0]} et al'
    else:
        family_names_text = ', '.join(family_names)
    output = f'{family_names_text} - {result.title}'
    return sanitize_filename(output)


# %% ../../nbs/49_helper.arxiv.ipynb 29
def analyze_arxiv_tarfile(
        filepath: PathLike # The path to the tar file.
        ) -> Literal["nested_archive", "direct_tex", "unknown_tar_structure", "plain_gz", "invalid_file"]:
    """
    Analyzes the contents of an arXiv download file, which can be either a tar.gz
    archive or a plain .gz file.

    This function attempts to determine the structure of the file downloaded from arXiv.
    It can identify several different types of content structures commonly found in arXiv
    downloads.

    Parameters:
    filepath (Union[str, Path]): The path to the file to be analyzed. Can be a string or a Path object.

    Returns:
    Literal["nested_archive", "direct_tex", "unknown_tar_structure", "plain_gz", "invalid_file"]: 
        - "nested_archive": If the tar.gz contains another compressed file
        - "direct_tex": If the tar.gz contains .tex files directly
        - "unknown_tar_structure": If the tar.gz structure doesn't match known patterns
        - "plain_gz": If the file is a plain .gz file (not a tar.gz)
        - "invalid_file": If the file is neither a valid tar.gz nor a valid .gz file

    Raises:
    No exceptions are raised; all errors are handled internally and returned as "invalid_file".
    Determine what kind of contents the 
    """
    try:
        with tarfile.open(filepath, "r:gz") as tar:
            members = tar.getmembers()
            
            if len(members) == 1 and members[0].name.endswith('.tar.gz'):
                return "nested_archive"
            
            tex_files = [m for m in members if m.name.endswith('.tex')]
            if tex_files:
                return "direct_tex"
            
            return "unknown_tar_structure"
    except tarfile.ReadError:
        # Check if it's a plain .gz file
        try:
            with gzip.open(filepath, 'rb') as gz_file:
                # Read a small portion to check if it's a valid gzip file
                gz_file.read(1024)
            return "plain_gz"
        except gzip.BadGzipFile:
            return "invalid_file"


def read_gz_file(filepath):
    with gzip.open(filepath, 'rt') as f:
        content = f.read()
    return content


def get_tex_filename_from_gz(filepath):
    with gzip.open(filepath, 'rt', encoding='utf-8') as f:
        content = f.read()
        
    # Look for a .tex filename in the content
    match = re.search(r'\b[\w-]+\.tex\b', content)
    if match:
        return match.group(0)
    else:
        return None


def extract_tex_from_gz(filepath):
    base_name = os.path.splitext(os.path.basename(filepath))[0]
    tex_filename = f"{base_name}.tex"
    
    with gzip.open(filepath, 'rb') as gz_file:
        content = gz_file.read()
    
    with open(tex_filename, 'wb') as tex_file:
        tex_file.write(content)
    
    return tex_filename

# %% ../../nbs/49_helper.arxiv.ipynb 31
def download_from_results(
        results: Result | list[Result],
        dir: PathLike, # The directory into which to download the files
        source: bool = True, # If `True`, download the source file. Otherweise, download a pdf file.
        # filename: Optional[str] = None, # The file name to save the file as. If `None`, then the filename is set to the arXiv id of the article.
        decompress_compressed_file: bool = True, # If `True`and if `source` is `True`, then decompress the source file after downloading it.
        file_or_folder_names: None | str | list[str] | Callable[Result, str] = folder_name_for_source, # If `None`, then the file/folder is named the arxiv id. If a `str` (in which case `results` must be a single `Result` or a `list[Result]` of length 1) or `list[str]` (whose length must equal that of `results`), then each file/folder is named by the specified corresponding `str`. If `Callable[Result, str]`, then each file/folder is named using the specified `Callable` 
        delete_compressed_file: bool = True, # If `True` and if `source` and `decompress_compressed_file` are `True`, then delete the compressed source file after downloading and then uncompressing it.
        download_metadata: bool = True, # If `True`, and if `source` is `True`, then create a file called `metadata.json` and put it into the newly created folder, unless a file called `metadata.json` already exists, in which case, a unique file name is created  
        verbose: bool = False,
        ) -> list[Path]: # Each `Path` is the folder in which the source files are newly downloaded or the path to the pdf file that is newly downloaded.
    """
    Download either the source files or pdfs of the arxiv article encoded in the results.

    - If `source = True` and `decompress_compressed_file = True`, then 
        - Download the source file/folder into a newly created folder (whose name is specified
          by `file_or_folder_names`) within `dir` and decompress the source (if applicable) in
          this newly created folder.
        - If `delete_compressed_file = True`, then delete the compressed file.
    - If `source = False`, then just download a pdf.

    For `file_or_folder_names`, the recommanded `Callable` arguments are `folder_name_for_source`
    for downloading source files and `file_name_for_pdf` for downloading pdf files.

    """
    if not isinstance(results, list):
        results = [results]
    if file_or_folder_names is not None and not isinstance(file_or_folder_names, (list,str)):
        file_or_folder_names = [file_or_folder_names(result) for result in results]
    elif file_or_folder_names is None:
        file_or_folder_names = [None for result in results]
    elif isinstance(file_or_folder_names, str):
        file_or_folder_names = [file_or_folder_names]
    downloaded_paths = []
    for result, file_or_folder_name in zip(results, file_or_folder_names):
        if not source:
            downloaded_paths.append(
                _download_pdf_with_name(result, dir, file_or_folder_name))
            continue
        downloaded_paths.append(
            _download_source(
                result, dir, file_or_folder_name, decompress_compressed_file,
                delete_compressed_file, download_metadata, verbose))
    return downloaded_paths



def _download_pdf_with_name(
        result: Result,
        dir: PathLike,
        file_or_folder_name: str|None # The name of the pdf file. If `None`, then the default name for the pdf file is used.
        ) -> Path:
    """
    Downloads the pdf for `result`.

    Helper function to `download_from_results`
    """
    if file_or_folder_name is not None:
        return Path(result.download_pdf(dir, filename=f'{file_or_folder_name}.pdf'))
    else:
        return Path(result.download_pdf(dir))
    

            

def _create_folder_for_source_download(
        result: Result,
        dir: PathLike,
        file_or_folder_name: str|None # The name of the folder in which the source should be downloaded. If `None`, then the arxiv id for `result` is used.
        ) -> Path: # The newly created folder 
    """
    Creates a new folder inside `dir` in which to download the source for `result`.

    Helper function to `download_from_results`
    """
    if file_or_folder_name is None:
        file_or_folder_name = result.entry_id
    new_folder = Path(dir) / file_or_folder_name
    if os.path.isdir(new_folder): #If folder exists
        # TODO: warn that folder exists
        if file_or_folder_name == result.entry_id:
            file_or_folder_name = f'{file_or_folder_name}_dupl'    
        else:
            file_or_folder_name = f'{file_or_folder_name}_{result.get_short_id()}'
        new_folder = Path(dir) / file_or_folder_name
    while os.path.isdir(new_folder): #If folder still exists
        file_or_folder_name = f'{file_or_folder_name}_dupl'
        new_folder = Path(dir) / file_or_folder_name
    os.mkdir(new_folder)
    return new_folder


def _download_source(
        result: Result,
        dir: PathLike,
        file_or_folder_name: str,
        decompress_compressed_file: bool,
        delete_compressed_file: bool,
        download_metadata: bool,
        verbose: bool
        ) -> Path:
    """
    Download source file into folder, decompress (as needed) the source, and download metadata

    Helper function to `download_from_results`
    """
    new_source_folder = _create_folder_for_source_download(result, dir, file_or_folder_name)
    source_file_path = result.download_source(new_source_folder)
    source_file_path = Path(new_source_folder) / source_file_path
    if verbose:
        print(source_file_path)
    compression_type = analyze_arxiv_tarfile(source_file_path)
    if decompress_compressed_file and file_is_compressed(source_file_path):
        if compression_type in ['plain_gz', 'direct_tex']:
            new_path = Path(f'{str(source_file_path)[:-7]}.gz')
            os.rename(source_file_path, new_path)
            source_file_path = new_path
        uncompressed = uncompress_file(source_file_path)
        if delete_compressed_file:
            os.remove(source_file_path)
        if len(uncompressed) == 1:
            if file_is_compressed(uncompressed[0]) and compression_type == 'nested_archive':
                uncompressed_again = uncompress_file(uncompressed[0])
                if delete_compressed_file:
                    os.remove(uncompressed[0])
            elif compression_type in ['plain_gz', 'direct_tex']:
                os.rename(uncompressed[0], f'{str(uncompressed[0])}.tex')
    if not download_metadata:
        return new_source_folder
    metadata_file_name = _unique_metadata_file_name(new_source_folder)
    metadata = extract_metadata(result)[0] 
    with open(new_source_folder / metadata_file_name, 'w') as json_file:
        json.dump(metadata, json_file, cls=ArxivMetadataEncoder, indent=4)
    return new_source_folder


def _unique_metadata_file_name(
        new_source_folder: PathLike # The folder in which to make the metadata file.
        ) -> str: #s tr: A unique file name for the metadata file within `new_source_folder`.

    """
    Identify a name to name the metadata file within `new_source_folder`; 
    the default name is `metadata.json` unless there is already a file with that name.
    
    If `metadata.json` exists, it appends a numeric suffix to create a unique file name 
    (e.g., `metadata_1.json`, `metadata_2.json`, etc.).
    """
    # Ensure the input is a Path object
    folder = Path(new_source_folder)
    
    # Default file name
    base_name = "metadata"
    extension = ".json"
    candidate = folder / f"{base_name}{extension}"
    
    # Check if the default file name exists
    counter = 1
    while candidate.exists():
        # TODO: warn that metadata.json exists
        # Generate a new candidate with a numeric suffix
        candidate = folder / f"{base_name}_{counter}{extension}"
        counter += 1
    
    return str(candidate.name)  # Return only the file name, not the full path

