import os
from typing import TypedDict, Literal
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from pydantic import BaseModel, Field

# 1. Define the State
class AgentState(TypedDict):
    shipment_id: str
    is_exception: str
    resolution: str
    rationale: str
    customer_message: str
    escalated: bool

# 2. Define Output Schemas
class ResolutionOutput(BaseModel):
    is_exception: Literal["YES", "NO"]
    resolution: Literal["RESCHEDULE", "REROUTE_TO_LOCKER", "REPLACE", "RETURN_TO_SENDER", "N/A"]
    rationale: str

class CommOutput(BaseModel):
    message: str

# 3. Define the Nodes
def resolution_agent(state: AgentState):
    llm = ChatOpenAI(model="gpt-4o")
    # In a real app, you'd fetch data from your SQLite/ChromaDB here
    prompt = f"Analyze shipment {state['shipment_id']} and provide a resolution based on logistics playbooks."
    res = llm.with_structured_output(ResolutionOutput).invoke(prompt)
    return {
        "is_exception": res.is_exception,
        "resolution": res.resolution,
        "rationale": res.rationale
    }

def communication_agent(state: AgentState):
    llm = ChatOpenAI(model="gpt-4o")
    prompt = f"Write a friendly notification for a {state['resolution']} action."
    res = llm.with_structured_output(CommOutput).invoke(prompt)
    return {"customer_message": res.message}

# 4. Build the Graph
workflow = StateGraph(AgentState)
workflow.add_node("resolver", resolution_agent)
workflow.add_node("communicator", communication_agent)

workflow.set_entry_point("resolver")
workflow.add_edge("resolver", "communicator")
workflow.add_edge("communicator", END)

app_logic = workflow.compile()

def run_delivery_process(sid):
    return app_logic.invoke({"shipment_id": sid, "escalated": False})