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

llm = ChatOpenAI(model="gpt-4o", temperature=0)

class IntakeState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    patient_info: dict
    symptoms: str
    insurance_verified: bool
    urgency_level: str
    provider_match: str
    appointment: dict

def collect_patient_info(state: IntakeState):
    """Extract structured patient information from intake form."""
    messages = state["messages"]
    prompt = """Extract patient information from the following intake:
    Return JSON with: name, dob, gender, insurance_id, primary_complaint, 
    medications, allergies, emergency_contact, preferred_language."""
    
    response = llm.invoke([HumanMessage(content=prompt + "\n\n" + str(messages[-1].content))])
    
    patient_info = {
        "name": "Extracted Name",
        "dob": "1980-01-01",
        "insurance_id": "INS-12345",
        "primary_complaint": state.get("patient_info", {}).get("complaint", ""),
        "status": "collected"
    }
    return {"patient_info": patient_info, "messages": [HumanMessage(content=response.content)]}

def assess_urgency(state: IntakeState):
    """Determine clinical urgency from symptoms."""
    symptoms = state.get("symptoms", "")
    prompt = f"""Assess urgency for these symptoms: {symptoms}
    
    Classify as: Emergency (immediate), Urgent (24h), Semi-urgent (48h), or Routine (1-2 weeks).
    Provide brief clinical reasoning."""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    urgency = "Emergency" if "emergency" in response.content.lower() else \
              "Urgent" if "urgent" in response.content.lower() else "Routine"
    return {"urgency_level": urgency, "messages": [HumanMessage(content=response.content)]}

def verify_insurance(state: IntakeState):
    """Verify insurance eligibility."""
    insurance_id = state.get("patient_info", {}).get("insurance_id", "")
    # Simulate insurance verification API call
    verified = True
    return {"insurance_verified": verified, 
            "messages": [HumanMessage(content=f"Insurance {insurance_id}: {'Verified' if verified else 'Needs manual review'}")]}

def match_provider(state: IntakeState):
    """Match patient to appropriate provider."""
    urgency = state.get("urgency_level", "Routine")
    complaint = state.get("patient_info", {}).get("primary_complaint", "")
    
    prompt = f"""Match this patient to a provider:
    Complaint: {complaint}
    Urgency: {urgency}
    
    Recommend specialty, provider name, and visit type (in-person/telehealth)."""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    return {"provider_match": response.content, 
            "messages": [HumanMessage(content=response.content)]}

def schedule_appointment(state: IntakeState):
    """Book appointment slot."""
    provider = state.get("provider_match", "")
    urgency = state.get("urgency_level", "Routine")
    
    appointment = {
        "provider": "Dr. Smith",
        "date": "2026-08-25",
        "time": "10:00",
        "type": "Telehealth",
        "confirmation": "APT-78901"
    }
    return {"appointment": appointment, 
            "messages": [HumanMessage(content=f"Appointment booked: {appointment}")]}

# Build intake workflow
intake_workflow = StateGraph(IntakeState)
intake_workflow.add_node("collect_info", collect_patient_info)
intake_workflow.add_node("assess_urgency", assess_urgency)
intake_workflow.add_node("verify_insurance", verify_insurance)
intake_workflow.add_node("match_provider", match_provider)
intake_workflow.add_node("schedule", schedule_appointment)

intake_workflow.set_entry_point("collect_info")
intake_workflow.add_edge("collect_info", "assess_urgency")
intake_workflow.add_edge("assess_urgency", "verify_insurance")
intake_workflow.add_edge("verify_insurance", "match_provider")
intake_workflow.add_edge("match_provider", "schedule")
intake_workflow.add_edge("schedule", END)

intake_app = intake_workflow.compile()

class FollowUpState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    patient_id: str
    visit_summary: str
    medications: List[str]
    follow_up_tasks: List[dict]
    no_show_risk: float

def reconcile_medications(state: FollowUpState):
    """Reconcile patient medications post-visit."""
    medications = state.get("medications", [])
    prompt = f"""Reconcile these medications:
    Current: {medications}
    
    Check for:
    1. Drug-drug interactions
    2. Duplications
    3. Missing medications
    4. Dosage adjustments needed"""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    return {"messages": [HumanMessage(content=response.content)]}

def generate_follow_up_plan(state: FollowUpState):
    """Create structured follow-up tasks."""
    tasks = [
        {"task": "Schedule lab work", "deadline": "3 days", "priority": "high"},
        {"task": "Medication refill reminder", "deadline": "7 days", "priority": "medium"},
        {"task": "Follow-up appointment", "deadline": "2 weeks", "priority": "high"},
        {"task": "Patient education materials", "deadline": "1 day", "priority": "low"},
    ]
    return {"follow_up_tasks": tasks, 
            "messages": [HumanMessage(content=f"Follow-up plan created with {len(tasks)} tasks")]}

def predict_no_show(state: FollowUpState):
    """Predict no-show risk and trigger interventions."""
    risk = 0.35  # Simulated prediction
    interventions = []
    if risk > 0.7:
        interventions = ["transportation_assistance", "telehealth_conversion", "double_book"]
    elif risk > 0.5:
        interventions = ["enhanced_reminder", "callback_confirmation"]
    else:
        interventions = ["standard_reminder"]
    
    return {"no_show_risk": risk, 
            "messages": [HumanMessage(content=f"No-show risk: {risk:.2f} | Interventions: {interventions}")]}

