voltmind.py

"""
VOLTMIND - Parent AI Controller
The central brain. Wires VehicleAgent + SchedulerAgent together,
calls the Claude API for intelligent decisions, and exposes a
FastAPI server the frontend talks to.

Run with:
    pip install fastapi uvicorn anthropic python-dotenv
    uvicorn voltmind:app --reload --port 8000
"""

import json
import os
from datetime import datetime, timedelta
from typing import Optional

import anthropic
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from vehicle_agent import VehicleAgent, UserTier, ChargingStatus
from scheduler_agent import SchedulerAgent


# ─── FastAPI app ───────────────────────────────────────────────────────────────

app = FastAPI(title="VOLTMIND API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],   # tighten this in production
    allow_methods=["*"],
    allow_headers=["*"],
)


# ─── Request / response models (what the frontend sends) ──────────────────────

class ConnectRequest(BaseModel):
    vehicle_id:     str
    current_soc:    float           # battery % now
    target_soc:     float           # battery % wanted
    max_charge_kw:  float           # charger hardware limit
    hours_until_departure: float    # e.g. 6.5
    user_tier:      str = "STANDARD"  # STANDARD | PRIORITY | EMERGENCY


class EmergencyRequest(BaseModel):
    vehicle_id: str
    reason:     str = "emergency override"


# ─── VOLTMIND core class ───────────────────────────────────────────────────────

