from fastapi import FastAPI, HTTPException
import urllib.request
from urllib.error import HTTPError
from habanero import Crossref


app = FastAPI()

BASE_URL = 'https://doi.org/'

@app.get("/title2ref/")
async def title2doi(title:str):
    cr = Crossref()

    result = cr.works(query = title)
    for i, item in enumerate(result['message']['items']):
        if 'group-title' in item and item['group-title'] != 'Preprints':
            doi = result['message']['items'][i-1]['DOI']
            break
    
    url = BASE_URL + doi
    req = urllib.request.Request(url)
    req.add_header('Accept', 'application/x-bibtex')

    try:
        with urllib.request.urlopen(req) as f:
            bibtex = f.read().decode()
        return {"bibtex": bibtex}
    except HTTPError as e:
        if e.code == 404:
            raise HTTPException(status_code=404, detail="DOI not found")
        else:
            raise HTTPException(status_code=503, detail="Service unavailable")
            
@app.get("/doi2ref/")
async def doi2ref(doi: str):
    url = BASE_URL + doi
    req = urllib.request.Request(url)
    req.add_header('Accept', 'application/x-bibtex')

    try:
        with urllib.request.urlopen(req) as f:
            bibtex = f.read().decode()
        return {"bibtex": bibtex}
    except HTTPError as e:
        if e.code == 404:
            raise HTTPException(status_code=404, detail="DOI not found")
        else:
            raise HTTPException(status_code=503, detail="Service unavailable")

if __name__ == "__main__":
    import uvicorn

    uvicorn.run("main:app", host="0.0.0.0", port=7860)