src/tools/fhir_memory.py

import os
import uuid
import datetime
from supabase import create_client, Client
from langchain_core.tools import tool
from src.utils.logger import setup_logger

logger = setup_logger("FHIRSupabaseMemory")

def _get_client() -> Client:
    url = os.getenv("SUPABASE_URL")
    key = os.getenv("SUPABASE_KEY")
    if not url or not key:
        logger.error("SUPABASE_URL or SUPABASE_KEY not found in environment")
        raise ValueError("Supabase configuration missing")
    return create_client(url, key)

def _ensure_patient_exists(patient_id: str, client: Client):
    """Ensure a patient record exists to prevent foreign key constraint violations."""
    try:
        # Check if patient exists
        res = client.table("patients").select("id").eq("id", patient_id).execute()
        if not res.data:
            logger.info(f"Auto-creating missing patient record for {patient_id}")
            fhir_patient = {
                "resourceType": "Patient",
                "id": patient_id,
                "active": True,
                "name": [{"text": "Auto-created Patient", "use": "official"}],
                "meta": {
                    "lastUpdated": datetime.datetime.now(datetime.timezone.utc).isoformat()
                }
            }
            client.table("patients").upsert({
                "id": patient_id,
                "resource": fhir_patient,
                "last_updated": datetime.datetime.now(datetime.timezone.utc).isoformat()
            }).execute()
    except Exception as e:
        logger.error(f"Failed to auto-create patient {patient_id}: {e}")

@tool
def save_patient(patient_id: str, name: str):
    """
    Save a new Patient record in strict FHIR format to Supabase.
    """
    logger.info(f"Saving FHIR Patient: {name} (ID: {patient_id})")
    client = _get_client()
    
    # Construct Strict FHIR Patient
    fhir_patient = {
        "resourceType": "Patient",
        "id": patient_id,
        "active": True,
        "name": [{"text": name, "use": "official"}],
        "meta": {
            "lastUpdated": datetime.datetime.now(datetime.timezone.utc).isoformat()
        }
    }
    
    data = {
        "id": patient_id,
        "resource": fhir_patient,
        "last_updated": datetime.datetime.now(datetime.timezone.utc).isoformat()
    }
    
    try:
        client.table("patients").upsert(data).execute()
        return f"Successfully saved FHIR Patient {name}."
    except Exception as e:
        logger.error(f"Error saving patient: {e}")
        return f"Error saving patient: {str(e)}"

@tool
def create_session(patient_id: str, title: str = "New Chat Session"):
    """
    Create a new chat session for a patient in Supabase.
    """
    logger.info(f"Creating new session for patient: {patient_id}")
    client = _get_client()
    _ensure_patient_exists(patient_id, client)
    session_id = str(uuid.uuid4())
    data = {
        "id": session_id,
        "patient_id": patient_id,
        "title": title
    }
    try:
        client.table("sessions").insert(data).execute()
        return session_id
    except Exception as e:
        logger.error(f"Error creating session: {e}")
        return None

@tool
def get_observations_by_patient(patient_id: str, loinc_code: str = None):
    """
    Retrieve FHIR Observation resources for a specific patient.
    """
    logger.info(f"Retrieving FHIR observations for patient: {patient_id}")
    client = _get_client()
    
    # Query using the top-level patient_id column for performance
    query = client.table("observations").select("resource").eq("patient_id", patient_id)
    
    try:
        response = query.execute()
        # Extract the 'resource' part of each record
        observations = [record["resource"] for record in response.data]
        
        if loinc_code:
            # Filter in Python for nested JSON matching if needed
            observations = [
                obs for obs in observations 
                if any(c.get("code") == loinc_code for c in obs.get("code", {}).get("coding", []))
            ]
            
        if not observations:
            return f"No FHIR observations found for patient {patient_id}."
            
        return observations
    except Exception as e:
        logger.error(f"Error retrieving observations: {e}")
        return f"Error retrieving observations: {str(e)}"

