# mcp/umls.py
import os, httpx
from functools import lru_cache

UMLS_API_KEY = os.getenv("UMLS_KEY")
AUTH_URL     = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
SEARCH_URL   = "https://uts-ws.nlm.nih.gov/rest/search/current"
CONTENT_URL  = "https://uts-ws.nlm.nih.gov/rest/content/current/CUI/{cui}"

async def _get_ticket() -> str:
    async with httpx.AsyncClient(timeout=10) as c:
        r1 = await c.post(AUTH_URL, data={"apikey": UMLS_API_KEY})
        r1.raise_for_status()
        tgt = r1.text.split('action="')[1].split('"')[0]
        r2 = await c.post(tgt, data={"service": "http://umlsks.nlm.nih.gov"})
        r2.raise_for_status()
        return r2.text

@lru_cache(maxsize=512)
async def lookup_umls(term: str) -> dict:
    ticket = await _get_ticket()
    params = {"string": term, "ticket": ticket, "pageSize": 1}
    async with httpx.AsyncClient(timeout=10) as c:
        r = await c.get(SEARCH_URL, params=params)
        r.raise_for_status()
        items = r.json().get("result", {}).get("results", [])
        if not items:
            return {"term": term}
        itm = items[0]
        cui, name = itm.get("ui"), itm.get("name")
        r2 = await c.get(CONTENT_URL.format(cui=cui), params={"ticket": ticket})
        r2.raise_for_status()
        entry = r2.json().get("result", {})
        types = [t["name"] for t in entry.get("semanticTypes", [])]
        definition = entry.get("definitions", [{}])[0].get("value", "")
        return {
            "term": term,
            "cui": cui,
            "name": name,
            "definition": definition,
            "types": types
        }