+--------------------------------------------------------+

| 1. STRATO DI INTERFACCIA ED ESPOSIZIONE (User / Voice) |
+--------------------------------------------------------+
                           │  (Pipeline Vocale / Chat)
                           ▼
+--------------------------------------------------------+

| 2. STRATO DI ORCHESTRAZIONE COGNITIVA (LangGraph AI)  |
+--------------------------------------------------------+
                           │  (Richiesta Dati / Proposte di Modifica via MCP)
                           ▼
+--------------------------------------------------------+

| 3. STRATO DI VALIDAZIONE E CONTROLLO (Il Middleware)   |
+--------------------------------------------------------+
                           │  (Dati Validati / Comandi Sicuri)
                           ▼
+--------------------------------------------------------+

| 4. STRATO DI OTTIMIZZAZIONE ALGORITMICA (EMHASS)       |
+--------------------------------------------------------+
                           │  (Stato Sensori / Attuazioni Hardware)
                           ▼
+--------------------------------------------------------+

| 5. STRATO DI INTEGRAZIONE FISICA (Home Assistant Core) |
+--------------------------------------------------------+

   [START] ➔ ( 1. NODO_PARSER )
                    │
                    ▼
           ( 2. NODO_MCP_GET )
                    │
                    ▼
         ( 3. NODO_CORE_REASONER ) ◄─── ( Loop di conversazione se mancano dati )
                    │
             [Router Condizionale]
             /                 \
    (Se propone modifiche)    (Se chiede solo info)
           /                     \
          ▼                       ▼
  ( 4. NODO_VALIDATOR )     ( 6. NODO_ANSWER_GEN ) ➔ [END]
          │                       ▲
   [Router di Esito]             │
     /           \                │
 (Approvato)   (Rifiutato)        │
   /               \              │
  ▼                 ▼             │
(5. NODO_EMHASS)  (Ri-ragionamento)
  │                        │
  └────────────────────────┘

import os
from typing import Annotated, Dict, Any, Literal
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
# Nota: Si assume l'uso del backend di messaggistica standard di LangGraph
# e l'integrazione con un LLM compatibile con le chiamate a funzioni/tool (es. OpenAI o Anthropic)

# ==========================================
# 1. DEFINIZIONE DELLO STATO DEL GRAFO
# ==========================================
class EnergyManagerState(TypedDict):
    # Elenco dei messaggi della chat (gestito automaticamente da add_messages)
    messages: Annotated[list, add_messages]
    # Intento rilevato (es. "modifica_piano", "richiesta_info")
    current_intent: str
    # Dati energetici correnti prelevati tramite ha-mcp
    energy_context: Dict[str, Any]
    # La modifica al piano energetico proposta dal LLM
    proposed_plan_mod: Dict[str, Any]
    # Risultato del controllo del Validator deterministico
    validation_result: Dict[str, Any]


# ==========================================
# 2. IMPLEMENTAZIONE DEI NODI (FUNZIONI)
# ==========================================

def nodo_parser(state: EnergyManagerState) -> Dict[str, Any]:
    """Analizza l'input dell'utente per estrarre l'intento."""
    last_message = state["messages"][-1].content.lower()
    
    # Logica di parsing (semplificata per il PoC, implementabile con un mini-LLM o regex)
    intent = "richiesta_info"
    if any(keyword in last_message for keyword in ["sposta", "ritarda", "anticipa", "cambia", "carica"]):
        intent = "modifica_piano"
        
    return {"current_intent": intent}


def nodo_mcp_get(state: EnergyManagerState) -> Dict[str, Any]:
    """Invocazione fittizia dei tool ha-mcp per scattare lo snapshot dell'impianto."""
    # Nella realtà qui chiameresti i tool esposti da ha-mcp verso Home Assistant
    # Esempio: mcp_client.call_tool("get_emhass_plan", {})
    mock_context = {
        "production_forecast_ok": True,
        "current_soc": 65, # % batteria
        "flexible_loads": {
            "lavatrice": {"scheduled_start": "11:00", "duration_hours": 2, "power_kw": 2.0},
            "auto_elettrica": {"scheduled_start": "22:00", "duration_hours": 4, "power_kw": 3.7}
        }
    }
    return {"energy_context": mock_context}


