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

UMLS_API_KEY = os.getenv("UMLS_API_KEY")
UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
UMLS_BASE_URL = "https://uts-ws.nlm.nih.gov/rest"


def get_tgt():
    """
    Get Ticket-Granting Ticket (TGT) for UMLS API authentication.
    """
    params = {"apikey": UMLS_API_KEY}
    r = requests.post(UMLS_AUTH_ENDPOINT, data=params)
    r.raise_for_status()

    # Extract TGT from HTML response
    tgt_url = r.text.split('action="')[1].split('"')[0]
    return tgt_url


def get_service_ticket(tgt):
    """
    Get Service Ticket (ST) from the TGT for making API requests.
    """
    params = {"service": "http://umlsks.nlm.nih.gov"}
    r = requests.post(tgt, data=params)
    r.raise_for_status()
    return r.text


def search_umls(term: str, max_results: int = 10):
    """
    Search UMLS for a term and return matching concepts.
    """
    tgt = get_tgt()
    st = get_service_ticket(tgt)

    url = f"{UMLS_BASE_URL}/search/current"
    params = {"string": term, "ticket": st, "pageSize": max_results}
    r = requests.get(url, params=params)
    r.raise_for_status()

    data = r.json()
    results = []

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

    return results


def get_concept_details(cui: str):
    """
    Retrieve detailed information for a concept by CUI (Concept Unique Identifier).
    """
    tgt = get_tgt()
    st = get_service_ticket(tgt)

    url = f"{UMLS_BASE_URL}/content/current/CUI/{cui}"
    params = {"ticket": st}
    r = requests.get(url, params=params)
    r.raise_for_status()
    return r.json()


def map_related_concepts(term: str):
    """
    Search UMLS and map related concepts with synonyms and semantic types.
    """
    concepts = search_umls(term, max_results=5)
    mapped = []

    for c in concepts:
        cui = c["ui"]
        if cui and cui != "NONE":
            details = get_concept_details(cui)
            entity = details.get("result", {})

            mapped.append({
                "cui": cui,
                "name": entity.get("name", ""),
                "synonyms": [syn.get("name", "") for syn in entity.get("atoms", [])],
                "semantic_types": [st.get("name", "") for st in entity.get("semanticTypes", [])],
                "rootSource": c.get("rootSource", "")
            })

    return mapped