from typing import TypedDict, Annotated, Literal
from datetime import datetime
from operator import add

class CustomerServiceState(TypedDict):
    """
    State compartido entre todos los agents customer service.
    Production: Usar Pydantic models para validation estricta.
    """
    # Customer Context (read-only para agents)
    customer_id: str
    customer_tier: Literal["free", "pro", "enterprise"]  # Para routing prioritario
    customer_history: list[dict]  # Previous tickets, interactions

    # Current Conversation
    messages: Annotated[list[dict], add]  # Conversación actual (reducer append)
    current_intent: Literal["faq", "troubleshooting", "billing", "escalate", "unknown"]
    confidence_score: float  # Triage agent confidence (0-1)

    # Issue Tracking
    ticket_id: str
    issue_category: str  # "login", "payment", "performance", "bug", etc.
    urgency: Literal["low", "medium", "high", "critical"]
    sentiment: Literal["positive", "neutral", "negative", "angry"]

    # Agent Coordination
    current_agent: str  # Qué agent está handling now
    agents_tried: Annotated[list[str], add]  # Track agents ya invocados (evitar loops)
    escalation_reason: str  # Si escalated, por qué

    # Resolution
    resolution_attempted: bool
    resolution_successful: bool
    resolution_details: str
    handoff_to_human: bool
    human_agent_id: str  # Si escalated

    # Metadata
    created_at: datetime
    last_updated: datetime
    total_tokens_used: int  # Track cost
    total_latency_ms: int  # Track performance


def create_initial_state(customer_id: str, initial_message: str) -> CustomerServiceState:
    """
    Factory function para inicializar state cuando ticket nuevo llega.
    """
    return {
        "customer_id": customer_id,
        "customer_tier": "pro",  # Lookup en database
        "customer_history": [],  # Fetch previous tickets
        "messages": [{"role": "user", "content": initial_message}],
        "current_intent": "unknown",
        "confidence_score": 0.0,
        "ticket_id": f"TICK-{datetime.now().strftime('%Y%m%d')}-{customer_id[:6]}",
        "issue_category": "",
        "urgency": "medium",
        "sentiment": "neutral",
        "current_agent": "triage",
        "agents_tried": [],
        "escalation_reason": "",
        "resolution_attempted": False,
        "resolution_successful": False,
        "resolution_details": "",
        "handoff_to_human": False,
        "human_agent_id": "",
        "created_at": datetime.now(),
        "last_updated": datetime.now(),
        "total_tokens_used": 0,
        "total_latency_ms": 0
    }
  

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.postgres import PostgresSaver
from langchain_openai import ChatOpenAI
import os

# State definition (ver arriba)
# ...

# ===== AGENT IMPLEMENTATIONS =====
def triage_agent(state: CustomerServiceState) -> CustomerServiceState:
    """
    Triage agent: Classifica intent y urgency del customer message.
    """
    llm = ChatOpenAI(model="gpt-4-turbo", temperature=0)
    user_message = state["messages"][-1]["content"]
    customer_tier = state["customer_tier"]

    system_prompt = f"""Eres un triage specialist. Analiza el mensaje del customer y:
    1. Clasifica intent: "faq", "troubleshooting", "billing", "escalate"
    2. Detecta urgency: "low", "medium", "high", "critical"
    3. Analiza sentiment: "positive", "neutral", "negative", "angry"

    Customer tier: {customer_tier}
    Message: {user_message}

    Responde en formato JSON:
    {{"intent": "...", "urgency": "...", "sentiment": "...", "confidence": 0.0-1.0}}
    """

    response = llm.invoke([{"role": "system", "content": system_prompt}])

    # Production: Parse structured output
    # Simplified:
    result = {
        "intent": "faq",  # Parse del LLM response
        "urgency": "medium",
        "sentiment": "neutral",
        "confidence": 0.85
    }

    return {
        **state,
        "current_intent": result["intent"],
        "urgency": result["urgency"],
        "sentiment": result["sentiment"],
        "confidence_score": result["confidence"],
        "agents_tried": state["agents_tried"] + ["triage"]
    }


