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

UMLS_API_KEY = os.getenv("UMLS_API_KEY")  # Your UMLS API key from NIH UTS
UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
UMLS_BASE = "https://uts-ws.nlm.nih.gov/rest"

session = requests.Session()

# -------------------------
# Authentication Helpers
# -------------------------
def get_umls_ticket_granting_ticket() -> str:
    """
    Get the Ticket Granting Ticket (TGT) for UMLS API authentication.
    """
    if not UMLS_API_KEY:
        raise ValueError("UMLS API key not found in environment variables.")

    params = {"apikey": UMLS_API_KEY}
    r = session.post(UMLS_AUTH_ENDPOINT, data=params)
    r.raise_for_status()

    # Extract TGT location from XML
    if "location" in r.headers:
        return r.headers["location"]

    # Fallback: parse XML body
    from xml.etree import ElementTree as ET
    root = ET.fromstring(r.text)
    form = root.find(".//form")
    if form is not None and "action" in form.attrib:
        return form.attrib["action"]

    raise RuntimeError("Unable to retrieve UMLS TGT. Check API key or response format.")

def get_umls_service_ticket(tgt: str) -> str:
    """
    Use TGT to request a one-time Service Ticket (ST).
    """
    params = {"service": "http://umlsks.nlm.nih.gov"}
    r = session.post(tgt, data=params)
    r.raise_for_status()
    return r.text.strip()

# -------------------------
# Search Concepts
# -------------------------
def search_umls(term: str, page_size: int = 5) -> List[Dict]:
    """
    Search the UMLS Metathesaurus for a given term.
    Returns a list of concepts with CUI and metadata.
    """
    tgt = get_umls_ticket_granting_ticket()
    st = get_umls_service_ticket(tgt)

    params = {
        "string": term,
        "ticket": st,
        "pageSize": page_size
    }
    r = session.get(f"{UMLS_BASE}/search/current", params=params)
    r.raise_for_status()

    results = r.json().get("result", {}).get("results", [])
    concepts = []
    for res in results:
        if res.get("ui") != "NONE":
            concepts.append({
                "name": res.get("name"),
                "cui": res.get("ui"),
                "rootSource": res.get("rootSource")
            })
    return concepts

# -------------------------
# Concept Details
# -------------------------
def get_concept_details(cui: str) -> Dict:
    """
    Fetch details for a UMLS CUI (definitions, atoms, semantic types).
    """
    tgt = get_umls_ticket_granting_ticket()
    st = get_umls_service_ticket(tgt)

    r = session.get(f"{UMLS_BASE}/content/current/CUI/{cui}", params={"ticket": st})
    r.raise_for_status()
    return r.json().get("result", {})

# -------------------------
# Related Concepts
# -------------------------
def get_related_concepts(cui: str) -> List[Dict]:
    """
    Fetch related concepts for a given CUI.
    """
    tgt = get_umls_ticket_granting_ticket()
    st = get_umls_service_ticket(tgt)

    r = session.get(f"{UMLS_BASE}/content/current/CUI/{cui}/relations", params={"ticket": st})
    r.raise_for_status()

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

# -------------------------
# Term → Semantic Network
# -------------------------
def umls_network(term: str) -> Dict:
    """
    Build a semantic network for a term: concepts, details, and related CUIs.
    """
    concepts = search_umls(term)
    network = []

    for concept in concepts:
        cui = concept["cui"]
        details = get_concept_details(cui)
        related = get_related_concepts(cui)
        network.append({
            "concept": concept,
            "details": details,
            "related": related
        })

    return {
        "term": term,
        "network": network
    }

# -------------------------
# Expansion Helper for pipeline.py
# -------------------------
def expand_with_umls(term: str, max_results: int = 5) -> List[str]:
    """
    Expand a biomedical term into related synonyms & preferred names from UMLS.
    Used by pipeline.py for ontology enrichment.
    """
    try:
        concepts = search_umls(term, page_size=max_results)
        expanded = []
        for c in concepts:
            if c["name"] and c["name"].lower() != term.lower():
                expanded.append(c["name"])
        return list(set(expanded))
    except Exception as e:
        print(f"[UMLS] Expansion failed for '{term}': {e}")
        return []