__version__ = '0.1.0'

from time import sleep
import re
import sys
import argparse
import pybtex.database
import requests
import pybtex.scanner
from pylatexenc.latex2text import LatexNodes2Text
from pyterrier_services import DblpApi

def normalize_title(title):
    title = title.lower()
    title = title.replace('-', ' ')
    return re.sub(r'  +', ' ', title)

def request_with_cooldown(fn):
    try:
        return fn()
    except requests.exceptions.HTTPError as ex:
        if ex.response.status_code == 429:
            sys.stderr.write('  too many requests. cooling down [30sec]...\n')
            sleep(30)
            return request_with_cooldown(fn)
        else:
            raise

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('input')
    args = parser.parse_args()
    dblp = DblpApi()
    latex_converter = LatexNodes2Text()
    buffer = ''
    misses = []
    hits = []
    with open(args.input) as fin:
        for line in fin:
            buffer += line
            try:
                data = pybtex.database.parse_string(buffer, 'bibtex')
            except pybtex.scanner.PrematureEOF:
                pass
            except pybtex.database.input.bibtex.DuplicateField as ex:
                sys.stderr.write(str(ex) + '\n')
                buffer = ''
            else:
                for entry in data.entries.values():
                    title = latex_converter.latex_to_text(entry.fields['title'])
                    sys.stderr.write(str(entry.key) + '\n')
                    norm_title = normalize_title(title)
                    if 'dblp' in entry.fields.get('bibsource', ''):
                        sys.stderr.write('  already dblp\n')
                        print(pybtex.database.BibliographyData({entry.key: entry}).to_string('bibtex'))
                        hits.append(entry.key)
                    else:
                        sys.stderr.write('  searching dblp...\n')
                        dblp_results = request_with_cooldown(lambda: dblp.search(title))
                        sys.stderr.write(f'  {len(dblp_results)} result(s) found\n')
                        for result in dblp_results.itertuples():
                            if normalize_title(result.title) == norm_title:
                                sys.stderr.write(f'  ✅ match on title found for {result.docno}\n')
                                sys.stderr.write('  fetching bibtex...\n')
                                dblp_bibtex = request_with_cooldown(lambda: dblp.load_bibtex(result.docno))
                                dblp_bibtex = pybtex.database.parse_string(dblp_bibtex, 'bibtex')
                                assert len(dblp_bibtex.entries) == 1
                                dblp_entry = list(dblp_bibtex.entries.values())[0]
                                dblp_entry.key = entry.key
                                print(pybtex.database.BibliographyData({entry.key: dblp_entry}).to_string('bibtex'))
                                hits.append(entry.key)
                                break
                        else:
                            sys.stderr.write('  ❌ no matches found. Using original bibtex.\n')
                            print('% TODO: Normalize this entry:')
                            print(f'%   DBLP Search: https://dblp.org/search?q={title.replace(" ", "%20")}')
                            print(pybtex.database.BibliographyData({entry.key: entry}).to_string('bibtex'))
                            misses.append(entry.key)
                buffer = ''
    if len(misses) > 0:
        sys.stderr.write(f'\n❌ Could not find dblp for {len(misses)} (of {len(misses) + len(hits)}) entries (marked with TODO comments):\n')
        for miss in misses:
            sys.stderr.write(f'  {miss}\n')
    else:
        sys.stderr.write(f'\n✅ Found dblp for all {len(hits)} entries.\n')

if __name__ == '__main__':
    main()
