"""Resilience utilities for agent runs: circuit breaker, structured errors, parking and resume.

Used by PB-AGT-09-02.

Provides:
- CircuitBreaker: CLOSED, OPEN and HALF_OPEN states for one tool. An operator override
  (the App Configuration feature flag `tool_<name>_force_open`) can force it OPEN.
  When the override is removed, the breaker goes to HALF_OPEN and lets one probe through.
- otel_state_listener: keeps the `agent.circuit_breaker.open` metric for the Section 7 alert.
- tool_error: the structured error a failed tool returns to the model or supervisor.
- park_run: parks an interrupted agent run on an Azure Service Bus queue.
- resume_run: resumes an interrupted LangGraph run from its latest checkpoint.
- resume_parked_runs: the resume worker; resumes parked runs in small batches.

The breaker state lives in process memory, so each replica keeps its own state.
The operator override is the shared control across replicas.

Run `python resilience.py` to self-test the breaker. The self-test needs no Azure SDK.
"""
from __future__ import annotations

import json
import threading
import time
from typing import Any, Callable

# Returns "OPEN" when the feature flag `tool_<name>_force_open` is enabled, or None otherwise.
# It is called on every allow_request(), so read the flag through the App Configuration
# provider with feature-flag refresh enabled (cached between refreshes), not on every call.
OverrideReader = Callable[[str], "str | None"]

# Called as (tool_name, old_state, new_state) on every state change.
# Use otel_state_listener() to turn it into the `agent.circuit_breaker.open` metric.
StateListener = Callable[[str, str, str], None]


class CircuitBreaker:
    """Circuit breaker for one tool or model dependency."""

    def __init__(
        self,
        tool_name: str,
        failure_threshold: int = 5,
        cool_off_seconds: float = 60.0,
        read_override: OverrideReader | None = None,
        on_state_change: StateListener | None = None,
    ) -> None:
        self.tool_name = tool_name
        self.failure_threshold = failure_threshold
        self.cool_off_seconds = cool_off_seconds
        self.read_override = read_override
        self.on_state_change = on_state_change
        self.state = "CLOSED"
        self.failure_count = 0
        self.opened_at: float | None = None
        self._forced_open = False
        self._probe_in_flight = False
        self._lock = threading.Lock()

    def _set_state(self, new_state: str) -> None:
        old_state, self.state = self.state, new_state
        if old_state != new_state and self.on_state_change is not None:
            self.on_state_change(self.tool_name, old_state, new_state)

    def allow_request(self) -> bool:
        """Return True if a call may go to the dependency now."""
        forced = self.read_override is not None and self.read_override(self.tool_name) == "OPEN"
        with self._lock:
            if forced:
                # The override drives the real state, so removing it never skips Half-Open.
                self._forced_open = True
                self.opened_at = time.monotonic()
                self._set_state("OPEN")
                return False
            if self._forced_open:
                # Override removed: the operator already probed the tool (Step 7), so go
                # straight to Half-Open and let exactly one call through.
                self._forced_open = False
                self._set_state("HALF_OPEN")
            if self.state == "CLOSED":
                return True
            if self.state == "OPEN":
                if self.opened_at is None or time.monotonic() - self.opened_at < self.cool_off_seconds:
                    return False
                self._set_state("HALF_OPEN")
            # HALF_OPEN: let exactly one probe through until its result is recorded.
            if self._probe_in_flight:
                return False
            self._probe_in_flight = True
            return True

    def record_result(self, success: bool) -> None:
        """Record the result of a call that allow_request() let through."""
        with self._lock:
            was_probe = self.state == "HALF_OPEN"
            self._probe_in_flight = False
            if success:
                self.failure_count, self.opened_at = 0, None
                self._set_state("CLOSED")
                return
            self.failure_count += 1
            if was_probe or self.failure_count >= self.failure_threshold:
                self.opened_at = time.monotonic()
                self._set_state("OPEN")


def otel_state_listener(counter: Any | None = None) -> StateListener:
    """Return an on_state_change listener that keeps the `agent.circuit_breaker.open` metric.

    The metric is an OpenTelemetry UpDownCounter with a `tool` attribute: +1 when a breaker
    opens and -1 when it leaves OPEN. Call configure_azure_monitor() at start-up (Azure Monitor
    OpenTelemetry Distro) so the metric reaches Application Insights.
    """
    if counter is None:
        from opentelemetry import metrics

        counter = metrics.get_meter("agent.resilience").create_up_down_counter(
            "agent.circuit_breaker.open", unit="1", description="Open circuit breakers per tool"
        )

    def listener(tool_name: str, old_state: str, new_state: str) -> None:
        if new_state == "OPEN" and old_state != "OPEN":
            counter.add(1, {"tool": tool_name})
        elif old_state == "OPEN" and new_state != "OPEN":
            counter.add(-1, {"tool": tool_name})

    return listener


def tool_error(
    error_type: str,
    request: dict[str, Any],
    retryable: bool,
    partial_results: list[Any] | None = None,
    alternatives: list[str] | None = None,
) -> dict[str, Any]:
    """Build the structured error a failed tool returns instead of a fake success.

    error_type is an access failure such as "timeout", "unavailable" or "throttled".
    A call that succeeds with no results is not an error: return an empty result instead.
    """
    return {
        "is_error": True,
        "error_type": error_type,
        "retryable": retryable,
        "request": request,
        "partial_results": partial_results or [],
        "alternatives": alternatives or [],
    }


