app.py

from fastapi import FastAPI, HTTPException, Depends, Request
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
import datetime
import os
import sys
import uuid
import csv
import datetime
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from src.utils.logger import setup_logger
from src.utils.auth import get_current_user, get_active_user

# Absolute import management
project_root = os.path.dirname(os.path.abspath(__file__))
if project_root not in sys.path:
    sys.path.append(project_root)

from src.core.graph import medical_pipeline
from src.core.graph_cdm import cdm_pipeline
from src.core.model_manager import model_manager
from src.agents.agent_instances import update_all_agents_llm
from src.tools.fhir_memory import (
    get_patient_summary_fhir, 
    save_observation, 
    save_patient, 
    create_session, 
    get_sessions_by_patient, 
    get_chat_history_by_session,
    save_chat_as_fhir,
    _get_client
)
from src.mcp.server import MedicalMCPServer

load_dotenv()
logger = setup_logger("FastAPI")

app = FastAPI(title="Medical AI Backend")
mcp_server = MedicalMCPServer()

class ChatMessage(BaseModel):
    role: str
    content: str

class PipelineRequest(BaseModel):
    prompt: str
    patient_id: Optional[str] = None
    session_id: Optional[str] = None
    mode: str = "Standard Triage"  # "Standard Triage" or "CDM Proactive"
    history: List[Dict[str, Any]] = []

class ClinicianScore(BaseModel):
    trace_id: str
    score: int
    feedback: str

class PipelineResponse(BaseModel):
    messages: List[Dict[str, Any]]
    final_state: Dict[str, Any]
    session_id: Optional[str]

def convert_to_langchain_messages(history):
    messages = []
    for msg in history:
        if msg["role"] == "user":
            messages.append(HumanMessage(content=msg["content"]))
        elif msg["role"] == "assistant":
            messages.append(AIMessage(content=msg["content"]))
    return messages

@app.post("/process", response_model=PipelineResponse)
async def process_pipeline(request: PipelineRequest, user_id: str = Depends(get_active_user)):
    # user_id from token is used as the default patient_id if not provided
    patient_id = request.patient_id or user_id
    logger.info(f"Processing request for user: {user_id}, patient: {patient_id}")
    
    # Session handling
    session_id = request.session_id
    if not session_id:
        session_id = create_session.invoke({"patient_id": patient_id, "title": f"Session {datetime.datetime.now().strftime('%Y-%m-%d %H:%M')}"})
        history = request.history
    else:
        # If session_id is provided but history is empty, fetch from Supabase
        history = request.history
        if not history:
            logger.info(f"Fetching history for session: {session_id}")
            raw_history = get_chat_history_by_session.invoke({"session_id": session_id})
            # Convert FHIR Communication resources back to simple role/content dicts
            for comm in raw_history:
                payload = comm.get("payload", [])
                for item in payload:
                    content = item.get("contentString", "")
                    if ":" in content:
                        role, text = content.split(":", 1)
                        history.append({"role": role.strip(), "content": text.strip()})
    
    # Select pipeline
    active_pipeline = cdm_pipeline if request.mode == "CDM Proactive" else medical_pipeline
    
    # Prepare initial state
    enhanced_prompt = f"[System: User's Patient ID is {patient_id}]\n\n{request.prompt}"
    initial_messages = convert_to_langchain_messages(history)
    initial_messages.append(HumanMessage(content=enhanced_prompt))
    
    trace_id = str(uuid.uuid4())
    
    initial_state = {
        "messages": initial_messages,
        "user_role": "unknown",
        "intent_type": "unknown",
        "trace_id": trace_id,
        "session_id": session_id,
        "is_valid": False,
        "is_safe": False,
        "attempts": 0,
        "clinician_outputs": [],
        "patient_response": "",
        "research_output": "",
        "sources": [],
        "logs": [],
        "metrics": []
    }

    initial_state["patient_id"] = patient_id

    try:
        final_state = await active_pipeline.ainvoke(
            initial_state,
            config={"run_name": "MedicalPipeline", "metadata": {"trace_id": trace_id}}
        )
        
        # Format messages for response
        resp_messages = []
        for msg in final_state["messages"][len(initial_messages):]:
            msg_type = "assistant" if isinstance(msg, AIMessage) else "tool" if isinstance(msg, ToolMessage) else "user"
            resp_messages.append({
                "role": msg_type,
                "content": msg.content,
                "type": msg.__class__.__name__
            })
            
        return PipelineResponse(
            messages=resp_messages,
            final_state={k: v for k, v in final_state.items() if k != "messages"},
            session_id=session_id
        )
    except Exception as e:
        logger.error(f"Pipeline error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/sessions")
async def list_sessions(user_id: str = Depends(get_active_user)):
    return get_sessions_by_patient.invoke({"patient_id": user_id})

@app.get("/sessions/{session_id}/history")
async def get_session_history(session_id: str, user_id: str = Depends(get_active_user)):
    return get_chat_history_by_session.invoke({"session_id": session_id})

@app.post("/feedback/score")
async def save_clinician_score(score_data: ClinicianScore, user_id: str = Depends(get_active_user)):
    file_exists = os.path.isfile("clinician_scores.csv")
    with open("clinician_scores.csv", mode="a", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        if not file_exists:
            writer.writerow(["trace_id", "user_id", "score", "feedback", "timestamp"])
        writer.writerow([score_data.trace_id, user_id, score_data.score, score_data.feedback, datetime.datetime.now().isoformat()])
    return {"status": "success", "message": "Score saved successfully"}

@app.get("/patient/summary")
async def get_patient_summary(user_id: str = Depends(get_active_user)):
    try:
        summary = get_patient_summary_fhir.invoke({"patient_id": user_id})
        return summary
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/patient/seed")
async def seed_patient_data(user_id: str = Depends(get_active_user)):
    try:
        save_patient.invoke({"patient_id": user_id, "name": "Authenticated Patient"})
        save_observation.invoke({"patient_id": user_id, "value": 110, "unit": "mg/dL", "display": "Glucose", "loinc_code": "2339-0"})
        return {"status": "success", "message": "Data seeded for user"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/config/llm")
async def set_llm_provider(provider: str, user_id: str = Depends(get_current_user)):
    try:
        update_all_agents_llm(provider)
        return {"status": "success", "provider": model_manager.provider}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/mcp")
async def mcp_endpoint(request: Request):
    body = await request.json()
    return await mcp_server.handle_request(body)

@app.get("/Patient/{patient_id}")
async def get_fhir_patient(patient_id: str):
    client = _get_client()
    res = client.table("patients").select("resource").eq("id", patient_id).execute()
    if res.data:
        return res.data[0]["resource"]
    raise HTTPException(status_code=404, detail="Patient not found")

@app.get("/Observation")
async def get_fhir_observation(patient: str, code: Optional[str] = None):
    client = _get_client()
    query = client.table("observations").select("resource").eq("patient_id", patient)
    res = query.execute()
    observations = [r["resource"] for r in res.data]
    if code:
        observations = [
            obs for obs in observations 
            if any(c.get("code") == code for c in obs.get("code", {}).get("coding", []))
        ]
    return observations

@app.get("/Communication")
async def get_fhir_communication(patient: str):
    client = _get_client()
    res = client.table("communications").select("resource").eq("patient_id", patient).execute()
    return [r["resource"] for r in res.data]

if __name__ == "__main__":
    import uvicorn
    import datetime
    uvicorn.run(app, host="0.0.0.0", port=8000)