# genesis/api_clients/umls_api.py
import requests
import os

UMLS_API_KEY = os.getenv("UMLS_API_KEY")
AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"

def get_tgt():
    """Get Ticket Granting Ticket for UMLS API."""
    data = {"apikey": UMLS_API_KEY}
    r = requests.post(AUTH_ENDPOINT, data=data)
    r.raise_for_status()
    # Extract TGT URL
    return r.headers.get("location") or r.text.split('action="')[1].split('"')[0]

def get_service_ticket(tgt):
    """Get single-use Service Ticket."""
    r = requests.post(tgt, data={"service": "http://umlsks.nlm.nih.gov"})
    r.raise_for_status()
    return r.text

def search_umls(query: str):
    """Search UMLS for terms."""
    tgt = get_tgt()
    ticket = get_service_ticket(tgt)
    url = f"https://uts-ws.nlm.nih.gov/rest/search/current"
    params = {"string": query, "ticket": ticket}
    r = requests.get(url, params=params)
    r.raise_for_status()
    return r.json()