@tool
def save_observation(patient_id: str, value: float, unit: str, display: str, loinc_code: str):
    """
    Save a new health observation in strict FHIR format to Supabase.
    """
    logger.info(f"Saving FHIR observation: {display} for patient: {patient_id}")
    client = _get_client()
    _ensure_patient_exists(patient_id, client)
    
    obs_id = str(uuid.uuid4())
    
    # Construct Strict FHIR Observation
    fhir_observation = {
        "resourceType": "Observation",
        "id": obs_id,
        "status": "final",
        "category": [{
            "coding": [{
                "system": "http://terminology.hl7.org/CodeSystem/observation-category",
                "code": "vital-signs",
                "display": "Vital Signs"
            }]
        }],
        "code": {
            "coding": [{
                "system": "http://loinc.org",
                "code": loinc_code,
                "display": display
            }]
        },
        "subject": {"reference": f"Patient/{patient_id}"},
        "effectiveDateTime": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "valueQuantity": {
            "value": value,
            "unit": unit,
            "system": "http://unitsofmeasure.org",
            "code": unit 
        }
    }
    
    data = {
        "id": obs_id,
        "patient_id": patient_id,
        "resource": fhir_observation,
        "last_updated": datetime.datetime.now(datetime.timezone.utc).isoformat()
    }
    
    try:
        client.table("observations").insert(data).execute()
        return f"Successfully saved FHIR {display} observation."
    except Exception as e:
        logger.error(f"Error saving observation: {e}")
        return f"Error saving observation: {str(e)}"

@tool
def get_patient_summary_fhir(patient_id: str):
    """
    Get a comprehensive health summary composed of strict FHIR resources.
    """
    logger.info(f"Generating FHIR summary for patient: {patient_id}")
    client = _get_client()
    
    try:
        # Get Patient Resource
        p_resp = client.table("patients").select("resource").eq("id", patient_id).single().execute()
        patient_resource = p_resp.data["resource"]
        
        # Get Observation Resources
        o_resp = client.table("observations").select("resource").eq("patient_id", patient_id).order("last_updated", desc=True).limit(10).execute()
        observation_resources = [r["resource"] for r in o_resp.data]
        
        summary = {
            "patient": patient_resource.get("name", [{"text": "Unknown"}])[0].get("text"),
            "full_patient_resource": patient_resource,
            "recent_observations": observation_resources,
            "active_medications": [] 
        }
        
        return summary
    except Exception as e:
        logger.error(f"Error generating FHIR summary: {e}")
        return f"Error generating FHIR summary: {str(e)}"

@tool
def save_chat_as_fhir(patient_id: str, messages: list, session_id: str = None):
    """
    Save chat history as a strict FHIR Communication resource.
    """
    logger.info(f"Saving FHIR Communication for patient: {patient_id}, session: {session_id}")
    client = _get_client()
    _ensure_patient_exists(patient_id, client)
    
    comm_id = str(uuid.uuid4())
    
    # Construct Strict FHIR Communication
    fhir_communication = {
        "resourceType": "Communication",
        "id": comm_id,
        "status": "completed",
        "subject": {"reference": f"Patient/{patient_id}"},
        "sent": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "payload": [
            {"contentString": f"{msg.get('role', 'unknown')}: {msg.get('content', '')}"}
            for msg in messages
        ]
    }
    
    data = {
        "id": comm_id,
        "patient_id": patient_id,
        "session_id": session_id,
        "resource": fhir_communication,
        "last_updated": datetime.datetime.now(datetime.timezone.utc).isoformat()
    }
    
    try:
        client.table("communications").insert(data).execute()
        return f"Successfully saved FHIR chat history."
    except Exception as e:
        logger.error(f"Error saving communication: {e}")
        return f"Error saving communication: {str(e)}"

@tool
def get_sessions_by_patient(patient_id: str):
    """
    Retrieve all chat sessions for a specific patient.
    """
    logger.info(f"Retrieving sessions for patient: {patient_id}")
    client = _get_client()
    try:
        response = client.table("sessions").select("*").eq("patient_id", patient_id).order("created_at", desc=True).execute()
        return response.data
    except Exception as e:
        logger.error(f"Error retrieving sessions: {e}")
        return []

