"""
Exception Agent — Miler stall detection + general exception handling.

Stall detection flow:
  JetStream TRACKING/miler.location.updated (pull) -> Redis HSET -> stall check -> miler.stalled
  JetStream TRACKING/miler.stalled (pull)          -> POST reassign + POST notify

StallDetector sweep (every 60 s):
  Postgres pickupbookings (active status) -> Redis movement.updated_at age -> miler.stalled
"""
import asyncio
import json
import os
import uuid
from datetime import datetime, timedelta, timezone
from typing import Dict, List, Any, Optional, Set
from dataclasses import dataclass, field
from enum import Enum

import asyncpg
import nats
import nats.errors
import redis.asyncio as aioredis

from core.agent import SpecializedAgent
from core.types import AgentTask, MessageType
from core.logger import logger
from core.http_client import api_post
from core.llm import decide_stall_response, build_stall_context, StallDecision
from config.system_config import (
    GO_API_BASE_URL, INTERNAL_API_KEY,
    DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD,
    NATS_HOST, NATS_PORT, NATS_USER, NATS_PASSWORD,
    REDIS_HOST, REDIS_PORT, REDIS_PASSWORD,
)

STALL_MINUTES    = 10
ACTIVE_STATUSES  = ["Miler_Assigned", "Pickup_Scheduled"]
TRACKING_STREAM  = "TRACKING"

# Reassignment is customer-visible and hard to reverse, so it is gated: the LLM
# only *proposes* a reassign unless the agent is explicitly put in autonomous
# mode AND the model is confident enough. Otherwise the proposal is escalated to
# a human via JARVIS.
AUTONOMOUS_REASSIGN     = os.getenv("EXCEPTION_AGENT_AUTONOMOUS", "false").lower() == "true"
REASSIGN_MIN_CONFIDENCE = float(os.getenv("EXCEPTION_AGENT_REASSIGN_CONFIDENCE", "0.7"))
# How long a booking stays "already alerted" so we don't re-alert the same stall.
# Self-expiring so dedup state is bounded and survives restarts (unlike a set).
STALL_NOTIFIED_TTL      = int(os.getenv("STALL_NOTIFIED_TTL_SEC", "21600"))  # 6h


def _utcnow() -> datetime:
    """Timezone-aware current time in UTC. All stall math uses this so it stays
    correct regardless of the container's local timezone, and never mixes naive
    with tz-aware timestamps."""
    return datetime.now(timezone.utc)


def _parse_ts(value: Optional[str]) -> Optional[datetime]:
    """Parse an ISO timestamp as tz-aware UTC. Naive values are assumed UTC so
    pre-existing (naive) Redis entries remain comparable during rollout."""
    if not value:
        return None
    try:
        dt = datetime.fromisoformat(value)
    except (ValueError, TypeError):
        return None
    return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)


class ExceptionType(str, Enum):
    DELAY                = "delay"
    CANCELLATION         = "cancellation"
    RESCHEDULE           = "reschedule"
    REROUTE              = "reroute"
    FAILED_DELIVERY      = "failed_delivery"
    VEHICLE_BREAKDOWN    = "vehicle_breakdown"
    HUB_OVERFLOW         = "hub_overflow"
    WEATHER_IMPACT       = "weather_impact"
    CUSTOMER_UNAVAILABLE = "customer_unavailable"
    MILER_STALLED        = "miler_stalled"


class ExceptionSeverity(str, Enum):
    LOW      = "low"
    MEDIUM   = "medium"
    HIGH     = "high"
    CRITICAL = "critical"


@dataclass
class ExceptionRecord:
    exception_id: str
    order_id: str
    exception_type: ExceptionType
    severity: ExceptionSeverity
    description: str
    detected_at: datetime
    resolved_at: Optional[datetime]
    resolution: Optional[str]
    assigned_to: Optional[str]
    status: str
    actions_taken: List[str] = field(default_factory=list)


