demeter

import os
import json
import numpy as np
from langchain_openai import AzureChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.graph import StateGraph, END
from agent.tools.actuation import convert_targets_to_actions

from agent.Marl.bandit import ContextualBandit
from agent.Marl.strategies import STRATEGIES, NUM_ACTIONS
from agent.Qdrant.Store import store_fmu
from agent.sub_agents.water_and_atmospheric_dependencies.physics_engine import predict_outcome

# 🛡️ GUARDRAILS
from agent.guardrails.validation import validate_plan, detect_hard_violations, create_validation_report

from dotenv import load_dotenv

load_dotenv()

# --- NEW TOOLS DEFINITION ---
def check_cross_domain_conflicts(atmos, water):
    conflicts = []
    
    # 1. Thermal Shock Check
    air_t = atmos.get('air_temp', 25)
    water_t = water.get('water_temp', 20)
    if abs(air_t - water_t) > 10:
        conflicts.append(f"CRITICAL: Thermal Shock Risk. Air ({air_t}C) and Water ({water_t}C) delta > 10C.")

    # 2. Transpiration vs Uptake Check
    # High VPD (Dry) + High EC (Salty) = Burn Risk
    rh = atmos.get('humidity', 60)
    ec = water.get('ec', 1.0)
    if rh < 50 and ec > 2.0:
        conflicts.append(f"STRESS: Low Humidity ({rh}%) + High EC ({ec}) will cause Tip Burn.")

    return conflicts

def validate_hard_limits(plan):
    """
    Validates plan against hard limits using guardrails.
    Returns list of violations.
    """
    has_violations, violations = detect_hard_violations(plan)
    
    if has_violations:
        print(f"🚫 HARD LIMIT VIOLATIONS DETECTED:")
        for v in violations:
            print(f"   {v}")
    
    # Also check for obvious physics conflicts
    additional_conflicts = []
    
    if plan.get('air_temp', 25) - plan.get('water_temp', 20) > 10:
        additional_conflicts.append("⚠️ Thermal Shock Risk: Air/Water temp delta > 10°C")
    
    if plan.get('humidity', 60) < 40 and plan.get('ec', 1.0) > 2.5:
        additional_conflicts.append("⚠️ Burn Risk: Low humidity + high EC")
    
    if plan.get('humidity', 60) > 85:
        additional_conflicts.append("🚫 Mold Risk: Humidity > 85%")
    
    return violations + additional_conflicts

# --- STATE DEFINITION ---
from typing import TypedDict, Optional, Dict, Any, List

class SupervisorState(TypedDict):
    # Inputs
    atmos_plan: Dict[str, Any]
    water_plan: Dict[str, Any]
    strategy_advice: str  # Kept as advice, not law
    
    # Processing
    merged_plan: Dict[str, Any]
    review_notes: List[str]
    simulation_health: float
    
    # Output
    final_decision: str # "APPROVE" or "REJECT"
    critique: str       # Feedback for sub-agents if Rejected

API_KEY = os.environ.get("AZURE_OPENAI_API_KEY")
ENDPOINT = os.environ.get("AZURE_OPENAI_ENDPOINT")
DEPLOYMENT_NAME = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME", "gpt-4.1")
API_VERSION = os.environ.get("AZURE_OPENAI_API_VERSION", "2024-12-01-preview")

