main.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
import datetime
import os
import sys
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from src.utils.logger import setup_logger

# 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_chat_history_by_session,
    _get_client,
)
from src.utils.auth import create_dev_token
import re
import time
import uuid

# Simple in-memory rate limiter: keys map to list of request timestamps
_RATE_LIMIT_WINDOW = 60  # seconds
_RATE_LIMIT_MAX = int(os.getenv("RATE_LIMIT_PER_MINUTE", "30"))
_rate_store = {}

def _check_rate_limit(key: str):
    now = time.time()
    bucket = _rate_store.get(key, [])
    # drop old
    bucket = [t for t in bucket if now - t < _RATE_LIMIT_WINDOW]
    if len(bucket) >= _RATE_LIMIT_MAX:
        return False
    bucket.append(now)
    _rate_store[key] = bucket
    return True

_PROMPT_INJECTION_PATTERNS = [
    r"ignore (system|instructions|previous|above)",
    r"disregard (previous|above|system)",
    r"do not follow (system|instructions)",
    r"override (system|instructions)",
]

def _detect_prompt_injection(text: str) -> bool:
    if not text:
        return False
    for p in _PROMPT_INJECTION_PATTERNS:
        if re.search(p, text, re.IGNORECASE):
            return True
    return False

load_dotenv()
logger = setup_logger("FastAPI")

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

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins for local development
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

@app.get("/")
async def root():
    return {"status": "healthy", "message": "Medical AI Backend is running"}

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 PipelineResponse(BaseModel):
    messages: List[Dict[str, Any]]
    final_state: Dict[str, Any]
    session_id: Optional[str] = None

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


def load_session_history(session_id: str):
    if not session_id:
        return []

    raw_history = get_chat_history_by_session.invoke({"session_id": session_id})
    history = []
    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()})
    return history

import json
from fastapi.responses import StreamingResponse

@app.post("/process_stream")
async def process_pipeline_stream(request: PipelineRequest):
    logger.info(f"Streaming request for mode: {request.mode}")

    history = list(request.history)
    session_id = request.session_id
    if not session_id and request.patient_id:
        session_id = create_session.invoke(
            {
                "patient_id": request.patient_id,
                "title": f"Session {datetime.datetime.now().strftime('%Y-%m-%d %H:%M')}",
            }
        )
    if session_id and not history:
        history = load_session_history(session_id)

    active_pipeline = cdm_pipeline if request.mode == "CDM Proactive" else medical_pipeline

    enhanced_prompt = request.prompt
    if request.patient_id:
        enhanced_prompt = f"[System: User's Patient ID is {request.patient_id}]\n\n{request.prompt}"

    initial_messages = convert_to_langchain_messages(history)
    initial_messages.append(HumanMessage(content=enhanced_prompt))

    initial_state = {
        "messages": initial_messages,
        "user_role": "unknown",
        "intent_type": "unknown",
        "session_id": session_id,
        "is_valid": False,
        "is_safe": False,
        "attempts": 0,
        "clinician_outputs": [],
        "patient_response": "",
        "research_output": "",
        "sources": [],
        "logs": [],
        "metrics": [],
    }

    if request.patient_id:
        initial_state["patient_id"] = request.patient_id

    final_state = None

    async def event_generator():
        saw_stream = False

        def chunk_text(text: str, size: int = 24):
            for start in range(0, len(text), size):
                yield text[start:start + size]

        nonlocal final_state
        try:
            async for event in active_pipeline.astream_events(initial_state, version="v2"):
                kind = event["event"]

                # Progress Update: Node start
                if kind == "on_chain_start" and event.get("name") in [
                    "role_classifier", "patient_llm", "caregiver_llm", "safety_check", "validator", 
                    "intent_classifier", "persistence_node", "tools_node",
                    "diagnosis_assist", "treatment_assist", "monitoring_assist", "general_assist",
                    "merge_outputs", "research_agent", "dietary_assist"
                ]:
                    yield f"data: {json.dumps({'type': 'node', 'node': event['name']})}\n\n"

                # Progress Update: Graph Nodes
                if kind == "on_chain_start" and event.get("name") in [
                    "role_classifier",
                    "patient_llm",
                    "safety_check",
                    "validator",
                    "intent_classifier",
                    "persistence_node",
                    "tools_node",
                ]:
                    yield f"data: {json.dumps({'type': 'node', 'node': event['name']})}\n\n"

                elif kind == "on_chain_end" and "node" in event.get("metadata", {}):
                    node_name = event["metadata"]["node"]
                    yield f"data: {json.dumps({'type': 'node_complete', 'node': node_name})}\n\n"

                elif kind == "on_chat_model_stream":
                    content = getattr(event["data"]["chunk"], "content", "")
                    if content:
                        saw_stream = True
                        yield f"data: {json.dumps({'type': 'token', 'content': content})}\n\n"

                # Final State: End of graph
                elif kind == "on_chain_end" and event["name"] == "LangGraph":
                    final_state = event["data"].get("output", {}) or {}
                    final_msg = ""
                    
                    # Extract the final message from the state
                    if "messages" in final_state and final_state["messages"]:
                        last_msg = final_state["messages"][-1]
                        final_msg = last_msg.content if hasattr(last_msg, "content") else str(last_msg)
                    
                    # Format state for frontend (exclude messages to save bandwidth)
                    clean_state = {k: v for k, v in final_state.items() if k != "messages"}
                    
                    yield f"data: {json.dumps({'type': 'end', 'final_state': clean_state, 'final_message': final_msg})}\n\n"

        except Exception as e:
            logger.error(f"Streaming error: {str(e)}")
            yield f"data: {json.dumps({'type': 'error', 'detail': str(e)})}\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")

