from fastapi import FastAPI, HTTPException
from rdkit import Chem
from rdkit.Chem import Descriptors
from rdkit.Chem import rdMolDescriptors
import requests
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()
# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

def name_to_smiles(name: str) -> str:
    url = f"https://cactus.nci.nih.gov/chemical/structure/{name}/smiles"
    response = requests.get(url)
    if response.status_code == 200:
        return response.text
    return None

def get_molecule_info(mol):
    mol_weight = Descriptors.MolWt(mol)
    num_atoms = mol.GetNumAtoms()
    num_bonds = mol.GetNumBonds()
    mol_formula = rdMolDescriptors.CalcMolFormula(mol)
    tpsa = Descriptors.TPSA(mol)
    mol_logp = Descriptors.MolLogP(mol)
    num_rotatable_bonds = Descriptors.NumRotatableBonds(mol)

    return {
        'molecular_weight': mol_weight,
        'number_of_atoms': num_atoms,
        'number_of_bonds': num_bonds,
        'molecular_formula': mol_formula,
        'tpsa': tpsa,
        'logP': mol_logp,
        'number_of_rotatable_bonds': num_rotatable_bonds
    }

@app.get("/molecule_info/")
async def read_molecule_info(name: str):
    if not name:
        raise HTTPException(status_code=400, detail="No molecule name provided")

    smiles = name_to_smiles(name)
    if not smiles:
        raise HTTPException(status_code=400, detail="Could not fetch SMILES string for provided name")

    mol = Chem.MolFromSmiles(smiles)

    if not mol:
        raise HTTPException(status_code=400, detail="Molecule not recognized")

    info = get_molecule_info(mol)

    return info

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)