def faq_agent(state: CustomerServiceState) -> CustomerServiceState:
    """
    FAQ agent: RAG retrieval + answer generation.
    """
    # Production: Vector DB retrieval (Pinecone, Weaviate)
    # Simplified para demo
    user_question = state["messages"][-1]["content"]

    # Simulated RAG retrieval
    retrieved_docs = [
        "To reset your password, visit /reset-password and follow instructions.",
        "For billing questions, email billing@company.com or call +1-800-XXX-XXXX"
    ]

    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.3)
    prompt = f"""Answer the customer question using the retrieved documents.
    Be helpful and concise. Include links when relevant.

    Question: {user_question}

    Retrieved docs:
    {chr(10).join(retrieved_docs)}

    Answer:
    """

    response = llm.invoke([{"role": "user", "content": prompt}])

    return {
        **state,
        "messages": state["messages"] + [{"role": "assistant", "content": response.content}],
        "resolution_attempted": True,
        "resolution_successful": True,  # Production: Verify con follow-up
        "resolution_details": "FAQ answered via RAG retrieval",
        "agents_tried": state["agents_tried"] + ["faq"]
    }


def troubleshooting_agent(state: CustomerServiceState) -> CustomerServiceState:
    """
    Troubleshooting agent: Diagnostic workflows.
    """
    # Production: Execute diagnostic tools (health checks, log analysis)
    issue_description = state["messages"][-1]["content"]

    # Simulated diagnostic workflow
    diagnostics = {
        "system_status": "healthy",
        "user_connectivity": "OK",
        "recent_errors": []
    }

    llm = ChatOpenAI(model="gpt-4-turbo", temperature=0)
    prompt = f"""You're a troubleshooting specialist. Based on diagnostics:

    Issue: {issue_description}
    Diagnostics: {diagnostics}

    Provide step-by-step troubleshooting steps or resolution.
    """

    response = llm.invoke([{"role": "user", "content": prompt}])

    return {
        **state,
        "messages": state["messages"] + [{"role": "assistant", "content": response.content}],
        "resolution_attempted": True,
        "resolution_successful": True,  # Production: Verify resolution
        "resolution_details": "Troubleshooting workflow completed",
        "agents_tried": state["agents_tried"] + ["troubleshooting"]
    }


def escalation_agent(state: CustomerServiceState) -> CustomerServiceState:
    """
    Escalation agent: Prep handoff a humano.
    """
    # Summarize conversation para humano
    conversation_summary = "\n".join([
        f"{msg['role']}: {msg['content']}"
        for msg in state["messages"][-5:]  # Last 5 messages
    ])

    context_summary = f"""
    ESCALATED TICKET: {state['ticket_id']}
    Customer: {state['customer_id']} ({state['customer_tier']} tier)
    Issue category: {state['issue_category']}
    Urgency: {state['urgency']}
    Sentiment: {state['sentiment']}
    Agents tried: {', '.join(state['agents_tried'])}
    Escalation reason: Complex issue requiring human expertise

    Recent conversation:
    {conversation_summary}

    ACTION REQUIRED: Human agent review and response.
    """

    # Production: Send to ticket system, notify human agent
    return {
        **state,
        "handoff_to_human": True,
        "escalation_reason": "AI confidence low / complex issue",
        "messages": state["messages"] + [{
            "role": "system",
            "content": context_summary
        }],
        "agents_tried": state["agents_tried"] + ["escalation"]
    }


# ===== ROUTING LOGIC =====
def route_based_on_intent(state: CustomerServiceState) -> str:
    """
    Router function: Decide qué agent llamar basado en triage result.
    """
    intent = state["current_intent"]
    confidence = state["confidence_score"]

    # Si confidence es baja, escalar inmediatamente
    if confidence < 0.7:
        return "escalate"

    # Route basado en intent
    intent_to_agent = {
        "faq": "faq",
        "troubleshooting": "troubleshooting",
        "billing": "billing",  # Simplificado (billing agent similar a faq)
        "escalate": "escalate"
    }

    return intent_to_agent.get(intent, "escalate")


def check_resolution(state: CustomerServiceState) -> str:
    """
    Después de agent execution, verificar si resuelto o necesita escalation.
    """
    if state["handoff_to_human"]:
        return "__end__"  # Terminar workflow, humano takeover

    if state["resolution_successful"]:
        return "__end__"  # Issue resuelto

    # Si agent no pudo resolver, escalar
    return "escalate"


# ===== BUILD GRAPH =====
workflow = StateGraph(CustomerServiceState)