def nodo_core_reasoner(state: EnergyManagerState) -> Dict[str, Any]:
    """L'LLM valuta la richiesta alla luce dei dati MCP e propone una strategia."""
    last_user_message = state["messages"][-1].content
    context = state["energy_context"]
    intent = state["current_intent"]
    
    proposed_mod = {}
    
    if intent == "modifica_piano":
        # Qui il modello decide una proposta da sottoporre al validator.
        # Esempio: l'utente ha chiesto di ritardare la lavatrice.
        proposed_mod = {
            "load_name": "lavatrice",
            "action": "reschedule",
            "new_start": "14:00"
        }
    
    return {"proposed_plan_mod": proposed_mod}


def nodo_validator(state: EnergyManagerState) -> Dict[str, Any]:
    """Strato 3: Il Middleware deterministico in Python puro (Zero AI).

    Valida la fisica e i vincoli impianto.
    """
    proposal = state["proposed_plan_mod"]
    context = state["energy_context"]
    
    # Logica di validazione fittizia:
    # Se lo spostamento porta il carico in una fascia oraria con sole calante (es. dopo le 16:00)
    # e il SoC è basso, rifiuta per proteggere la batteria.
    if not proposal:
        return {"validation_result": {"approved": True}}
        
    new_hour = int(proposal["new_start"].split(":")[0])
    
    if new_hour >= 16 and context["current_soc"] < 70:
        return {
            "validation_result": {
                "approved": False,
                "reason": "La batteria è al 65% e dopo le 16:00 la produzione solare cala. "
                          "Rischieresti di scaricare l'accumulo prima di sera."
            }
        }
        
    return {"validation_result": {"approved": True, "reason": "Ottimizzazione sicura."}}


def nodo_emhass_commit(state: EnergyManagerState) -> Dict[str, Any]:
    """Applica fisicamente la modifica inviandola a EMHASS tramite servizio HA (via MCP)."""
    proposal = state["proposed_plan_mod"]
    # Qui invocheresti il tool MCP di scrittura:
    # mcp_client.call_tool("set_emhass_load_window", {"load": proposal["load_name"], "time": proposal["new_start"]})
    print(f"[EMHASS COMMIT] Schedulazione aggiornata per {proposal['load_name']} alle ore {proposal['new_start']}")
    return {}


def nodo_answer_gen(state: EnergyManagerState) -> Dict[str, Any]:
    """Genera la risposta finale (testo per la pipeline vocale TTS di HA)."""
    intent = state["current_intent"]
    validation = state.get("validation_result", {"approved": True})
    proposal = state.get("proposed_plan_mod", {})
    
    if intent == "richiesta_info":
        response = "Oggi l'obiettivo di autoconsumo è confermato al 92%. C'è un surplus di energia previsto tra le 13:00 e le 15:00."
    else:
        if validation["approved"]:
            response = f"Ottima idea! Ho ricalcolato il piano energetico spostando la {proposal['load_name']} alle {proposal['new_start']}. L'autoconsumo resta ottimizzato."
        else:
            response = f"Non posso spostare la {proposal['load_name']} a quell'ora. {validation['reason']}"
            
    # Creazione del messaggio di risposta dell'AI (Formato LangGraph standard)
    from langchain_core.messages import AIMessage
    return {"messages": [AIMessage(content=response)]}


# ==========================================
# 3. LOGICA DI ROUTING CONDIZIONALE
# ==========================================