@tool
def get_chat_history_by_session(session_id: str):
    """
    Retrieve all communications (chat history) for a specific session.
    """
    logger.info(f"Retrieving chat history for session: {session_id}")
    client = _get_client()
    try:
        response = client.table("communications").select("resource").eq("session_id", session_id).order("last_updated", asc=True).execute()
        return [r["resource"] for r in response.data]
    except Exception as e:
        logger.error(f"Error retrieving chat history: {e}")
        return []

@tool
def get_medications_by_patient(patient_id: str):
    """
    Retrieve current medications for a patient in FHIR MedicationRequest format.
    Bug 7.3: Retrieve medications to support CDM pipeline context.
    """
    logger.info(f"Retrieving medications for patient: {patient_id}")
    client = _get_client()
    
    try:
        # Query medications table if it exists
        query = client.table("medications").select("resource").eq("patient_id", patient_id)
        response = query.execute()
        
        if response.data:
            medications = [record["resource"] for record in response.data]
            return medications
        
        # Return empty list with informative message if no medications found
        return []
    except Exception as e:
        logger.error(f"Error retrieving medications: {e}")
        # Return empty list on error to prevent pipeline failures
        return []

@tool
def save_medication(patient_id: str, medication_name: str, dosage: str, frequency: str, status: str = "active"):
    """
    Save a medication record for a patient in FHIR MedicationRequest format.
    Bug 7.3: Store medications to support CDM pipeline context.
    """
    logger.info(f"Saving medication for patient: {patient_id}")
    client = _get_client()
    _ensure_patient_exists(patient_id, client)
    
    med_id = str(uuid.uuid4())
    
    # Construct FHIR MedicationRequest
    fhir_medication_request = {
        "resourceType": "MedicationRequest",
        "id": med_id,
        "status": status,
        "intent": "order",
        "subject": {"reference": f"Patient/{patient_id}"},
        "authoredOn": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "medication": {
            "coding": [{
                "system": "http://www.nlm.nih.gov/research/umls/rxnorm",
                "display": medication_name
            }]
        },
        "dosageInstruction": [{
            "text": f"{dosage} {frequency}"
        }]
    }
    
    data = {
        "id": med_id,
        "patient_id": patient_id,
        "resource": fhir_medication_request,
        "last_updated": datetime.datetime.now(datetime.timezone.utc).isoformat()
    }
    
    try:
        # Check if medications table exists, if not create in memory
        client.table("medications").insert(data).execute()
        return f"Successfully saved medication: {medication_name}"
    except Exception as e:
        logger.error(f"Error saving medication: {e}")
        return f"Error saving medication: {str(e)}"

@tool
def ingest_fhir_bundle(bundle: dict):
    """
    Ingest a FHIR Bundle into Supabase by splitting it into individual resources.
    """
    if bundle.get("resourceType") != "Bundle":
        return "Error: Input must be a FHIR Bundle."
    
    client = _get_client()
    entries = bundle.get("entry", [])
    count = 0
    
    for entry in entries:
        resource = entry.get("resource")
        if not resource: continue
        
        rtype = resource.get("resourceType").lower()
        res_id = resource.get("id") or str(uuid.uuid4())
        resource["id"] = res_id
        
        table_map = {
            "patient": "patients",
            "observation": "observations",
            "communication": "communications"
        }
        
        target_table = table_map.get(rtype)
        if not target_table: continue
        
        data = {
            "id": res_id,
            "resource": resource,
            "last_updated": datetime.datetime.now(datetime.timezone.utc).isoformat()
        }
        
        if rtype in ["observation", "communication"]:
            ref = resource.get("subject", {}).get("reference", "")
            if "Patient/" in ref:
                data["patient_id"] = ref.split("/")[-1]
        
        client.table(target_table).upsert(data).execute()
        count += 1
        
    return f"Successfully ingested {count} resources from FHIR Bundle."