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

# UMLS API
UMLS_BASE = "https://uts-ws.nlm.nih.gov/rest"
UMLS_API_KEY = os.getenv("UMLS_API_KEY")  # Add this to your Hugging Face secrets or .env

if not UMLS_API_KEY:
    raise EnvironmentError("UMLS_API_KEY is missing. Add it to your environment variables or Hugging Face secrets.")

# UMLS requires ticket-granting service for authentication
def get_tgt() -> str:
    """
    Obtain a Ticket Granting Ticket (TGT) for UMLS authentication.
    """
    r = requests.post("https://utslogin.nlm.nih.gov/cas/v1/api-key", data={"apikey": UMLS_API_KEY})
    r.raise_for_status()
    from xml.etree import ElementTree
    root = ElementTree.fromstring(r.text)
    return root.attrib.get("action")


def get_service_ticket(tgt: str) -> str:
    """
    Get a single-use service ticket for API requests.
    """
    r = requests.post(tgt, data={"service": "http://umlsks.nlm.nih.gov"})
    r.raise_for_status()
    return r.text


def search_umls(query: str, max_results: int = 10) -> List[Dict]:
    """
    Search UMLS Metathesaurus for a term and return CUIs and names.
    """
    tgt = get_tgt()
    ticket = get_service_ticket(tgt)
    params = {"string": query, "ticket": ticket, "pageSize": max_results}
    
    r = requests.get(f"{UMLS_BASE}/search/current", params=params)
    r.raise_for_status()
    results = r.json().get("result", {}).get("results", [])
    
    return [
        {
            "ui": item.get("ui"),
            "name": item.get("name"),
            "rootSource": item.get("rootSource")
        }
        for item in results if item.get("ui") != "NONE"
    ]


def get_concept_details(cui: str) -> Dict:
    """
    Get full details for a given CUI (Concept Unique Identifier).
    Includes synonyms, definitions, and mappings.
    """
    tgt = get_tgt()
    ticket = get_service_ticket(tgt)
    
    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}", params={"ticket": ticket})
    r.raise_for_status()
    return r.json().get("result", {})


def map_between_ontologies(cui: str) -> List[Dict]:
    """
    Given a CUI, return cross-ontology mappings (e.g., SNOMED → MeSH → ICD).
    """
    tgt = get_tgt()
    ticket = get_service_ticket(tgt)
    
    r = requests.get(f"{UMLS_BASE}/crosswalk/current/source/{cui}", params={"ticket": ticket})
    r.raise_for_status()
    return r.json().get("result", [])


def normalize_term(term: str) -> Optional[Dict]:
    """
    Normalize a free-text biomedical term to its canonical CUI and return metadata.
    """
    matches = search_umls(term, max_results=1)
    return matches[0] if matches else None