# Add nodes
workflow.add_node("triage", triage_agent)
workflow.add_node("faq", faq_agent)
workflow.add_node("troubleshooting", troubleshooting_agent)
workflow.add_node("billing", faq_agent)  # Simplificado: billing similar a FAQ
workflow.add_node("escalate", escalation_agent)

# Set entry point
workflow.set_entry_point("triage")

# Conditional routing desde triage
workflow.add_conditional_edges(
    "triage",
    route_based_on_intent,
    {
        "faq": "faq",
        "troubleshooting": "troubleshooting",
        "billing": "billing",
        "escalate": "escalate"
    }
)

# Después de cada agent, verificar resolution
workflow.add_conditional_edges("faq", check_resolution, {"escalate": "escalate", "__end__": END})
workflow.add_conditional_edges("troubleshooting", check_resolution, {"escalate": "escalate", "__end__": END})
workflow.add_conditional_edges("billing", check_resolution, {"escalate": "escalate", "__end__": END})
workflow.add_edge("escalate", END)

# ===== PERSISTENCE (PRODUCTION) =====
# Postgres checkpointer para persistence
DB_URI = os.getenv("POSTGRES_URI", "postgresql://user:pass@localhost:5432/langgraph")
checkpointer = PostgresSaver.from_conn_string(DB_URI)

# Compile con checkpointing
app = workflow.compile(
    checkpointer=checkpointer,
    interrupt_before=["escalate"]  # HITL: Pause antes de escalation
)

# ===== USAGE =====
if __name__ == "__main__":
    # Ejemplo: Customer submits ticket
    initial_state = create_initial_state(
        customer_id="CUST-12345",
        initial_message="I can't log in to my account. Getting 'invalid password' error."
    )

    # Execute workflow
    config = {"configurable": {"thread_id": initial_state["ticket_id"]}}
    result = app.invoke(initial_state, config)

    print("=== RESULT ===")
    print(f"Ticket: {result['ticket_id']}")
    print(f"Intent: {result['current_intent']}")
    print(f"Resolution: {result['resolution_successful']}")
    print(f"Agents used: {result['agents_tried']}")
    print(f"Handoff to human: {result['handoff_to_human']}")
    print(f"Final message: {result['messages'][-1]['content']}")
  

from typing import Annotated, TypedDict, Literal
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
  from langchain_core.messages import HumanMessage

  # Define State compartido entre agentes
  class AgentState(TypedDict):
      messages: Annotated[list, "Messages compartidos entre agentes"]
      next_agent: str  # Qué agent ejecutar siguiente
      final_answer: str  # Resultado final consolidado


  # Supervisor Agent: Decide routing
  def supervisor_node(state: AgentState) -> AgentState:
      """
      Supervisor analiza contexto y decide qué worker llamar.
      Production: Usar structured output con function calling.
      """
      llm = ChatOpenAI(model="gpt-4-turbo", temperature=0)

      system_prompt = """Eres un supervisor coordinando workers especializados.

      Workers disponibles:
      - researcher: Busca información en web/docs
      - analyzer: Analiza datos y genera insights
      - writer: Escribe reportes finales

      Basándote en el último mensaje, decide qué worker llamar.
      Si el task está completo, responde 'FINISH'.
      """

      messages = state["messages"] + [
          HumanMessage(content=system_prompt)
      ]

      response = llm.invoke(messages)

      # Parse response para decidir next agent
      content = response.content.lower()
      if "researcher" in content:
          next_agent = "researcher"
      elif "analyzer" in content:
          next_agent = "analyzer"
      elif "writer" in content:
          next_agent = "writer"
      else:
          next_agent = "FINISH"

      return {
          **state,
          "next_agent": next_agent,
          "messages": state["messages"] + [response]
      }


  # Worker Agent: Researcher
  def researcher_node(state: AgentState) -> AgentState:
      """Worker especializado en research."""
      llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.3)

      # En producción: Añadir tools (web search, database lookup, etc.)
      response = llm.invoke([
          HumanMessage(content=f"Research task: {state['messages'][-1].content}")
      ])

      return {
          **state,
          "messages": state["messages"] + [response]
      }


  # Worker Agent: Analyzer
  def analyzer_node(state: AgentState) -> AgentState:
      """Worker especializado en análisis."""
      llm = ChatOpenAI(model="gpt-4-turbo", temperature=0)

      response = llm.invoke([
          HumanMessage(content=f"Analyze findings: {state['messages'][-3].content}")
      ])

      return {
          **state,
          "messages": state["messages"] + [response]
      }


  # Conditional routing basado en supervisor decision
  def router(state: AgentState) -> Literal["researcher", "analyzer", "writer", "__end__"]:
      """Route al siguiente agent basado en supervisor decision."""
      next_agent = state.get("next_agent", "FINISH")

      if next_agent == "FINISH":
          return "__end__"

      return next_agent


  # Construir graph
  workflow = StateGraph(AgentState)

  # Add nodes
  workflow.add_node("supervisor", supervisor_node)
  workflow.add_node("researcher", researcher_node)
  workflow.add_node("analyzer", analyzer_node)
  # workflow.add_node("writer", writer_node)  # Similarmente

  # Add edges
  workflow.set_entry_point("supervisor")

  # Conditional edges desde supervisor
  workflow.add_conditional_edges(
      "supervisor",
      router,
      {
          "researcher": "researcher",
          "analyzer": "analyzer",
          "writer": "writer",
          "__end__": END
      }
  )

  # Después de cada worker, volver a supervisor
  workflow.add_edge("researcher", "supervisor")
  workflow.add_edge("analyzer", "supervisor")
  workflow.add_edge("writer", "supervisor")

  app = workflow.compile()

  # Ejecutar
  result = app.invoke({
      "messages": [HumanMessage(content="Research AI market trends 2025")],
      "next_agent": "",
      "final_answer": ""
  })
  

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph
from operator import add