class VoltMind:
    """
    The Parent AI. Holds all active VehicleAgents, owns the
    SchedulerAgent, and calls Claude to make intelligent decisions.
    """

    SYSTEM_PROMPT = """
You are VOLTMIND, an AI controller for smart EV charging infrastructure.
You receive the current state of all connected vehicles and their charging
schedule, then return a JSON decision object.

Your goals (in priority order):
  1. Emergency vehicles must always charge immediately at full rate
  2. Keep total grid load below 80% of capacity at all times
  3. Ensure every vehicle reaches its target SoC before departure
  4. Minimise cost by preferring off-peak hours when possible

You must respond ONLY with a valid JSON object in this exact format:
{
  "assessment": "one sentence summary of the current grid situation",
  "risk_level": "LOW | MEDIUM | HIGH | CRITICAL",
  "actions": [
    {
      "agent_id": "<id>",
      "vehicle_id": "<id>",
      "action": "CHARGE_FULL | THROTTLE | PAUSE | EMERGENCY_OVERRIDE",
      "recommended_kw": <number>,
      "reason": "<short reason>"
    }
  ],
  "grid_health": {
    "current_load_kw": <number>,
    "capacity_kw": <number>,
    "load_percent": <number>,
    "recommendation": "<one sentence>"
  },
  "alerts": ["<alert string>", ...]
}

Be decisive. Never leave a vehicle without a clear action.
"""

    def __init__(self, total_grid_capacity_kw: float = 100.0):
        self.grid_capacity_kw = total_grid_capacity_kw
        self.agents:    dict[str, VehicleAgent]  = {}   # vehicle_id → agent
        self.scheduler: SchedulerAgent           = SchedulerAgent(total_grid_capacity_kw)
        self.client:    anthropic.Anthropic      = anthropic.Anthropic()
        self.decision_log: list[dict]            = []
        print(f"[VOLTMIND] Initialised. Grid capacity: {total_grid_capacity_kw} kW")

    # ── Vehicle lifecycle ──────────────────────────────────────────────────────

    def connect_vehicle(
        self,
        vehicle_id:            str,
        current_soc:           float,
        target_soc:            float,
        max_charge_kw:         float,
        hours_until_departure: float,
        user_tier:             UserTier = UserTier.STANDARD,
    ) -> dict:
        """Called when a vehicle plugs in."""
        if vehicle_id in self.agents:
            raise ValueError(f"{vehicle_id} is already connected.")

        departure = datetime.now() + timedelta(hours=hours_until_departure)
        agent = VehicleAgent(
            vehicle_id     = vehicle_id,
            current_soc    = current_soc,
            target_soc     = target_soc,
            max_charge_kw  = max_charge_kw,
            departure_time = departure,
            user_tier      = user_tier,
        )
        self.agents[vehicle_id] = agent
        print(f"[VOLTMIND] {vehicle_id} connected. Urgency: {agent.urgency_score}")

        # Get AI decision for this new state
        decision = self.think()
        self._apply_decision(decision)
        return decision

    def disconnect_vehicle(self, vehicle_id: str) -> dict:
        """Called when a vehicle unplugs."""
        if vehicle_id not in self.agents:
            raise ValueError(f"{vehicle_id} is not connected.")

        summary = self.agents[vehicle_id].disconnect()
        del self.agents[vehicle_id]
        print(f"[VOLTMIND] {vehicle_id} disconnected.")

        decision = self.think()
        self._apply_decision(decision)
        return {"disconnected": summary, "new_schedule": decision}

    def emergency_override(self, vehicle_id: str, reason: str) -> dict:
        """Instantly reprioritises one vehicle — bypasses normal scheduling."""
        if vehicle_id not in self.agents:
            raise ValueError(f"{vehicle_id} not found.")

        agent = self.agents[vehicle_id]
        agent.user_tier = UserTier.EMERGENCY
        agent.urgency_score = 100.0
        print(f"[VOLTMIND] EMERGENCY OVERRIDE → {vehicle_id}: {reason}")

        decision = self.think(context=f"EMERGENCY: {vehicle_id} needs immediate full charge. Reason: {reason}")
        self._apply_decision(decision)
        return decision

    # ── The AI brain: call Claude ──────────────────────────────────────────────

    def think(self, context: str = "") -> dict:
        """
        Builds a snapshot of current grid state, sends it to Claude,
        and gets back a structured JSON decision.
        This is where VOLTMIND becomes truly intelligent.
        """
        if not self.agents:
            return {
                "assessment": "No vehicles connected. Grid idle.",
                "risk_level": "LOW",
                "actions": [],
                "grid_health": {
                    "current_load_kw": 0,
                    "capacity_kw": self.grid_capacity_kw,
                    "load_percent": 0,
                    "recommendation": "System ready."
                },
                "alerts": []
            }

        # Build schedule from Scheduler Agent
        requests = [a.build_charging_request() for a in self.agents.values()]
        schedule = self.scheduler.build_schedule(requests)
        payload  = self.scheduler.to_voltmind_payload(schedule)

        # Current total load
        current_load = sum(
            a.current_charge_kw for a in self.agents.values()
        )

        # Build the prompt for Claude
        prompt = f"""
Current grid state:
- Total capacity: {self.grid_capacity_kw} kW
- Current load: {current_load:.1f} kW
- Load %: {(current_load / self.grid_capacity_kw * 100):.1f}%
- Connected vehicles: {len(self.agents)}
- Time: {datetime.now().strftime("%H:%M")}

{"Context: " + context if context else ""}

Vehicle states:
{json.dumps([a.to_dict() for a in self.agents.values()], indent=2)}

Proposed schedule from Scheduler Agent:
{json.dumps(payload, indent=2)}

Analyse this situation and return your JSON decision.
"""

        print(f"[VOLTMIND] Calling Claude API...")

        try:
            response = self.client.messages.create(
                model      = "claude-sonnet-4-20250514",
                max_tokens = 1000,
                system     = self.SYSTEM_PROMPT,
                messages   = [{"role": "user", "content": prompt}],
            )
            raw = response.content[0].text.strip()

            # Strip markdown fences if present
            if raw.startswith("```"):
                raw = raw.split("```")[1]
                if raw.startswith("json"):
                    raw = raw[4:]

            decision = json.loads(raw)
            self.decision_log.append({
                "time":     datetime.now().isoformat(),
                "decision": decision,
            })
            print(f"[VOLTMIND] Decision: {decision.get('assessment','')}")
            return decision

        except Exception as e:
            print(f"[VOLTMIND] Claude API error: {e}")
            # Fallback: use scheduler output directly without AI reasoning
            return self._fallback_decision(payload)

    # ── Apply Claude's decision to all agents ──────────────────────────────────

    def _apply_decision(self, decision: dict) -> None:
        """
        Reads Claude's action list and calls apply_throttle()
        on each VehicleAgent with the recommended kW.
        """
        for action in decision.get("actions", []):
            vid = action.get("vehicle_id")
            kw  = action.get("recommended_kw", 0)
            if vid and vid in self.agents:
                self.agents[vid].apply_throttle(kw)
                print(f"[VOLTMIND] Applied: {vid} → {kw} kW "
                      f"({action.get('action','')}) — {action.get('reason','')}")

    # ── Status snapshot ────────────────────────────────────────────────────────

    def get_status(self) -> dict:
        """Full system status — what the dashboard polls."""
        current_load = sum(a.current_charge_kw for a in self.agents.values())
        return {
            "timestamp":        datetime.now().isoformat(),
            "grid_capacity_kw": self.grid_capacity_kw,
            "current_load_kw":  round(current_load, 2),
            "load_percent":     round(current_load / self.grid_capacity_kw * 100, 1),
            "vehicles_connected": len(self.agents),
            "vehicles": [a.to_dict() for a in self.agents.values()],
            "last_decision": self.decision_log[-1] if self.decision_log else None,
        }

    # ── Fallback (if Claude API unavailable) ──────────────────────────────────

    def _fallback_decision(self, payload: dict) -> dict:
        """Rule-based fallback used if Claude API call fails."""
        actions = []
        for v in payload.get("vehicles", []):
            kw = v.get("slots", [{}])[0].get("charge_kw", 0) if v.get("slots") else 0
            actions.append({
                "agent_id":      v["agent_id"],
                "vehicle_id":    v["vehicle_id"],
                "action":        "CHARGE_FULL" if kw > 0 else "PAUSE",
                "recommended_kw": kw,
                "reason":        "fallback rule-based decision",
            })
        current_load = sum(a.get("recommended_kw", 0) for a in actions)
        return {
            "assessment": "Fallback mode: rule-based scheduling active.",
            "risk_level": "MEDIUM",
            "actions":    actions,
            "grid_health": {
                "current_load_kw": current_load,
                "capacity_kw":     self.grid_capacity_kw,
                "load_percent":    round(current_load / self.grid_capacity_kw * 100, 1),
                "recommendation":  "Restore Claude API connection for full AI control.",
            },
            "alerts": ["Claude API unavailable — running in fallback mode."],
        }


