"""
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 uuid
from datetime import datetime, timedelta
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 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"


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._notified_stalls: Set[str] = set()
        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    = datetime.now().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_str = prev.get("position_unchanged_since", now_ts)
            try:
                unchanged_since = datetime.fromisoformat(unchanged_since_str)
            except ValueError:
                unchanged_since = datetime.now()

            minutes_stalled = (datetime.now() - unchanged_since).total_seconds() / 60

            if minutes_stalled >= STALL_MINUTES:
                booking = await self._get_active_booking(miler_id)
                if booking:
                    booking_id = booking["booking_id"]
                    if booking_id not in self._notified_stalls:
                        await self._publish_stall(miler_id, 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 = datetime.now()
        for row in rows:
            miler_id   = row["miler_id"]
            booking_id = row["booking_id"]

            if booking_id in self._notified_stalls:
                continue

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

            try:
                updated_at    = datetime.fromisoformat(movement["updated_at"])
                minutes_stale = (now - updated_at).total_seconds() / 60
            except ValueError:
                continue

            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 ───────────────────────────────────────────────

    async def _publish_stall(self, miler_id: str, booking_id: str, minutes_stalled: float):
        self._notified_stalls.add(booking_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):
        """Reassign booking and notify customer via Go API."""
        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)")

        headers = {"X-Internal-Key": INTERNAL_API_KEY}

        reassign_result = await api_post(
            f"{GO_API_BASE_URL}/api/v1/internal/bookings/{booking_id}/reassign",
            json={"reason": "miler_stalled"},
            headers=headers,
        )
        logger.info(f"Reassign {'OK' if reassign_result is not None else 'FAILED'} for booking {booking_id}")

        notify_result = await api_post(
            f"{GO_API_BASE_URL}/api/v1/internal/notify",
            json={
                "booking_id": booking_id,
                "message":    "We detected a delay, finding you a new miler",
                "target":     "customer",
            },
            headers=headers,
        )
        logger.info(f"Notify {'OK' if notify_result is not None else 'FAILED'} for booking {booking_id}")

        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=booking_id,
            exception_type=ExceptionType.MILER_STALLED,
            severity=ExceptionSeverity.HIGH,
            description=f"Miler {miler_id} stalled {minutes_stalled} min",
            detected_at=datetime.now(), resolved_at=datetime.now(),
            resolution="Reassignment triggered + customer notified",
            assigned_to=None, status="resolved",
            actions_taken=["reassign", "notify_customer"],
        )

    # ── 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="DISPATCH_AGENT", message_type=MessageType.AGENT_TASK,
            payload={"task_type": "reschedule_route", "order_id": order_id, "new_date": new_date},
        )
        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}"