# Shared State con reducers para merge updates
class CollaborationState(TypedDict):
    task: str
    findings: Annotated[list[str], add]  # Reducer: Append findings
    analysis: str
    draft_report: str


def research_agent(state: CollaborationState) -> CollaborationState:
    """Agente 1: Investiga y añade findings."""
    # Simulated research
    new_findings = [
        "AI market growing 48.6% CAGR",
        "Multi-agent adoption increased 200% in 2024"
    ]

    return {
        **state,
        "findings": new_findings  # Reducer 'add' append automáticamente
    }


def analysis_agent(state: CollaborationState) -> CollaborationState:
    """Agente 2: Analiza findings acumulados."""
    findings = state.get("findings", [])
    analysis = f"Analysis of {len(findings)} findings: " + "; ".join(findings)

    return {
        **state,
        "analysis": analysis
    }


def writer_agent(state: CollaborationState) -> CollaborationState:
    """Agente 3: Escribe report basado en analysis."""
    draft = f"Report on {state['task']}\n\n{state['analysis']}\n\nRecommendations: ..."

    return {
        **state,
        "draft_report": draft
    }


# Build graph sin supervisor
workflow = StateGraph(CollaborationState)

workflow.add_node("research", research_agent)
workflow.add_node("analyze", analysis_agent)
workflow.add_node("write", writer_agent)

# Sequential flow
workflow.set_entry_point("research")
workflow.add_edge("research", "analyze")
workflow.add_edge("analyze", "write")

app = workflow.compile()

result = app.invoke({
    "task": "AI market trends",
    "findings": [],
    "analysis": "",
    "draft_report": ""
})
  

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

class RoutingState(TypedDict):
    user_query: str
    intent: str  # "faq", "troubleshooting", "billing", "escalate"
    response: str


def intent_classifier(state: RoutingState) -> RoutingState:
    """
    Clasifica intent del usuario usando LLM con structured output.
    Production: Fine-tune classifier para intents específicos.
    """
    query = state["user_query"].lower()

    # Simplified classification (producción usar LLM)
    if "password" in query or "login" in query:
        intent = "faq"
    elif "not working" in query or "error" in query:
        intent = "troubleshooting"
    elif "charge" in query or "refund" in query:
        intent = "billing"
    else:
        intent = "escalate"

    return {**state, "intent": intent}


def faq_agent(state: RoutingState) -> RoutingState:
    """Agent especializado en FAQs con RAG sobre knowledge base."""
    # Production: RAG retrieval + response generation
    response = "FAQ response: Para resetear password, visita /reset-password"
    return {**state, "response": response}


def troubleshooting_agent(state: RoutingState) -> RoutingState:
    """Agent con diagnostic tools y workflows."""
    response = "Troubleshooting: He verificado tu cuenta. El issue es X. Solución: Y"
    return {**state, "response": response}