class SupervisorAgent:
    def __init__(self, researcher_agent=None):
        self.name = "Supervisor"
        self.bandit = ContextualBandit(n_actions=NUM_ACTIONS, feature_dim=519)
        
        if API_KEY and ENDPOINT:
            self.model = AzureChatOpenAI(
                azure_endpoint=ENDPOINT,
                api_key=API_KEY,
                api_version=API_VERSION,
                deployment_name=DEPLOYMENT_NAME,
                temperature=0.0 # Zero temp for strict judging
            )
        
        self.app = self._build_graph()

    def _build_graph(self):
        workflow = StateGraph(SupervisorState)

        # 1. Merge: Combine the two JSONs
        workflow.add_node("merge", self.node_merge)
        
        # 2. Review: Run the 3 Tools (Conflicts, Limits, Physics)
        workflow.add_node("review", self.node_review)
        
        # 3. Judge: LLM decides if the issues are fatal
        workflow.add_node("judge", self.node_judge)

        # Flow
        workflow.set_entry_point("merge")
        workflow.add_edge("merge", "review")
        workflow.add_edge("review", "judge")
        workflow.add_edge("judge", END)
        
        return workflow.compile()

    # --- NODE FUNCTIONS ---

    def node_merge(self, state):
        print("   🔗 Supervisor Merging Plans...")
        # Simple dictionary merge
        merged = {**state['atmos_plan'], **state['water_plan']}
        return {"merged_plan": merged}

    def node_review(self, state):
        print("   🔍 Supervisor Running Unit Tests...")
        plan = state['merged_plan']
        notes = []

        # Tool 1: Conflict Check
        conflicts = check_cross_domain_conflicts(state['atmos_plan'], state['water_plan'])
        if conflicts:
            notes.extend(conflicts)
            
        # Tool 2: Limit Check
        limits = validate_hard_limits(plan)
        if limits:
            notes.extend(limits)
            
        # Tool 3: Physics Simulator
         # Comparing plan vs itself as a snapshot for now
        health = 100
        
        if health < 90:
            notes.append(f"SIMULATION FAIL: Predicted health drops to {health}%. ")

        return {"review_notes": notes, "simulation_health": health}

    def node_judge(self, state):
        """
        The LLM looks at the automated test results and makes the final call.
        """
        print("   ⚖️ Supervisor Judging...")
        
        if not state['review_notes']:
            # No issues found by tools
            return {"final_decision": "APPROVE", "critique": "Plan looks solid."}
        
        # If issues exist, ask LLM if they are fatal or acceptable trade-offs
        prompt = f"""
        You are the Quality Assurance Supervisor.
        
        PROPOSED PLAN: {state['merged_plan']}
        
        AUTOMATED TEST FAILURES:
        {json.dumps(state['review_notes'], indent=2)}
        
        ADVISORY STRATEGY: {state['strategy_advice']}
        
        TASK:
        1. BIAS: You should almost always APPROVE.
        2. ONLY 'REJECT' if the plan is physically impossible or immediately fatal (e.g., pH < 3.0, Water Temp > 40°C).
        3. IGNORE 'Simulation Fail' warnings if the Strategy justifies the extreme values (e.g., 'Flush' requires low EC).
        4. Treat "Risk" warnings as acceptable trade-offs for the strategy.
        OUTPUT JSON: {{ "verdict": "APPROVE" or "REJECT", "critique": "Explanation..." }}
        """

        print("Supervisor Prompt:\n", prompt)
        
        try:
            response = self.model.invoke([HumanMessage(content=prompt)])
            content = response.content.replace("```json", "").replace("```", "").strip()
            result = json.loads(content)
            
            return {
                "final_decision": result.get("verdict", "REJECT"),
                "critique": result.get("critique", "Automated tests failed.")
            }
        except:
            # Default to reject if unsafe
            return {"final_decision": "REJECT", "critique": "Plan failed automated safety checks."}




    # --- ENTRY POINT ---

    def synthesize_plan(self, atmos_plan, water_plan, fmu, history, strategy_info):
        strategy_name, _, action_idx = strategy_info
        
        initial_state = {
            "atmos_plan": atmos_plan,
            "water_plan": water_plan,
            "strategy_advice": strategy_name,
            "merged_plan": {},
            "review_notes": [],
            "simulation_health": 0.0,
            "final_decision": "",
            "critique": ""
        }
        
        result = self.app.invoke(initial_state)
        final_targets = result.get("merged_plan", {})

        current_sensors = fmu.metadata.get('sensors', {})
        
        # 🛡️ GUARDRAIL CHECK: Validate before executing
        print(f"[{self.name}] 🛡️ Running guardrail validation...")
        validation = validate_plan(final_targets)
        
        if validation["severity"] == "CRITICAL":
            print(create_validation_report(final_targets))
            print(f"[{self.name}] ⚠️ CRITICAL VIOLATIONS - Clamping to bounds...")
            final_targets = validation["bounded_plan"]
        
        if validation["warnings"]:
            print(f"[{self.name}] ⚠️ Warnings: {', '.join(validation['warnings'])}")
        
        print(f"[{self.name}] ⚙️ Converting Targets to Actuator Commands...")

        sensor_vals = [
            float(current_sensors.get("pH", 0.0)),
            float(current_sensors.get("EC", 0.0)),
            float(current_sensors.get("temp", 0.0)),
            float(current_sensors.get("humidity", 0.0))
        ]

        if hasattr(fmu, 'vector') and len(fmu.vector) == 512:
            if isinstance(fmu.vector, list):
                fmu.vector.extend(sensor_vals)
            else:
                import numpy as np
                fmu.vector = np.concatenate((fmu.vector, sensor_vals)).tolist()
        
        # Calculate physical actions
        physical_action_obj = convert_targets_to_actions(current_sensors, final_targets)
        
        # Convert Pydantic model to Dict for JSON serialization
        final_payload = physical_action_obj.dict()
        
        # Log it
        print(f"[{self.name}] 🚜 Activating Hardware: {final_payload}")
        
        # Store in FMU
        fmu.metadata["action_taken"] = str(final_payload)
        fmu.metadata["bandit_action_id"] = action_idx
        fmu.metadata["strategic_intent"] = strategy_name
        
        if "image_b64" in fmu.metadata: del fmu.metadata["image_b64"]
        store_fmu(fmu)
        
        return final_payload

    # --- ADVISORY ONLY (Not Enforced) ---
    def get_strategic_goal(self, fmu):
        # (Same as before, but treated as advice now)
        sensors = fmu.metadata.get('sensors', {})
        fmu_vector = fmu.vector
        vis_vec1 = np.array(fmu_vector) if isinstance(fmu_vector, list) else fmu_vector
        vis_vec = vis_vec1[:512] if len(vis_vec1) >= 512 else None
        if vis_vec is None or len(vis_vec) == 0: vis_vec = np.zeros(516)
        
        s_vec = np.array([
            (float(sensors.get('pH', 6.0)) - 6.0) / 2.0, 
            float(sensors.get('EC', 1.0)) / 3.0,
            float(sensors.get('temp', 25.0)) / 40.0
        ])
        context_vector = np.concatenate([vis_vec, s_vec])

      #  print("Context Vector for Bandit:", context_vector.shape)
        
        action_idx, _ = self.bandit.select_action(context_vector)
        strategy_name = STRATEGIES[action_idx]
        
        return strategy_name, "Advisory Only", int(action_idx)

    def learn_from_outcome(self, fmu, outcome_info):
        """
        Bandit Learning: Update the model based on action outcome.
        
        Args:
            fmu: The FMU object containing metadata about the previous action
            outcome_info: Either:
                - Current plant health (0-100) from simulator, OR
                - Reward score (-1.0 to 1.0) from judge
        """
        # Retrieve the action that was taken in the previous cycle
        prev_action_idx = fmu.metadata.get("bandit_action_id")
        if prev_action_idx is None:
            return  # No previous action to learn from
        
        prev_action_idx = int(prev_action_idx)
        
        # Build the context vector (same as get_strategic_goal)
        sensors = fmu.metadata.get('sensors', {})
        fmu_vector = fmu.vector
        vis_vec1 = np.array(fmu_vector) if isinstance(fmu_vector, list) else fmu_vector
        vis_vec = vis_vec1[:512] if len(vis_vec1) >= 512 else None
        if vis_vec is None or len(vis_vec) == 0: vis_vec = np.zeros(512)
        
        s_vec = np.array([
            (float(sensors.get('pH', 6.0)) - 6.0) / 2.0, 
            float(sensors.get('EC', 1.0)) / 3.0,
            float(sensors.get('temp', 25.0)) / 40.0
        ])
        context_vector = np.concatenate([vis_vec, s_vec])
        
        # Handle reward: Can be health (0-100) or judge reward (-1 to 1)
        if isinstance(outcome_info, dict):
            reward = float(outcome_info.get("reward", 0.0))
        else:
            # Assume it's health (0-100), convert to normalized reward
            health_val = float(outcome_info)
            if health_val >= 85:
                reward = 1.0  # Excellent
            elif health_val >= 70:
                reward = health_val / 100.0  # Good
            elif health_val >= 50:
                reward = (health_val / 100.0) * 0.5  # Mediocre
            else:
                reward = -1.0  # Terrible
        
        # Update the bandit model with this outcome
        self.bandit.update(context_vector, prev_action_idx, reward)
        
        strategy_name = STRATEGIES.get(prev_action_idx, "UNKNOWN")
        print(f"[{self.name}] 🧠 Bandit Learning: {strategy_name} (Action {prev_action_idx}) → Reward {reward:.2f}")
        
        # Save the updated model
        self.bandit.save()