from langchain_core.tools import tool
from pydantic import BaseModel, Field
from typing import List, Optional

class SymptomInput(BaseModel):
    """Patient symptom assessment."""
    age: int = Field(description="Patient age")
    gender: str = Field(description="Male or Female")
    symptoms: List[str] = Field(description="List of symptoms")
    duration: str = Field(description="How long symptoms have lasted")
    severity: str = Field(description="Mild, Moderate, Severe")
    vital_signs: Optional[dict] = Field(description="BP, HR, Temp, RR, SpO2")

@tool
def assess_symptoms(symptoms: str, age: int, gender: str, duration: str) -> str:
    """Assess patient symptoms and determine urgency level.
    
    Args:
        symptoms: Comma-separated list of symptoms
        age: Patient age in years
        gender: Patient gender (Male/Female)
        duration: How long symptoms have persisted
    """
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    prompt = f"""Assess these patient symptoms and determine urgency:
    
    Age: {age}, Gender: {gender}
    Symptoms: {symptoms}
    Duration: {duration}
    
    Provide:
    1. Urgency level (Emergency/Urgent/Routine)
    2. Key concerns
    3. Recommended immediate actions
    
    Format as structured assessment:"""
    
    response = llm.invoke(prompt)
    return response.content


@tool
def generate_differential(symptoms: str, age: int, gender: str, history: str) -> str:
    """Generate differential diagnosis based on symptoms.
    
    Args:
        symptoms: Patient symptoms
        age: Patient age
        gender: Patient gender
        history: Relevant medical history
    """
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    prompt = f"""Generate a differential diagnosis for this patient:
    
    Demographics: {age}yo {gender}
    Symptoms: {symptoms}
    History: {history}
    
    Provide:
    1. Top 5 differential diagnoses (ranked by probability)
    2. For each: Key supporting features, Against features
    3. Recommended diagnostic tests
    4. Red flags to watch for
    
    Use Bayesian reasoning based on prevalence and symptom presentation:"""
    
    response = llm.invoke(prompt)
    return response.content


@tool
def recommend_treatment(diagnosis: str, patient_age: int, allergies: str) -> str:
    """Recommend treatment based on diagnosis.
    
    Args:
        diagnosis: Primary diagnosis
        patient_age: Patient age
        allergies: Known drug allergies
    """
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    prompt = f"""Recommend treatment for:
    
    Diagnosis: {diagnosis}
    Patient Age: {patient_age}
    Allergies: {allergies}
    
    Provide:
    1. First-line treatment
    2. Alternative treatments
    3. Medication dosing
    4. Follow-up plan
    5. Patient education points
    
    Follow current clinical guidelines:"""
    
    response = llm.invoke(prompt)
    return response.content

from typing import TypedDict, Annotated, List
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

class ClinicalState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    patient_data: dict
    urgency_level: str
    differential: str
    treatment: str
    follow_up: str

def triage_node(state: ClinicalState):
    """Assess patient urgency."""
    patient = state["patient_data"]
    result = assess_symptoms.invoke({
        "symptoms": patient.get("symptoms", ""),
        "age": patient.get("age", 0),
        "gender": patient.get("gender", ""),
        "duration": patient.get("duration", ""),
    })
    
    urgency = "Emergency" if any(w in result.lower() for w in ["emergency", "immediate", "critical"]) else \
              "Urgent" if any(w in result.lower() for w in ["urgent", "24 hours"]) else "Routine"
    
    return {"urgency_level": urgency, "messages": [HumanMessage(content=result)]}

def diagnosis_node(state: ClinicalState):
    """Generate differential diagnosis."""
    patient = state["patient_data"]
    result = generate_differential.invoke({
        "symptoms": patient.get("symptoms", ""),
        "age": patient.get("age", 0),
        "gender": patient.get("gender", ""),
        "history": patient.get("history", ""),
    })
    return {"differential": result, "messages": [HumanMessage(content=result)]}

def treatment_node(state: ClinicalState):
    """Recommend treatment."""
    patient = state["patient_data"]
    result = recommend_treatment.invoke({
        "diagnosis": state["differential"][:500],
        "patient_age": patient.get("age", 0),
        "allergies": patient.get("allergies", "None known"),
    })
    return {"treatment": result, "messages": [HumanMessage(content=result)]}

def emergency_check(state: ClinicalState):
    """Route based on urgency."""
    if state["urgency_level"] == "Emergency":
        return "emergency_alert"
    return "diagnosis"

def emergency_alert(state: ClinicalState):
    """Handle emergency cases."""
    return {"messages": [HumanMessage(content="⚠️ EMERGENCY: Activate emergency protocols. Call code team.")]}

# Build graph
workflow = StateGraph(ClinicalState)
workflow.add_node("triage", triage_node)
workflow.add_node("emergency_alert", emergency_alert)
workflow.add_node("diagnosis", diagnosis_node)
workflow.add_node("treatment", treatment_node)

workflow.set_entry_point("triage")
workflow.add_conditional_edges("triage", emergency_check, {
    "emergency_alert": "emergency_alert",
    "diagnosis": "diagnosis",
})
workflow.add_edge("emergency_alert", END)
workflow.add_edge("diagnosis", "treatment")
workflow.add_edge("treatment", END)

app = workflow.compile()

# Run
result = app.invoke({
    "messages": [HumanMessage(content="Patient assessment")],
    "patient_data": {
        "age": 55,
        "gender": "Male",
        "symptoms": "Chest pain, shortness of breath, sweating",
        "duration": "2 hours",
        "history": "Hypertension, diabetes, smoker",
        "allergies": "Penicillin",
    },
    "urgency_level": "",
    "differential": "",
    "treatment": "",
    "follow_up": "",
})
print(result["urgency_level"])
print(result["differential"][:500])