def park_run(
    servicebus_fqdn: str,
    queue_name: str,
    thread_id: str,
    tool_name: str,
    error_code: str,
    credential: Any | None = None,
) -> None:
    """Park a suspended run on Service Bus.

    message_id is derived from thread_id, so duplicate detection on the queue
    (if enabled) drops a second park of the same run.
    """
    from azure.identity import DefaultAzureCredential
    from azure.servicebus import ServiceBusClient, ServiceBusMessage

    payload = {"thread_id": thread_id, "tool_name": tool_name, "error_code": error_code, "parked_at": time.time()}
    message = ServiceBusMessage(
        body=json.dumps(payload),
        message_id=f"park-{thread_id}",
        application_properties={"reason": "dependency_outage", "tool": tool_name},
    )
    with ServiceBusClient(servicebus_fqdn, credential or DefaultAzureCredential()) as client:
        with client.get_queue_sender(queue_name=queue_name) as sender:
            sender.send_messages(message)


def resume_run(graph: Any, thread_id: str) -> Any:
    """Resume a LangGraph run from its latest checkpoint on the same thread.

    Successful nodes of the interrupted super-step are not re-run, but the node that
    failed runs again from its start. Every write in that node needs an idempotency key.
    """
    return graph.invoke(None, {"configurable": {"thread_id": thread_id}})


def resume_parked_runs(
    graph: Any,
    servicebus_fqdn: str,
    queue_name: str,
    batch_size: int = 5,
    delay_seconds: float = 2.0,
    max_batches: int | None = None,
    credential: Any | None = None,
) -> int:
    """Resume parked runs in small batches and return how many completed.

    Each message is completed only after its run resumes. At the first failure the message
    is abandoned (it is redelivered, and goes to the dead-letter queue after MaxDeliveryCount)
    and the worker stops, so the operator can roll back (Section 8). Locks are renewed
    automatically while a run resumes. Run with max_batches=1 first.
    """
    from azure.identity import DefaultAzureCredential
    from azure.servicebus import AutoLockRenewer, ServiceBusClient

    resumed, batches = 0, 0
    renewer = AutoLockRenewer(max_lock_renewal_duration=600)
    try:
        with ServiceBusClient(servicebus_fqdn, credential or DefaultAzureCredential()) as client:
            with client.get_queue_receiver(queue_name=queue_name, auto_lock_renewer=renewer) as receiver:
                while max_batches is None or batches < max_batches:
                    messages = receiver.receive_messages(max_message_count=batch_size, max_wait_time=5)
                    if not messages:
                        break
                    for index, message in enumerate(messages):
                        thread_id = json.loads(str(message))["thread_id"]
                        try:
                            resume_run(graph, thread_id)
                        except Exception:
                            # Return this message and the rest of the batch to the queue.
                            for unsettled in messages[index:]:
                                receiver.abandon_message(unsettled)
                            raise
                        receiver.complete_message(message)
                        resumed += 1
                    batches += 1
                    time.sleep(delay_seconds)
    finally:
        renewer.close()
    return resumed


if __name__ == "__main__":
    # Self-test of the breaker state machine. No network calls.
    override: dict[str, str | None] = {"value": None}
    changes: list[tuple[str, str]] = []
    breaker = CircuitBreaker(
        "mock_tool", failure_threshold=2, cool_off_seconds=0.05,
        read_override=lambda _name: override["value"],
        on_state_change=lambda _name, old, new: changes.append((old, new)),
    )

    assert breaker.allow_request()
    breaker.record_result(False)
    assert breaker.state == "CLOSED"
    breaker.record_result(False)
    assert breaker.state == "OPEN" and not breaker.allow_request(), "opens at the threshold"

    time.sleep(0.06)
    assert breaker.allow_request() and breaker.state == "HALF_OPEN", "half-opens after cool-off"
    assert not breaker.allow_request(), "only one probe while half-open"
    breaker.record_result(False)
    assert breaker.state == "OPEN", "a failed probe reopens the breaker"

    time.sleep(0.06)
    assert breaker.allow_request()
    breaker.record_result(True)
    assert breaker.state == "CLOSED", "a successful probe closes the breaker"

    # Override while the breaker is CLOSED (below threshold), as in Step 2 of the runbook.
    override["value"] = "OPEN"
    assert not breaker.allow_request() and breaker.state == "OPEN", "the override sets the real state to OPEN"
    override["value"] = None
    assert breaker.allow_request() and breaker.state == "HALF_OPEN", "removing the override goes to Half-Open"
    assert not breaker.allow_request(), "only one probe after the override is removed"
    breaker.record_result(True)
    assert breaker.state == "CLOSED"

    assert ("CLOSED", "OPEN") in changes and ("OPEN", "HALF_OPEN") in changes, "state changes are reported"

    class FakeCounter:
        def __init__(self) -> None:
            self.value = 0

        def add(self, amount: int, attributes: dict[str, str]) -> None:
            self.value += amount

    gauge = FakeCounter()
    metric_breaker = CircuitBreaker("metric_tool", failure_threshold=1, cool_off_seconds=0.05,
                                    on_state_change=otel_state_listener(gauge))
    metric_breaker.record_result(False)
    assert gauge.value == 1, "metric is 1 while the breaker is open"
    time.sleep(0.06)
    metric_breaker.allow_request()
    assert gauge.value == 0, "metric returns to 0 when the breaker leaves OPEN"
    error = tool_error("timeout", {"query": "q"}, retryable=True)
    assert error["is_error"] and error["partial_results"] == []
    print("CircuitBreaker self-test passed.")