@app.post("/process", response_model=PipelineResponse)
async def process_pipeline(request: PipelineRequest):
    logger.info(f"Processing request for mode: {request.mode}")

    history = list(request.history)
    session_id = request.session_id
    if not session_id and request.patient_id:
        session_id = create_session.invoke(
            {
                "patient_id": request.patient_id,
                "title": f"Session {datetime.datetime.now().strftime('%Y-%m-%d %H:%M')}",
            }
        )
    if session_id and not history:
        history = load_session_history(session_id)

    active_pipeline = cdm_pipeline if request.mode == "CDM Proactive" else medical_pipeline

    enhanced_prompt = request.prompt
    if request.patient_id:
        enhanced_prompt = f"[System: User's Patient ID is {request.patient_id}]\n\n{request.prompt}"

    initial_messages = convert_to_langchain_messages(history)
    initial_messages.append(HumanMessage(content=enhanced_prompt))

    initial_state = {
        "messages": initial_messages,
        "user_role": "unknown",
        "intent_type": "unknown",
        "session_id": session_id,
        "is_valid": False,
        "is_safe": False,
        "attempts": 0,
        "clinician_outputs": [],
        "patient_response": "",
        "research_output": "",
        "sources": [],
        "logs": [],
        "metrics": [],
    }

    if request.patient_id:
        initial_state["patient_id"] = request.patient_id

    try:
        final_state = await active_pipeline.ainvoke(initial_state)

        resp_messages = []
        for msg in final_state["messages"][len(initial_messages):]:
            from langchain_core.messages import AIMessage, ToolMessage
            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("/patient/{patient_id}")
async def get_patient_summary(patient_id: str):
    try:
        summary = get_patient_summary_fhir.invoke({"patient_id": patient_id})
        return summary
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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

@app.post("/config/llm")
async def set_llm_provider(provider: str):
    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))

if __name__ == "__main__":
    import uvicorn
    import os
    port = int(os.environ.get("PORT", 8000))
    logger.info(f"Starting server on port {port}")
    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=False)


class RegisterRequest(BaseModel):
    username: str
    first_name: str
    last_name: str

@app.post("/auth/register")
async def dev_register(req: RegisterRequest):
    if not req.username.strip():
        raise HTTPException(status_code=400, detail="Username is required")
    if not req.first_name.strip():
        raise HTTPException(status_code=400, detail="First name is required")
    if not req.last_name.strip():
        raise HTTPException(status_code=400, detail="Last name is required")
    
    client = _get_client()
    try:
        # Check if username already exists
        res = client.table("patients").select("id").eq("resource->>username", req.username.strip()).execute()
        if res.data:
            raise HTTPException(status_code=400, detail="Username already exists")
        
        pid = str(uuid.uuid4())
        full_name = f"{req.first_name.strip()} {req.last_name.strip()}"
        fhir_patient = {
            "resourceType": "Patient",
            "id": pid,
            "active": True,
            "name": [{
                "text": full_name,
                "use": "official",
                "given": [req.first_name.strip()],
                "family": req.last_name.strip()
            }],
            "username": req.username.strip(),
            "meta": {
                "lastUpdated": datetime.datetime.now(datetime.timezone.utc).isoformat()
            }
        }
        data = {
            "id": pid,
            "resource": fhir_patient,
            "last_updated": datetime.datetime.now(datetime.timezone.utc).isoformat()
        }
        client.table("patients").insert(data).execute()
        
        token = create_dev_token(pid, expires_minutes=24 * 60)
        return {"status": "ok", "patient_id": pid, "token": token}
    except HTTPException as he:
        raise he
    except Exception as e:
        logger.error(f"Failed to register: {e}")
        raise HTTPException(status_code=500, detail=str(e))


class LoginRequest(BaseModel):
    username: str

@app.post("/auth/login")
async def dev_login(req: LoginRequest):
    """Development-only login endpoint that verifies username.
    """
    if not req.username:
        raise HTTPException(status_code=400, detail="Username required")
    
    client = _get_client()
    try:
        res = client.table("patients").select("*").eq("resource->>username", req.username.strip()).execute()
        if not res.data:
            raise HTTPException(status_code=404, detail="Username not found. Please register first.")
        
        patient = res.data[0]
        pid = patient["id"]
        token = create_dev_token(pid, expires_minutes=24 * 60)
        return {"status": "ok", "patient_id": pid, "token": token}
    except HTTPException as he:
        raise he
    except Exception as e:
        logger.error(f"Failed to login: {e}")
        raise HTTPException(status_code=500, detail=str(e))