# genesis/api_clients/umls_api.py
import os
import requests
from typing import List, Dict, Optional

UMLS_API_KEY = os.getenv("UMLS_API_KEY")  # Set in Hugging Face secrets or .env
UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
UMLS_BASE = "https://uts-ws.nlm.nih.gov/rest"

# -------------------------
# AUTHENTICATION
# -------------------------
def get_tgt() -> str:
    """
    Get Ticket Granting Ticket (TGT) from UMLS API key.
    """
    params = {"apikey": UMLS_API_KEY}
    r = requests.post(UMLS_AUTH_ENDPOINT, data=params)
    r.raise_for_status()
    # Extract TGT from HTML form action
    tgt_url = r.text.split('action="')[1].split('"')[0]
    return tgt_url

def get_service_ticket(tgt_url: str) -> str:
    """
    Get a one-time Service Ticket (ST) for API calls.
    """
    params = {"service": "http://umlsks.nlm.nih.gov"}
    r = requests.post(tgt_url, data=params)
    r.raise_for_status()
    return r.text

# -------------------------
# SEARCH FUNCTIONS
# -------------------------
def search_umls(term: str, max_results: int = 10) -> List[Dict]:
    """
    Search UMLS for a medical concept by term and return CUIs, preferred names, and semantic types.
    """
    tgt = get_tgt()
    st = get_service_ticket(tgt)
    params = {
        "string": term,
        "ticket": st,
        "pageSize": max_results
    }
    r = requests.get(f"{UMLS_BASE}/search/current", params=params)
    r.raise_for_status()
    data = r.json()

    results = []
    for item in data.get("result", {}).get("results", []):
        results.append({
            "ui": item.get("ui"),  # CUI
            "name": item.get("name"),
            "rootSource": item.get("rootSource")
        })
    return results

def get_cui_details(cui: str) -> Dict:
    """
    Retrieve details for a specific CUI.
    """
    tgt = get_tgt()
    st = get_service_ticket(tgt)
    params = {"ticket": st}
    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}", params=params)
    r.raise_for_status()
    return r.json()

def get_cui_relations(cui: str) -> List[Dict]:
    """
    Get relationships for a CUI (e.g., associated diseases, drugs, symptoms).
    """
    tgt = get_tgt()
    st = get_service_ticket(tgt)
    params = {"ticket": st}
    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}/relations", params=params)
    r.raise_for_status()
    data = r.json()

    relations = []
    for rel in data.get("result", []):
        relations.append({
            "relationLabel": rel.get("relationLabel"),
            "relatedId": rel.get("relatedId"),
            "relatedName": rel.get("relatedIdName")
        })
    return relations