# ─── Singleton instance ────────────────────────────────────────────────────────

voltmind = VoltMind(total_grid_capacity_kw=100.0)


# ─── FastAPI routes ────────────────────────────────────────────────────────────

@app.get("/")
def root():
    return {"system": "VOLTMIND", "status": "online"}


@app.get("/status")
def status():
    """Dashboard polls this every 5 seconds."""
    return voltmind.get_status()


@app.post("/vehicle/connect")
def connect(req: ConnectRequest):
    """EV plugs in → register with VOLTMIND."""
    try:
        tier = UserTier[req.user_tier.upper()]
        return voltmind.connect_vehicle(
            vehicle_id            = req.vehicle_id,
            current_soc           = req.current_soc,
            target_soc            = req.target_soc,
            max_charge_kw         = req.max_charge_kw,
            hours_until_departure = req.hours_until_departure,
            user_tier             = tier,
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/vehicle/disconnect/{vehicle_id}")
def disconnect(vehicle_id: str):
    """EV unplugs → remove from schedule."""
    try:
        return voltmind.disconnect_vehicle(vehicle_id)
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


@app.post("/emergency")
def emergency(req: EmergencyRequest):
    """Emergency override — instantly reprioritises a vehicle."""
    try:
        return voltmind.emergency_override(req.vehicle_id, req.reason)
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


@app.post("/think")
def think():
    """Manually trigger a VOLTMIND re-evaluation (useful for demos)."""
    decision = voltmind.think()
    voltmind._apply_decision(decision)
    return decision


# ─── Local demo (no server needed) ────────────────────────────────────────────

if __name__ == "__main__":

    print("=" * 60)
    print("  VOLTMIND — Full System Demo (no server)")
    print("  Set ANTHROPIC_API_KEY env var to use Claude AI")
    print("=" * 60)

    vm = VoltMind(total_grid_capacity_kw=100.0)

    # Connect 3 vehicles
    print("\n--- Step 1: Vehicles connect ---\n")
    vm.connect_vehicle("EV-001", current_soc=12.0, target_soc=80.0,
                       max_charge_kw=11.0, hours_until_departure=6.0)

    vm.connect_vehicle("EV-002", current_soc=60.0, target_soc=90.0,
                       max_charge_kw=7.4,  hours_until_departure=10.0)

    vm.connect_vehicle("AMBULANCE-01", current_soc=35.0, target_soc=100.0,
                       max_charge_kw=22.0, hours_until_departure=2.0,
                       user_tier=UserTier.EMERGENCY)

    # Check status
    print("\n--- Step 2: System status ---\n")
    status = vm.get_status()
    print(json.dumps({
        "load_percent":       status["load_percent"],
        "vehicles_connected": status["vehicles_connected"],
        "current_load_kw":    status["current_load_kw"],
    }, indent=2))

    # Simulate emergency override
    print("\n--- Step 3: Emergency override on EV-001 ---\n")
    vm.emergency_override("EV-001", "driver reports medical emergency")

    # EV-002 disconnects
    print("\n--- Step 4: EV-002 disconnects ---\n")
    vm.disconnect_vehicle("EV-002")

    print("\n--- Final status ---\n")
    final = vm.get_status()
    for v in final["vehicles"]:
        print(f"  {v['vehicle_id']:15s}  {v['current_charge_kw']:5.1f} kW  "
              f"SoC={v['current_soc']}%  status={v['status']}")

    print(f"\n  Total load: {final['current_load_kw']} kW "
          f"({final['load_percent']}% of grid)")
    print("\n[VOLTMIND] System nominal.")