process_latex.py

import traceback
import requests
import tarfile
import io
import os
import re
import time
from typing import List, Tuple, Set, Dict, Optional
from diskcache import Cache
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry


from parse_bbl import extract_citations_and_titles



cache_dir = os.path.join(os.path.dirname(__file__), 'cache', 'content')
os.makedirs(cache_dir, exist_ok=True)
cache = Cache(cache_dir, size_limit=100*1000**3)


session = requests.session()
retry_times = 5
retry_backoff_factor = 2
retry = Retry(total=retry_times, backoff_factor=retry_backoff_factor)
adapter = HTTPAdapter(max_retries=retry)
session.mount('http://', adapter)
session.mount('https://', adapter)


def download_file(url: str) -> bytes:
    cache_key = url
    
    if cache_key in cache:
        print(f"Load form cache: {url}")
        return cache[cache_key]
    
    print(f"Downloading file: {url}")
    response = session.get(url, timeout=30)
    response.raise_for_status()
    content = response.content
    
    cache[cache_key] = content
    
    time.sleep(1)
    
    return content


def find_file_in_contents(file_path: str, file_contents: Dict[str, str], base_dir: str = '') -> Optional[str]:
    file_path = normalize_path(file_path)
    base_dir = normalize_path(base_dir)
    
    possible_paths = [
        file_path,
        os.path.join(base_dir, file_path),
        file_path.lstrip('/'),
        os.path.join(base_dir, file_path.lstrip('/'))
    ]
    
    possible_paths = [normalize_path(p) for p in possible_paths]
    possible_paths = list(dict.fromkeys(possible_paths))
    
    for path in possible_paths:
        if path in file_contents:
            return path
            
    return None


def find_tex_includes(content: str, file_contents: Dict[str, str] = None, base_dir: str = '') -> List[str]:
    pattern = r'\\(?:input|include)\s*\{([^}]+)\}'
    matches = re.finditer(pattern, content)
    
    included_files = []
    for match in matches:
        filename = match.group(1)
        if not filename.lower().endswith('.tex'):
            filename += '.tex'
            
        if file_contents is not None:
            actual_path = find_file_in_contents(filename, file_contents, base_dir)
            if actual_path:
                included_files.append(actual_path)
                continue
                
        included_files.append(filename)
    
    return included_files


def is_main_tex_file(content: str) -> bool:
    has_documentclass = bool(re.search(r'\\documentclass', content))
    has_document_env = bool(re.search(r'\\begin\s*{\s*document\s*}.*\\end\s*{\s*document\s*}', 
                                    content, re.DOTALL))
    return has_documentclass and has_document_env


def normalize_path(path: str) -> str:
    return path.replace('\\', '/')


def load_all_tex_files(
    tar_file: tarfile.TarFile,
    load_bib=True,
) -> tuple[Dict[str, str], list]:
    file_contents = {}
    bib_records = []
    
    
    for member in tar_file.getmembers():
        if not member.isfile():
            continue
        
        if load_bib and member.name.lower().endswith('.bbl'):
            try:
                with tar_file.extractfile(member) as bibtex_file:
                    data = bibtex_file.read().decode('utf-8')
                    titles = extract_citations_and_titles(data)
                    bib_records += [
                        {
                            'ID': key,
                            'title': candidates,
                        }
                        for key, candidates in titles.items()
                    ]
            except:
                print(f"Warning: Cannot read file {member.name}")
                continue
        elif member.name.lower().endswith('.tex'):
            try:
                with tar_file.extractfile(member) as f:
                    content = f.read().decode('utf-8')
                    file_contents[normalize_path(member.name)] = content
            except:
                print(f"Warning: Cannot read file {member.name}")
                traceback.print_exc()
                continue
    
    return file_contents, bib_records


def find_main_tex_file(file_contents: Dict[str, str]) -> Optional[str]:
    references: Dict[str, List[str]] = {}
    referenced_by: Dict[str, List[str]] = {}
    main_candidates = []
    
    for filepath, content in file_contents.items():
        if is_main_tex_file(content):
            main_candidates.append(filepath)
        
        included_files = find_tex_includes(content, file_contents, os.path.dirname(filepath))
        resolved_includes = [
            normalize_path(inc)
            for inc in included_files
        ]
        
        references[filepath] = resolved_includes
        
        for included in resolved_includes:
            if included not in referenced_by:
                referenced_by[included] = []
            referenced_by[included].append(filepath)
    
    main_file = None
    if main_candidates:
        unreferenced_candidates = [
            f for f in main_candidates
            if f not in referenced_by or not referenced_by[f]
        ]
        if unreferenced_candidates:
            main_file = max(
                unreferenced_candidates,
                key=lambda f: len(references.get(f, []))
            )
        else:
            main_file = min(
                main_candidates,
                key=lambda f: len(referenced_by.get(f, []))
            )
    
    return main_file


def merge_tex_files(main_content: str, file_contents: Dict[str, str], base_dir: str = '', 
                   processed: Set[str] = None) -> str:
    if processed is None:
        processed = set()
    
    def replace_include(match: re.Match) -> str:
        filename = match.group(1)
        if not filename.lower().endswith('.tex'):
            filename += '.tex'
        path = find_file_in_contents(filename, file_contents, base_dir)
        
        if path is None:
            print(f"Warning: File not found {filename}")
            return ''
        
        if path not in processed:
            processed.add(path)
            content = file_contents[path]
            merged = merge_tex_files(content, file_contents, 
                                    os.path.dirname(path), processed)
            return merged
        
    
    pattern = r'\\(?:input|include)\s*\{([^}]+)\}'
    merged_content = re.sub(pattern, replace_include, main_content)
    
    return merged_content


def process_tex_content(content: str) -> List[str]:
    ESCAPED_PERCENT_PLACEHOLDER = '__ESCAPED_PERCENT_dafjbrsgbjse__'
    
    lines = content.split('\n')
    processed_lines = []
    for line in lines:
        temp = line.replace(r'\%', ESCAPED_PERCENT_PLACEHOLDER)
        temp = re.sub(r'%.*$', '', temp)
        temp = temp.replace(ESCAPED_PERCENT_PLACEHOLDER, r'\%')
        processed_lines.append(temp.rstrip())
    return processed_lines


def download_and_process_latex(url: str, load_bib=True) -> Tuple[List[str], str, list]:
    targz_content = download_file(url)
    fileobj = io.BytesIO(targz_content)
    
    with tarfile.open(fileobj=fileobj, mode='r:*') as tar_file:
        file_contents, bib_records = load_all_tex_files(tar_file, load_bib=load_bib)
        
        main_tex_path = find_main_tex_file(file_contents)
        
        if not main_tex_path:
            raise Exception("Failed to find main tex file")
        
        print(f"Found main tex file: {main_tex_path}")
        
        main_content = file_contents.get(main_tex_path, '')
        base_dir = os.path.dirname(main_tex_path)
        merged_content = merge_tex_files(main_content, file_contents, base_dir)
        
        processed_lines = process_tex_content(merged_content)
        full_text = '\n'.join(processed_lines)
        
        print(f"Processed files:")
        successful_files = sorted(file_contents.keys())
        for file in successful_files:
            print(f"- {file}")
        
        print(full_text, file=open('tmp.txt', 'w'))
        
        return processed_lines, full_text, bib_records