def route_dopo_reasoner(state: EnergyManagerState) -> Literal["nodo_validator", "nodo_answer_gen"]:
    """Decide se andare al controllo di sicurezza o rispondere direttamente."""
    if state["current_intent"] == "modifica_piano":
        return "nodo_validator"
    return "nodo_answer_gen"


def route_dopo_validator(state: EnergyManagerState) -> Literal["nodo_emhass_commit", "nodo_answer_gen"]:
    """Se approvato scrive su EMHASS, altrimenti va direttamente alla risposta di rifiuto."""
    if state["validation_result"]["approved"]:
        return "nodo_emhass_commit"
    return "nodo_answer_gen"


# ==========================================
# 4. COSTRUZIONE E COMPILAZIONE DEL GRAFO
# ==========================================

workflow = StateGraph(EnergyManagerState)

# Aggiunta dei Nodi
workflow.add_node("nodo_parser", nodo_parser)
workflow.add_node("nodo_mcp_get", nodo_mcp_get)
workflow.add_node("nodo_core_reasoner", nodo_core_reasoner)
workflow.add_node("nodo_validator", nodo_validator)
workflow.add_node("nodo_emhass_commit", nodo_emhass_commit)
workflow.add_node("nodo_answer_gen", nodo_answer_gen)

# Definizione dei Collegamenti (Edges)
workflow.add_edge(START, "nodo_parser")
workflow.add_edge("nodo_parser", "nodo_mcp_get")
workflow.add_edge("nodo_mcp_get", "nodo_core_reasoner")

# Router dopo il ragionamento dell'LLM
workflow.add_conditional_edges(
    "nodo_core_reasoner",
    route_dopo_reasoner
)

# Router dopo la validazione del Middleware
workflow.add_conditional_edges(
    "nodo_validator",
    route_dopo_validator
)

# Chiusura dei flussi verso la generazione risposta
workflow.add_edge("nodo_emhass_commit", "nodo_answer_gen")
workflow.add_edge("nodo_answer_gen", END)

# Compilazione dell'applicazione finale
app = workflow.compile()

# ==========================================
# 5. ESEMPIO DI ESECUZIONE (TEST)
# ==========================================
if __name__ == "__main__":
    from langchain_core.messages import HumanMessage
    
    # Simuliamo un input utente critico (rifiutato dal validator)
    inputs = {"messages": [HumanMessage(content="Sposta la lavatrice stasera tardi verso le 18")]}
    config = {"configurable": {"thread_id": "casa_utente_1"}}
    
    print("--- Avvio esecuzione Grafo ---")
    for output in app.stream(inputs, config):
        for key, value in output.items():
            print(f"\n[Nodo Eseguito]: {key}")
            if "messages" in value:
                print(f"Risposta parziale: {value['messages'][-1].content}")

import asyncio
from typing import List, Dict, Any
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import ToolNode

# =====================================================================
# 1. CONFIGURAZIONE CONNESSIONE AL SERVER MCP (ha-mcp)
# =====================================================================
# Definiamo come l'agente spawna o si connette al server ha-mcp.
# Se ha-mcp è un container separato, si userà un trasporto di tipo SSE (HTTP).
# Se viene eseguito in locale via CLI, si usa Stdio. Qui simuliamo lo Stdio.
server_params = StdioServerParameters(
    command="npx",
    args=["-y", "@mcp-server/homeassistant"] # Sostituire con il pacchetto specifico ha-mcp se custom
)

# =====================================================================
# 2. DEFINIZIONE DEI TOOL ATTRAVERSO IL PROTOCOLLO MCP
# =====================================================================
# Creiamo dei wrapper LangChain attorno alle chiamate di sessione MCP.
# Nota: In un'implementazione reale, puoi automatizzare questo mapping 
# ciclando su session.list_tools() fornito dal server MCP.