def billing_agent(state: RoutingState) -> RoutingState:
    """Agent con acceso a billing system."""
    response = "Billing: He revisado tu última factura. El cargo corresponde a..."
    return {**state, "response": response}


def escalation_agent(state: RoutingState) -> RoutingState:
    """Agent que escala a humano con context summary."""
    response = "Escalated to human agent. Summary: " + state["user_query"]
    return {**state, "response": response}


# Routing function
def route_by_intent(state: RoutingState) -> Literal["faq", "troubleshooting", "billing", "escalate"]:
    return state["intent"]


# Build graph
workflow = StateGraph(RoutingState)

workflow.add_node("classifier", intent_classifier)
workflow.add_node("faq", faq_agent)
workflow.add_node("troubleshooting", troubleshooting_agent)
workflow.add_node("billing", billing_agent)
workflow.add_node("escalate", escalation_agent)

workflow.set_entry_point("classifier")

# Conditional routing basado en intent
workflow.add_conditional_edges(
    "classifier",
    route_by_intent,
    {
        "faq": "faq",
        "troubleshooting": "troubleshooting",
        "billing": "billing",
        "escalate": "escalate"
    }
)

# Todos terminan después de responder
workflow.add_edge("faq", END)
workflow.add_edge("troubleshooting", END)
workflow.add_edge("billing", END)
workflow.add_edge("escalate", END)

app = workflow.compile()
  

from langgraph.checkpoint import MemorySaver
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import tools_condition

class HandoffState(TypedDict):
    messages: list
    handoff_to: str  # "tier2_agent", "billing_specialist", "human"


def tier1_agent(state: HandoffState) -> HandoffState:
    """
    Tier 1 agent intenta resolver. Si no puede, handoff.
    """
    user_message = state["messages"][-1]

    # Simplified logic (producción usar LLM decision)
    if "refund" in user_message.lower():
        # Tier 1 no puede procesar refunds
        return {
            **state,
            "handoff_to": "billing_specialist",
            "messages": state["messages"] + [
                {"role": "assistant", "content": "Transferring to billing specialist..."}
            ]
        }
    elif "technical issue" in user_message.lower():
        return {
            **state,
            "handoff_to": "tier2_agent",
            "messages": state["messages"] + [
                {"role": "assistant", "content": "Escalating to technical support..."}
            ]
        }
    else:
        # Tier 1 puede resolver
        return {
            **state,
            "handoff_to": "RESOLVED",
            "messages": state["messages"] + [
                {"role": "assistant", "content": "Here's the solution..."}
            ]
        }


def tier2_agent(state: HandoffState) -> HandoffState:
    """Tier 2 agent con más tools y knowledge."""
    # Intenta resolver issue técnico
    # Si aún no puede, handoff a humano
    return {
        **state,
        "handoff_to": "human",
        "messages": state["messages"] + [
            {"role": "assistant", "content": "Complex issue. Connecting to human expert..."}
        ]
    }


def human_handoff_node(state: HandoffState) -> HandoffState:
    """
    HITL: Pausa workflow hasta que humano intervenga.
    Production: Interrupt() here, resume cuando human responda.
    """
    # Prep context para humano
    summary = "Context summary: " + "\n".join([m.get("content", "") for m in state["messages"][-5:]])

    return {
        **state,
        "messages": state["messages"] + [
            {"role": "system", "content": f"HUMAN TAKEOVER\n\n{summary}"}
        ]
    }


# Routing
def handoff_router(state: HandoffState):
    handoff = state.get("handoff_to", "RESOLVED")

    if handoff == "RESOLVED":
        return END

    return handoff


# Build graph
workflow = StateGraph(HandoffState)

workflow.add_node("tier1", tier1_agent)
workflow.add_node("tier2_agent", tier2_agent)
workflow.add_node("billing_specialist", lambda s: s)  # Simplified
workflow.add_node("human", human_handoff_node)

workflow.set_entry_point("tier1")

workflow.add_conditional_edges(
    "tier1",
    handoff_router,
    {
        "tier2_agent": "tier2_agent",
        "billing_specialist": "billing_specialist",
        "human": "human",
        END: END
    }
)

workflow.add_conditional_edges("tier2_agent", handoff_router)

# Checkpointing para HITL
memory = MemorySaver()
app = workflow.compile(checkpointer=memory, interrupt_before=["human"])