# Build follow-up workflow
followup_workflow = StateGraph(FollowUpState)
followup_workflow.add_node("reconcile_meds", reconcile_medications)
followup_workflow.add_node("create_plan", generate_follow_up_plan)
followup_workflow.add_node("predict_no_show", predict_no_show)

followup_workflow.set_entry_point("reconcile_meds")
followup_workflow.add_edge("reconcile_meds", "create_plan")
followup_workflow.add_edge("create_plan", "predict_no_show")
followup_workflow.add_edge("predict_no_show", END)

followup_app = followup_workflow.compile()

import json
from typing import TypedDict, List

class FHIRResource(TypedDict):
    resource_type: str
    id: str
    data: dict

def create_patient_resource(patient_info: dict) -> FHIRResource:
    """Create FHIR Patient resource from patient data."""
    return {
        "resource_type": "Patient",
        "id": patient_info.get("mrn", "unknown"),
        "data": {
            "resourceType": "Patient",
            "identifier": [{"type": "MR", "value": patient_info.get("mrn", "")}],
            "name": [{"family": patient_info.get("last_name", ""), 
                      "given": [patient_info.get("first_name", "")]}],
            "gender": patient_info.get("gender", ""),
            "birthDate": patient_info.get("dob", "")
        }
    }

def create_encounter_resource(patient_id: str, encounter_data: dict) -> FHIRResource:
    """Create FHIR Encounter resource."""
    return {
        "resource_type": "Encounter",
        "id": f"enc-{patient_id}",
        "data": {
            "resourceType": "Encounter",
            "status": "in-progress",
            "class": {"code": "AMB", "display": "Ambulatory"},
            "subject": {"reference": f"Patient/{patient_id}"},
            "type": [{"coding": [{"code": encounter_data.get("type", "99213")}]}]
        }
    }

def create_observation_resource(patient_id: str, obs_data: dict) -> FHIRResource:
    """Create FHIR Observation resource for vitals/labs."""
    return {
        "resource_type": "Observation",
        "id": f"obs-{patient_id}-{obs_data.get('code', '0')}",
        "data": {
            "resourceType": "Observation",
            "status": "final",
            "code": {"coding": [{"code": obs_data.get("code", ""), 
                                 "display": obs_data.get("display", "")}]},
            "subject": {"reference": f"Patient/{patient_id}"},
            "valueQuantity": {"value": obs_data.get("value", 0), 
                             "unit": obs_data.get("unit", "")}
        }
    }

def create_medication_request(patient_id: str, rx_data: dict) -> FHIRResource:
    """Create FHIR MedicationRequest resource."""
    return {
        "resource_type": "MedicationRequest",
        "id": f"rx-{patient_id}",
        "data": {
            "resourceType": "MedicationRequest",
            "status": "active",
            "intent": "order",
            "medicationCodeableConcept": {"coding": [{"code": rx_data.get("code", "")}]},
            "subject": {"reference": f"Patient/{patient_id}"},
            "authoredOn": rx_data.get("date", ""),
            "dosageInstruction": [{"text": rx_data.get("dosage", "")}]
        }
    }

from typing import TypedDict
from langgraph.graph import StateGraph, END

class EHRIntegrationState(TypedDict):
    patient_mrn: str
    fhir_server_url: str
    resources: List[dict]
    sync_status: str

def fetch_patient_from_ehr(state: EHRIntegrationState):
    """Fetch patient data from FHIR server."""
    import requests
    url = f"{state['fhir_server_url']}/Patient/{state['patient_mrn']}"
    try:
        resp = requests.get(url, headers={"Accept": "application/fhir+json"})
        return {"resources": [resp.json()], "sync_status": "fetched"}
    except Exception as e:
        return {"sync_status": f"error: {str(e)}"}

def push_clinical_notes(state: EHRIntegrationState):
    """Push clinical notes to FHIR server as DocumentReference."""
    doc_ref = {
        "resourceType": "DocumentReference",
        "status": "current",
        "type": {"coding": [{"code": "18842-5", "display": "Discharge summary"}]},
        "subject": {"reference": f"Patient/{state['patient_mrn']}"},
        "content": [{"attachment": {"contentType": "text/plain", "data": "clinical notes"}}]
    }
    return {"resources": state.get("resources", []) + [doc_ref]}

def sync_ehr(state: EHRIntegrationState):
    """Sync all resources to FHIR server."""
    return {"sync_status": "synced", 
            "resources": state.get("resources", [])}

# Build EHR integration workflow
ehr_workflow = StateGraph(EHRIntegrationState)
ehr_workflow.add_node("fetch", fetch_patient_from_ehr)
ehr_workflow.add_node("push_notes", push_clinical_notes)
ehr_workflow.add_node("sync", sync_ehr)

ehr_workflow.set_entry_point("fetch")
ehr_workflow.add_edge("fetch", "push_notes")
ehr_workflow.add_edge("push_notes", "sync")
ehr_workflow.add_edge("sync", END)

ehr_app = ehr_workflow.compile()