class ExceptionAgent(SpecializedAgent):
    """
    Handles miler stall detection via JetStream pull-subscribe + Postgres + Redis,
    and all existing order exception types.
    """

    def __init__(self):
        super().__init__(
            agent_id="EXCEPTION_AGENT",
            domain="exception_handling",
            description="Stall detection, exception resolution, reassignment",
        )

        self._nc: Optional[nats.aio.client.Client] = None
        self._js = None
        self._pg: Optional[asyncpg.Pool] = None
        self._redis = aioredis.Redis(
            host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD,
            decode_responses=True,
        )

        self._location_sub = None
        self._stalled_sub  = None
        self._exceptions: Dict[str, ExceptionRecord] = {}
        self._strategies       = self._init_strategies()
        self._escalation_rules = self._init_escalation_rules()

    # ── Startup ───────────────────────────────────────────────────────────────

    async def start(self):
        await self._connect_infra()
        await self._setup_nats_pull_subs()
        asyncio.create_task(self._pull_location_loop())
        asyncio.create_task(self._pull_stalled_loop())
        asyncio.create_task(self._stall_detector_loop())
        await super().start()

    async def _connect_infra(self):
        try:
            self._nc = await nats.connect(
                servers=[f"nats://{NATS_HOST}:{NATS_PORT}"],
                user=NATS_USER,
                password=NATS_PASSWORD,
                max_reconnect_attempts=10,
            )
            self._js = self._nc.jetstream()
            logger.info("EXCEPTION_AGENT connected to NATS JetStream")
        except Exception as e:
            logger.warning(f"EXCEPTION_AGENT NATS connect failed: {e}")

        try:
            self._pg = await asyncpg.create_pool(
                host=DB_HOST, port=DB_PORT, database=DB_NAME,
                user=DB_USER, password=DB_PASSWORD,
                min_size=1, max_size=5,
            )
            logger.info("EXCEPTION_AGENT connected to Postgres")
        except Exception as e:
            logger.warning(f"EXCEPTION_AGENT Postgres connect failed: {e}")

    async def _setup_nats_pull_subs(self):
        if not self._js:
            logger.warning("EXCEPTION_AGENT: skipping NATS pull subs (no JetStream connection)")
            return

        try:
            self._location_sub = await self._js.pull_subscribe(
                "miler.location.updated",
                "tracking_location_consumer",
                stream=TRACKING_STREAM,
            )
            logger.info("EXCEPTION_AGENT pull-subscribed: miler.location.updated")
        except Exception as e:
            logger.warning(f"pull_subscribe miler.location.updated failed: {e}")

        try:
            self._stalled_sub = await self._js.pull_subscribe(
                "miler.stalled",
                "tracking_stall_consumer",
                stream=TRACKING_STREAM,
            )
            logger.info("EXCEPTION_AGENT pull-subscribed: miler.stalled")
        except Exception as e:
            logger.warning(f"pull_subscribe miler.stalled failed: {e}")

    # ── Pull loops ────────────────────────────────────────────────────────────

    async def _pull_location_loop(self):
        while self._running:
            if not self._location_sub:
                await asyncio.sleep(5)
                continue
            try:
                msgs = await self._location_sub.fetch(batch=20, timeout=1.0)
                for msg in msgs:
                    try:
                        await self._on_location_update(msg)
                    except Exception as e:
                        logger.error(f"Location handler error: {e}")
                    finally:
                        await msg.ack()
            except nats.errors.TimeoutError:
                pass
            except Exception as e:
                logger.error(f"Location pull loop error: {e}")
                await asyncio.sleep(2)

    async def _pull_stalled_loop(self):
        while self._running:
            if not self._stalled_sub:
                await asyncio.sleep(5)
                continue
            try:
                msgs = await self._stalled_sub.fetch(batch=10, timeout=1.0)
                for msg in msgs:
                    try:
                        await self._on_miler_stalled(msg)
                    except Exception as e:
                        logger.error(f"Stalled handler error: {e}")
                    finally:
                        await msg.ack()
            except nats.errors.TimeoutError:
                pass
            except Exception as e:
                logger.error(f"Stalled pull loop error: {e}")
                await asyncio.sleep(2)

    # ── Stall detection: per location ping ───────────────────────────────────

    async def _on_location_update(self, msg):
        try:
            data = json.loads(msg.data.decode())
        except Exception:
            return

        miler_id  = str(data.get("miler_id", ""))
        lat       = str(round(float(data.get("lat", 0)), 6))
        lon       = str(round(float(data.get("lon", 0)), 6))
        now_ts    = _utcnow().isoformat()
        redis_key = f"miler:{miler_id}:movement"

        prev     = await self._redis.hgetall(redis_key)
        prev_lat = prev.get("lat", "")
        prev_lon = prev.get("lon", "")

        if lat != prev_lat or lon != prev_lon:
            await self._redis.hset(redis_key, mapping={
                "lat": lat, "lon": lon,
                "updated_at": now_ts,
                "position_unchanged_since": now_ts,
            })
        else:
            await self._redis.hset(redis_key, mapping={"lat": lat, "lon": lon, "updated_at": now_ts})

            unchanged_since = _parse_ts(prev.get("position_unchanged_since")) or _utcnow()
            minutes_stalled = (_utcnow() - unchanged_since).total_seconds() / 60

            if minutes_stalled >= STALL_MINUTES:
                booking = await self._get_active_booking(miler_id)
                if booking:
                    await self._publish_stall(miler_id, booking["booking_id"], minutes_stalled)

    # ── Stall detection: background sweep ────────────────────────────────────

    async def _stall_detector_loop(self):
        while self._running:
            await asyncio.sleep(60)
            try:
                await self._sweep_active_bookings()
            except Exception as e:
                logger.error(f"StallDetector error: {e}")

    async def _sweep_active_bookings(self):
        if not self._pg:
            return

        logger.debug("StallDetector: sweeping active bookings")

        try:
            async with self._pg.acquire() as conn:
                rows = await conn.fetch(
                    """
                    SELECT DISTINCT
                        assignedmileruserid::text AS miler_id,
                        bookingid::text           AS booking_id
                    FROM pickupbookings
                    WHERE status = ANY($1)
                    """,
                    ACTIVE_STATUSES,
                )
        except Exception as e:
            logger.error(f"StallDetector Postgres error: {e}")
            return

        now = _utcnow()
        for row in rows:
            miler_id   = row["miler_id"]
            booking_id = row["booking_id"]

            movement = await self._redis.hgetall(f"miler:{miler_id}:movement")
            if not movement or "updated_at" not in movement:
                continue

            updated_at = _parse_ts(movement.get("updated_at"))
            if updated_at is None:
                continue
            minutes_stale = (now - updated_at).total_seconds() / 60

            if minutes_stale >= STALL_MINUTES:
                logger.warning(f"StallDetector: miler {miler_id} stale {minutes_stale:.1f} min (booking {booking_id})")
                await self._publish_stall(miler_id, booking_id, minutes_stale)

    # ── Stall: publish + handle ───────────────────────────────────────────────

    def _stall_counter_key(self, miler_id) -> str:
        return f"miler:{miler_id}:stalls:{_utcnow().strftime('%Y-%m-%d')}"

    def _stall_notified_key(self, booking_id) -> str:
        return f"stall_notified:{booking_id}"

    def _stall_handled_key(self, booking_id) -> str:
        return f"stall_handled:{booking_id}"

    async def _claim_stall(self, booking_id) -> bool:
        """Atomically claim the first stall ALERT for a booking (Redis SET NX with
        TTL). True = publish this alert; False = already alerted within the window.
        Replaces the old unbounded in-memory set: self-expiring (bounded memory),
        durable across restarts, and race-free between the ping and sweep paths.
        Fail-open on Redis error so a genuine stall is never silently muted."""
        try:
            got = await self._redis.set(
                self._stall_notified_key(booking_id), "1", ex=STALL_NOTIFIED_TTL, nx=True,
            )
            return bool(got)
        except Exception as e:
            logger.warning(f"stall claim failed for {booking_id}: {e}")
            return True

    async def _claim_stall_handled(self, booking_id) -> bool:
        """Atomically claim HANDLING of a stall event (consumer side). Guards
        against JetStream at-least-once redelivery re-running the decision and
        repeating an irreversible reassign. Fail-open on Redis error."""
        try:
            got = await self._redis.set(
                self._stall_handled_key(booking_id), "1", ex=STALL_NOTIFIED_TTL, nx=True,
            )
            return bool(got)
        except Exception as e:
            logger.warning(f"stall handled-claim failed for {booking_id}: {e}")
            return True

    async def _incr_stall_counter(self, miler_id):
        """Bump this miler's per-day stall count (Redis, best-effort). 48h TTL
        covers the day rollover; failures never block stall handling."""
        try:
            key = self._stall_counter_key(miler_id)
            await self._redis.incr(key)
            await self._redis.expire(key, 172800)
        except Exception as e:
            logger.warning(f"stall counter incr failed for miler {miler_id}: {e}")

    async def _publish_stall(self, miler_id: str, booking_id: str, minutes_stalled: float):
        if not await self._claim_stall(booking_id):
            return  # already alerted for this booking within the dedup window
        await self._incr_stall_counter(miler_id)
        payload = json.dumps({
            "miler_id":        miler_id,
            "booking_id":      booking_id,
            "minutes_stalled": round(minutes_stalled, 1),
        }).encode()
        try:
            await self._js.publish("miler.stalled", payload)
            logger.info(f"miler.stalled published: miler={miler_id} booking={booking_id} ({minutes_stalled:.1f} min)")
        except Exception as e:
            logger.error(f"Failed to publish miler.stalled: {e}")

    async def _on_miler_stalled(self, msg):
        """Decide (via Claude) how to handle a stalled miler, then act.

        Claude chooses one of wait / notify_only / reassign / escalate. The only
        irreversible, customer-visible action (reassign) is gated behind an
        autonomy flag and a confidence threshold; otherwise it is escalated to a
        human via JARVIS. If the LLM is unavailable, we fall back to the previous
        deterministic behaviour (reassign + notify) so detection never silently
        stops acting.
        """
        try:
            data = json.loads(msg.data.decode())
        except Exception:
            return

        miler_id        = data.get("miler_id", "")
        booking_id      = data.get("booking_id", "")
        minutes_stalled = data.get("minutes_stalled", 0)

        logger.warning(f"EXCEPTION_AGENT stall: miler={miler_id} booking={booking_id} ({minutes_stalled} min)")

        if not await self._claim_stall_handled(booking_id):
            logger.info(f"Stall for booking {booking_id} already handled; skipping redelivery")
            return

        facts = await self._gather_stall_facts(miler_id, booking_id)
        if facts:
            logger.info(f"[EXCEPTION] booking={booking_id} context facts: {facts}")

        context = build_stall_context(
            minutes_stalled, stall_threshold_min=STALL_MINUTES, now=_utcnow(), facts=facts,
        )
        decision = await decide_stall_response(context)

        if decision is None:
            logger.warning(f"No LLM decision for booking {booking_id}; falling back to reassign + notify")
            await self._reassign(booking_id, "miler_stalled")
            await self._notify_customer(booking_id, "We detected a delay, finding you a new miler")
            self._record_stall_exception(
                miler_id, booking_id, minutes_stalled,
                resolution="Fallback (LLM unavailable): reassignment triggered + customer notified",
                actions=["reassign", "notify_customer"],
            )
            return

        logger.info(
            f"[EXCEPTION] booking={booking_id} decision={decision.action} "
            f"confidence={decision.confidence:.2f} reasoning={decision.reasoning!r}"
        )

        actions: List[str] = []

        if decision.action == "wait":
            actions.append("wait")

        elif decision.action == "notify_only":
            await self._notify_customer(
                booking_id, "We're keeping an eye on your delivery — thanks for your patience."
            )
            actions.append("notify_customer")

        elif decision.action == "reassign":
            if AUTONOMOUS_REASSIGN and decision.confidence >= REASSIGN_MIN_CONFIDENCE:
                await self._reassign(booking_id, "miler_stalled")
                await self._notify_customer(booking_id, "We detected a delay, finding you a new miler")
                actions += ["reassign", "notify_customer"]
            else:
                # Not authorised to auto-reassign (mode off or low confidence):
                # hand the proposal to a human rather than taking the action.
                await self._escalate_to_human(miler_id, booking_id, minutes_stalled, decision)
                actions.append("escalate_reassign_for_review")

        else:  # "escalate"
            await self._escalate_to_human(miler_id, booking_id, minutes_stalled, decision)
            actions.append("escalate")

        self._record_stall_exception(
            miler_id, booking_id, minutes_stalled,
            resolution=f"[{decision.action}] {decision.reasoning}",
            actions=actions,
        )

    # ── Stall: action helpers ─────────────────────────────────────────────────

    async def _reassign(self, booking_id: str, reason: str) -> bool:
        result = await api_post(
            f"{GO_API_BASE_URL}/api/v1/internal/bookings/{booking_id}/reassign",
            json={"reason": reason},
            headers={"X-Internal-Key": INTERNAL_API_KEY},
        )
        ok = result is not None
        logger.info(f"Reassign {'OK' if ok else 'FAILED'} for booking {booking_id}")
        return ok

    async def _notify_customer(self, booking_id: str, message: str) -> bool:
        result = await api_post(
            f"{GO_API_BASE_URL}/api/v1/internal/notify",
            json={"booking_id": booking_id, "message": message, "target": "customer"},
            headers={"X-Internal-Key": INTERNAL_API_KEY},
        )
        ok = result is not None
        logger.info(f"Notify {'OK' if ok else 'FAILED'} for booking {booking_id}")
        return ok

    async def _escalate_to_human(self, miler_id, booking_id, minutes_stalled, decision: StallDecision):
        logger.warning(
            f"[EXCEPTION] Escalating booking {booking_id} to human dispatcher "
            f"(proposed={decision.action}, confidence={decision.confidence:.2f}): {decision.reasoning}"
        )
        await self.send_message(
            recipient="JARVIS",
            message_type=MessageType.EXCEPTION_DETECTED,
            payload={
                "order_id": booking_id,
                "miler_id": miler_id,
                "exception_type": ExceptionType.MILER_STALLED.value,
                "severity": ExceptionSeverity.HIGH.value,
                "proposed_action": decision.action,
                "confidence": decision.confidence,
                "reasoning": decision.reasoning,
                "minutes_stalled": minutes_stalled,
                "urgency": "high",
            },
        )

    async def _gather_stall_facts(self, miler_id, booking_id) -> Dict[str, Any]:
        """Read-only context gathering: pull real signals from Postgres + Redis
        to inform the decision. Best-effort — any failure is logged and skipped,
        never raised, so a thin-context decision still runs (which skews safe)."""
        facts: Dict[str, Any] = {}
        now = _utcnow()

        # Booking phase + age (Postgres, read-only).
        if self._pg:
            try:
                async with self._pg.acquire() as conn:
                    row = await conn.fetchrow(
                        "SELECT status, createdat FROM pickupbookings WHERE bookingid::text = $1",
                        str(booking_id),
                    )
                if row:
                    facts["booking_status"] = row["status"]
                    created = row["createdat"]
                    if created is not None:
                        if created.tzinfo is None:
                            created = created.replace(tzinfo=timezone.utc)
                        facts["minutes_since_booking_created"] = round((now - created).total_seconds() / 60, 1)
            except Exception as e:
                logger.warning(f"stall facts: booking query failed for {booking_id}: {e}")

        # Movement freshness (Redis, read-only).
        try:
            movement = await self._redis.hgetall(f"miler:{miler_id}:movement")
        except Exception as e:
            logger.warning(f"stall facts: redis read failed for miler {miler_id}: {e}")
            movement = {}

        if movement:
            last_ping = _parse_ts(movement.get("updated_at"))
            if last_ping:
                facts["last_gps_ping_minutes_ago"] = round((now - last_ping).total_seconds() / 60, 1)
            unchanged = _parse_ts(movement.get("position_unchanged_since"))
            if unchanged:
                facts["position_unchanged_minutes"] = round((now - unchanged).total_seconds() / 60, 1)

        # Per-miler stall count today (Redis, read-only). Already includes the
        # current stall, since _publish_stall increments before the consumer runs.
        try:
            raw = await self._redis.get(self._stall_counter_key(miler_id))
            if raw is not None:
                facts["stalls_today"] = int(raw)
        except Exception as e:
            logger.warning(f"stall facts: stall counter read failed for miler {miler_id}: {e}")

        return facts

    def _record_stall_exception(self, miler_id, booking_id, minutes_stalled, resolution, actions):
        escalated = any(a.startswith("escalate") for a in actions)
        exc_id = f"EXC-{_utcnow().strftime('%Y%m%d')}-{uuid.uuid4().hex[:6].upper()}"
        self._exceptions[exc_id] = ExceptionRecord(
            exception_id=exc_id, order_id=booking_id,
            exception_type=ExceptionType.MILER_STALLED,
            severity=ExceptionSeverity.HIGH,
            description=f"Miler {miler_id} stalled {minutes_stalled} min",
            detected_at=_utcnow(),
            resolved_at=None if escalated else _utcnow(),
            resolution=resolution,
            assigned_to=None,
            status="escalated" if escalated else "resolved",
            actions_taken=list(actions),
        )

    # ── Postgres helper ───────────────────────────────────────────────────────

    async def _get_active_booking(self, miler_id: str) -> Optional[Dict]:
        if not self._pg:
            return None
        try:
            async with self._pg.acquire() as conn:
                row = await conn.fetchrow(
                    """
                    SELECT bookingid::text AS booking_id, status
                    FROM pickupbookings
                    WHERE assignedmileruserid::text = $1
                      AND status = ANY($2)
                    ORDER BY createdat DESC
                    LIMIT 1
                    """,
                    miler_id,
                    ACTIVE_STATUSES,
                )
            if row:
                return {"booking_id": row["booking_id"], "status": row["status"]}
        except Exception as e:
            logger.error(f"Postgres get_active_booking error: {e}")
        return None

    # ── Task handler ──────────────────────────────────────────────────────────

    async def handle_task(self, task: AgentTask) -> Dict[str, Any]:
        handlers = {
            "detect_exception":      self._detect_exception,
            "analyze_exception":     self._analyze_exception,
            "resolve_exception":     self._resolve_exception,
            "cancel_order":          self._cancel_order,
            "reschedule_delivery":   self._reschedule_delivery,
            "handle_delay":          self._handle_delay,
            "get_exception_status":  self._get_exception_status,
            "get_exception_history": self._get_exception_history,
            "escalate_exception":    self._escalate_exception,
        }
        handler = handlers.get(task.task_type, self._unknown_task)
        return await handler(task)

    async def _detect_exception(self, task: AgentTask) -> Dict[str, Any]:
        order_id       = task.data.get("order_id")
        exception_type = task.data.get("exception_type")
        severity       = task.data.get("severity", "medium")
        description    = task.data.get("description", "")

        logger.warning(f"Exception Agent: {exception_type} for {order_id} (severity={severity})")

        exc_id = f"EXC-{datetime.now().strftime('%Y%m%d')}-{uuid.uuid4().hex[:6].upper()}"
        record = ExceptionRecord(
            exception_id=exc_id, order_id=order_id,
            exception_type=ExceptionType(exception_type),
            severity=ExceptionSeverity(severity), description=description,
            detected_at=datetime.now(), resolved_at=None,
            resolution=None, assigned_to=None, status="detected",
        )
        self._exceptions[exc_id] = record

        if severity in ("high", "critical"):
            await self._analyze_exception(AgentTask(
                task_id=f"{exc_id}_analyze", agent_type="exception",
                task_type="analyze_exception", data={"exception_id": exc_id},
            ))

        await self.send_message(
            recipient="JARVIS", message_type=MessageType.EXCEPTION_DETECTED,
            payload={"exception_id": exc_id, "order_id": order_id,
                     "exception_type": exception_type, "severity": severity},
        )

        return {"status": "detected", "exception_id": exc_id,
                "exception_type": exception_type, "severity": severity}

    async def _analyze_exception(self, task: AgentTask) -> Dict[str, Any]:
        exc_id = task.data.get("exception_id")
        if exc_id not in self._exceptions:
            return {"status": "error", "message": f"Exception {exc_id} not found"}
        record = self._exceptions[exc_id]
        record.status     = "analyzing"
        strategy          = self._strategies.get(record.exception_type, {})
        auto_actions      = strategy.get("auto_actions", [])
        record.actions_taken.append("analyzed")
        return {"status": "analyzed", "exception_id": exc_id, "strategy": auto_actions}

    async def _resolve_exception(self, task: AgentTask) -> Dict[str, Any]:
        exc_id = task.data.get("exception_id")
        if exc_id not in self._exceptions:
            return {"status": "error", "message": f"Exception {exc_id} not found"}
        record = self._exceptions[exc_id]
        record.status     = "resolving"
        for action in self._strategies.get(record.exception_type, {}).get("auto_actions", []):
            record.actions_taken.append(action)
        record.status      = "resolved"
        record.resolved_at = datetime.now()
        record.resolution  = f"Completed {len(record.actions_taken)} actions"

        await self.send_message(
            recipient="CUSTOMER_AGENT", message_type=MessageType.NOTIFICATION_SENT,
            payload={"order_id": record.order_id, "message_type": "exception_resolved",
                     "data": {"exception_type": record.exception_type.value}},
        )
        return {"status": "resolved", "exception_id": exc_id, "resolution": record.resolution}

    async def _cancel_order(self, task: AgentTask) -> Dict[str, Any]:
        order_id = task.data.get("order_id")
        reason   = task.data.get("reason", "Customer requested cancellation")
        exc_id   = f"EXC-{datetime.now().strftime('%Y%m%d')}-{uuid.uuid4().hex[:6].upper()}"
        self._exceptions[exc_id] = ExceptionRecord(
            exception_id=exc_id, order_id=order_id,
            exception_type=ExceptionType.CANCELLATION, severity=ExceptionSeverity.MEDIUM,
            description=reason, detected_at=datetime.now(), resolved_at=datetime.now(),
            resolution="Order cancelled", assigned_to=None, status="resolved",
            actions_taken=["release_vehicle", "notify_customer"],
        )
        await self.send_message(
            recipient="FLEET_AGENT", message_type=MessageType.AGENT_TASK,
            payload={"task_type": "release_vehicle_for_cancel", "order_id": order_id},
        )
        await self.send_message(
            recipient="CUSTOMER_AGENT", message_type=MessageType.ORDER_CANCELLED,
            payload={"order_id": order_id, "reason": reason,
                     "refund_info": "Refund within 5-7 business days"},
        )
        return {"status": "cancelled", "exception_id": exc_id, "order_id": order_id}

    async def _reschedule_delivery(self, task: AgentTask) -> Dict[str, Any]:
        order_id = task.data.get("order_id")
        new_date = task.data.get("new_date")
        reason   = task.data.get("reason", "Operational constraints")
        exc_id   = f"EXC-{datetime.now().strftime('%Y%m%d')}-{uuid.uuid4().hex[:6].upper()}"
        self._exceptions[exc_id] = ExceptionRecord(
            exception_id=exc_id, order_id=order_id,
            exception_type=ExceptionType.RESCHEDULE, severity=ExceptionSeverity.MEDIUM,
            description=f"Rescheduled to {new_date}: {reason}",
            detected_at=datetime.now(), resolved_at=None, resolution=None,
            assigned_to=None, status="resolving",
        )
        await self.send_message(
            recipient="CUSTOMER_AGENT", message_type=MessageType.NOTIFICATION_SENT,
            payload={"order_id": order_id, "message_type": "rescheduled",
                     "data": {"new_date": new_date, "reason": reason}},
        )
        return {"status": "rescheduled", "exception_id": exc_id, "order_id": order_id, "new_date": new_date}

    async def _handle_delay(self, task: AgentTask) -> Dict[str, Any]:
        order_id     = task.data.get("order_id")
        delay_reason = task.data.get("reason", "Unknown")
        new_eta      = task.data.get("new_eta")
        severity     = task.data.get("severity", "medium")
        exc_id       = f"EXC-{datetime.now().strftime('%Y%m%d')}-{uuid.uuid4().hex[:6].upper()}"
        self._exceptions[exc_id] = ExceptionRecord(
            exception_id=exc_id, order_id=order_id,
            exception_type=ExceptionType.DELAY, severity=ExceptionSeverity(severity),
            description=f"Delay: {delay_reason}", detected_at=datetime.now(),
            resolved_at=None, resolution=None, assigned_to=None, status="resolving",
        )
        if new_eta:
            await self.send_message(
                recipient="DISPATCH_AGENT", message_type=MessageType.ORDER_DELAYED,
                payload={"order_id": order_id, "new_eta": new_eta, "reason": delay_reason},
            )
        await self.send_message(
            recipient="CUSTOMER_AGENT", message_type=MessageType.NOTIFICATION_SENT,
            payload={"order_id": order_id, "message_type": "delayed",
                     "data": {"delay_reason": delay_reason, "new_eta": new_eta}},
        )
        return {"status": "delay_handled", "exception_id": exc_id,
                "order_id": order_id, "new_eta": new_eta, "customer_notified": True}

    async def _get_exception_status(self, task: AgentTask) -> Dict[str, Any]:
        exc_id = task.data.get("exception_id")
        if exc_id not in self._exceptions:
            return {"status": "error", "message": f"Exception {exc_id} not found"}
        r = self._exceptions[exc_id]
        return {
            "exception_id": exc_id, "order_id": r.order_id,
            "type": r.exception_type.value, "severity": r.severity.value,
            "status": r.status, "description": r.description,
            "detected_at": r.detected_at.isoformat(),
            "resolved_at": r.resolved_at.isoformat() if r.resolved_at else None,
            "actions_taken": r.actions_taken,
        }

    async def _get_exception_history(self, task: AgentTask) -> Dict[str, Any]:
        exc_type = task.data.get("exception_type")
        status_f = task.data.get("status")
        items    = list(self._exceptions.values())
        if exc_type:
            items = [e for e in items if e.exception_type.value == exc_type]
        if status_f:
            items = [e for e in items if e.status == status_f]
        return {
            "total_exceptions": len(items),
            "exceptions": [
                {"exception_id": e.exception_id, "order_id": e.order_id,
                 "type": e.exception_type.value, "severity": e.severity.value,
                 "status": e.status, "detected_at": e.detected_at.isoformat()}
                for e in items[-50:]
            ],
        }

    async def _escalate_exception(self, task: AgentTask) -> Dict[str, Any]:
        exc_id = task.data.get("exception_id")
        if exc_id not in self._exceptions:
            return {"status": "error", "message": f"Exception {exc_id} not found"}
        record      = self._exceptions[exc_id]
        escalate_to = self._escalation_rules.get(record.exception_type, {}).get("escalate_to", "JARVIS")
        record.actions_taken.append(f"escalated_to_{escalate_to}")
        await self.send_message(
            recipient=escalate_to, message_type=MessageType.EXCEPTION_DETECTED,
            payload={"exception_id": exc_id, "order_id": record.order_id,
                     "exception_type": record.exception_type.value,
                     "severity": record.severity.value, "urgency": "high"},
        )
        return {"status": "escalated", "exception_id": exc_id, "escalated_to": escalate_to}

    async def _unknown_task(self, task: AgentTask) -> Dict[str, Any]:
        return {"status": "error", "message": f"Unknown task: {task.task_type}"}

    def _init_strategies(self) -> Dict[str, Any]:
        return {
            ExceptionType.DELAY:             {"auto_actions": ["notify_customer", "update_eta", "reroute_if_possible"]},
            ExceptionType.CANCELLATION:      {"auto_actions": ["release_vehicle", "refund_check", "notify_customer"]},
            ExceptionType.RESCHEDULE:        {"auto_actions": ["find_new_slot", "update_route", "notify_customer"]},
            ExceptionType.REROUTE:           {"auto_actions": ["recalculate_route", "update_vehicle", "notify_customer"]},
            ExceptionType.FAILED_DELIVERY:   {"auto_actions": ["retry_attempt", "return_to_hub"]},
            ExceptionType.VEHICLE_BREAKDOWN: {"auto_actions": ["find_replacement_vehicle", "reassign_route", "notify_customer"]},
            ExceptionType.HUB_OVERFLOW:      {"auto_actions": ["reroute_to_alternate", "increase_processing"]},
            ExceptionType.WEATHER_IMPACT:    {"auto_actions": ["delay_notification", "route_avoidance"]},
            ExceptionType.MILER_STALLED:     {"auto_actions": ["reassign", "notify_customer"]},
        }

    def _init_escalation_rules(self) -> Dict[str, Any]:
        return {
            ExceptionType.VEHICLE_BREAKDOWN: {"escalate_to": "FLEET_AGENT"},
            ExceptionType.HUB_OVERFLOW:      {"escalate_to": "HUB_AGENT"},
            ExceptionType.FAILED_DELIVERY:   {"escalate_to": "DISPATCH_AGENT"},
            ExceptionType.DELAY:             {"escalate_to": "JARVIS"},
            ExceptionType.MILER_STALLED:     {"escalate_to": "JARVIS"},
        }

    async def think(self, context: str, options: List[str] = None) -> str:
        return f"[EXCEPTION_AGENT reasoning]: {context}"