class MCPToolWrapper:
    def __init__(self, session: ClientSession):
        self.session = session

    def get_tools(self) -> List[Any]:
        
        @tool
        async def get_energy_status() -> str:
            """Recupera lo stato energetico in tempo reale della casa: SoC batteria, produzione FV attuale e consumi."""
            # Effettua la chiamata standardizzata MCP verso il server di Home Assistant
            result = await self.session.call_tool("get_states", arguments={"entity_ids": ["sensor.battery_soc", "sensor.power_production", "sensor.power_consumption"]})
            return str(result.content)

        @tool
        async def get_emhass_plan() -> str:
            """Recupera il piano di ottimizzazione orario corrente generato da EMHASS per i carichi differibili."""
            result = await self.session.call_tool("get_emhass_production_plan", arguments={})
            return str(result.content)

        @tool
        async def update_emhass_schedule(load_name: str, requested_start_time: str) -> str:
            """Invia una richiesta di riprogrammazione oraria per un carico differibile. 
            Passa prima dal Validator prima di essere consolidata su EMHASS.
            Args:
                load_name: Il nome del carico (es. 'lavatrice', 'lavastoviglie').
                requested_start_time: L'orario richiesto in formato HH:MM (es. '14:30').
            """
            # Chiama il servizio di Home Assistant esposto tramite MCP che punta al nostro Middleware/Validator
            result = await self.session.call_tool(
                "call_service", 
                arguments={
                    "domain": "energy_manager",
                    "service": "validate_and_schedule",
                    "service_data": {
                        "load": load_name,
                        "time": requested_start_time
                    }
                }
            )
            return str(result.content)

        return [get_energy_status, get_emhass_plan, update_emhass_schedule]

# =====================================================================
# 3. INTEGRAZIONE DEI TOOL MCP NEL RAGIONAMENTO DEL GRAFO
# =====================================================================
async def main():
    # Inizializziamo il client MCP
    async with stdio_client(server_params) as (read_stream, write_stream):
        async with ClientSession(read_stream, write_stream) as session:
            # Inizializza la sessione protocollo
            await session.initialize()
            
            # Estraiamo i tool mappati su MCP
            mcp_tools = MCPToolWrapper(session).get_tools()
            
            # Creiamo il ToolNode nativo di LangGraph per gestire le esecuzioni dei tool
            tool_node = ToolNode(mcp_tools)
            
            # Configuriamo il modello LLM agganciandogli i tool MCP
            model = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(mcp_tools)
            
            # --- ESEMPIO DI COINVOLGIMENTO DEI TOOL IN UN NODO ---
            # Questo frammento simula ciò che accade dentro il NODO_CORE_REASONER
            system_prompt = SystemMessage(
                content="Sei l'Energy Manager vocale della casa. Hai accesso ai tool MCP per controllare "
                        "lo stato dell'impianto, i piani EMHASS e richiedere spostamenti di carico. "
                        "Rispondi in modo conciso e focalizzato sull'autoconsumo."
            )
            user_message = HumanMessage(content="Posso spostare la lavatrice alle 14:00 oggi?")
            
            print(f"User: {user_message.content}\n")
            print("[LangGraph Core Reasoner] L'agente valuta la richiesta e decide quali tool MCP invocare...")
            
            # Primo step di ragionamento dell'LLM
            response = await model.ainvoke([system_prompt, user_message])
            
            # Se l'LLM decide che ha bisogno di dati, genererà delle 'tool_calls'
            if response.tool_calls:
                for tool_call in response.tool_calls:
                    print(f"➔ L'AI decide di chiamare il Tool MCP: '{tool_call['name']}' con argomenti: {tool_call['args']}")
                    
                # Il grafo passerebbe questi comandi al ToolNode per l'esecuzione reale su Home Assistant
                # Es: tool_output = await tool_node.ainvoke({"messages": [response]})
            else:
                print(f"AI: {response.content}")

# Esecuzione dell'ambiente asincrono (necessario per MCP)
if __name__ == "__main__":
    asyncio.run(main())