"""
cortex/hippo/tools/recursion_guard.py — Expansion recursion guard.

Full parity with tools/lcm-expansion-recursion-guard.ts (286 lines).
Prevents infinite delegation loops in lcm_expand_query → lcm_expand chains.

Issue: #1118
"""

from __future__ import annotations

import logging
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, Literal, Optional, Set

logger = logging.getLogger(__name__)

EXPANSION_RECURSION_ERROR_CODE = "EXPANSION_RECURSION_BLOCKED"
EXPANSION_DELEGATION_DEPTH_CAP = 1

TelemetryEvent = Literal["start", "block", "timeout", "success"]


# ═══════════════════════════════════════════════════════════════════════════
# Data types
# ═══════════════════════════════════════════════════════════════════════════


@dataclass
class DelegatedExpansionContext:
    """Metadata stamped onto a child session for recursion/depth tracking."""
    request_id: str
    expansion_depth: int
    origin_session_key: str
    stamped_by: str
    created_at: str  # ISO 8601
    allowed_summary_ids: Optional[frozenset[str]] = None  # #1231: scoped grants


ExpansionRecursionBlockReason = Literal["depth_cap", "idempotent_reentry"]


@dataclass
class ExpansionRecursionGuardDecision:
    """Result of evaluating the expansion recursion guard.

    When blocked=False, the session is allowed to delegate.
    When blocked=True, code/reason/message explain the block.
    """
    blocked: bool
    request_id: str
    expansion_depth: int
    origin_session_key: str
    code: Optional[str] = None
    reason: Optional[ExpansionRecursionBlockReason] = None
    message: Optional[str] = None


# ═══════════════════════════════════════════════════════════════════════════
# In-memory state
# ═══════════════════════════════════════════════════════════════════════════

_delegated_context_by_session_key: Dict[str, DelegatedExpansionContext] = {}
_blocked_request_ids_by_session_key: Dict[str, Set[str]] = {}

_telemetry_counters: Dict[TelemetryEvent, int] = {
    "start": 0,
    "block": 0,
    "timeout": 0,
    "success": 0,
}


# ═══════════════════════════════════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════════════════════════════════


def _normalize_session_key(session_key: Optional[str] = None) -> str:
    return session_key.strip() if isinstance(session_key, str) else ""


def _get_or_init_blocked_request_ids(session_key: str) -> Set[str]:
    existing = _blocked_request_ids_by_session_key.get(session_key)
    if existing is not None:
        return existing
    created: Set[str] = set()
    _blocked_request_ids_by_session_key[session_key] = created
    return created


def _build_expansion_recursion_recovery_guidance(origin_session_key: str) -> str:
    return (
        "Recovery: In delegated sub-agent sessions, call `lcm_expand` directly and synthesize "
        "your answer from that result. Do NOT call `lcm_expand_query` from delegated context. "
        f"If deeper delegation is required, return to the origin session ({origin_session_key}) "
        "and call `lcm_expand_query` there."
    )


# ═══════════════════════════════════════════════════════════════════════════
# Public API
# ═══════════════════════════════════════════════════════════════════════════


def create_expansion_request_id() -> str:
    """Create a stable request identifier for delegated expansion orchestration."""
    return str(uuid.uuid4())


def resolve_expansion_request_id(session_key: Optional[str] = None) -> str:
    """Resolve the active expansion request ID for a session.

    Inherits from any stamped delegated context when present.
    """
    key = _normalize_session_key(session_key)
    ctx = _delegated_context_by_session_key.get(key)
    if ctx is not None:
        return ctx.request_id
    return create_expansion_request_id()


def resolve_next_expansion_depth(session_key: Optional[str] = None) -> int:
    """Resolve the next delegated expansion depth to stamp onto a child session."""
    key = _normalize_session_key(session_key)
    if not key:
        return 1
    existing = _delegated_context_by_session_key.get(key)
    if existing is not None:
        return existing.expansion_depth + 1
    return 1


def stamp_delegated_expansion_context(
    *,
    session_key: str,
    request_id: str,
    expansion_depth: int,
    origin_session_key: str,
    stamped_by: str,
    allowed_summary_ids: Optional[list[str]] = None,
) -> DelegatedExpansionContext:
    """Stamp delegated expansion metadata for a child session.

    Enables re-entry checks to enforce recursion and depth policies.
    When *allowed_summary_ids* is provided, the grant is scoped to only
    those summary IDs — sub-agents should not access summaries outside
    this set (#1231).
    """
    from datetime import datetime, timezone

    key = _normalize_session_key(session_key)
    scoped_ids = (
        frozenset(allowed_summary_ids) if allowed_summary_ids else None
    )
    context = DelegatedExpansionContext(
        request_id=request_id,
        expansion_depth=max(0, int(expansion_depth)),
        origin_session_key=origin_session_key.strip() or "main",
        stamped_by=stamped_by,
        created_at=datetime.now(timezone.utc).isoformat(),
        allowed_summary_ids=scoped_ids,
    )
    if key:
        _delegated_context_by_session_key[key] = context
    return context


def check_summary_id_allowed(
    session_key: Optional[str],
    summary_id: str,
) -> bool:
    """Check whether *summary_id* is allowed by the scoped grant for *session_key*.

    Returns True if:
    - No delegated context exists (no scope restrictions), OR
    - The context has no allowed_summary_ids set (unscoped grant), OR
    - The summary_id is in the allowed set.

    Issue: #1231.
    """
    key = _normalize_session_key(session_key)
    ctx = _delegated_context_by_session_key.get(key)
    if ctx is None or ctx.allowed_summary_ids is None:
        return True
    return summary_id in ctx.allowed_summary_ids


def clear_delegated_expansion_context(session_key: str) -> None:
    """Remove delegated expansion metadata for a child session after cleanup."""
    key = _normalize_session_key(session_key)
    if not key:
        return
    _delegated_context_by_session_key.pop(key, None)
    _blocked_request_ids_by_session_key.pop(key, None)


def evaluate_expansion_recursion_guard(
    *,
    session_key: Optional[str] = None,
    request_id: str,
) -> ExpansionRecursionGuardDecision:
    """Evaluate whether a session is allowed to delegate expansion work.

    Delegated contexts are blocked at depth >= EXPANSION_DELEGATION_DEPTH_CAP,
    with repeated request ID re-entry mapped to an explicit idempotency block reason.
    """
    key = _normalize_session_key(session_key)
    rid = request_id.strip()

    delegated_context = _delegated_context_by_session_key.get(key)

    if delegated_context is None:
        # No delegated context — allowed
        return ExpansionRecursionGuardDecision(
            blocked=False,
            request_id=rid,
            expansion_depth=0,
            origin_session_key=key or "main",
        )

    if delegated_context.expansion_depth < EXPANSION_DELEGATION_DEPTH_CAP:
        return ExpansionRecursionGuardDecision(
            blocked=False,
            request_id=rid,
            expansion_depth=delegated_context.expansion_depth,
            origin_session_key=delegated_context.origin_session_key,
        )

    # At or above depth cap — blocked
    seen = _get_or_init_blocked_request_ids(key)
    is_idempotent_reentry = rid in seen
    seen.add(rid)

    reason: ExpansionRecursionBlockReason = (
        "idempotent_reentry" if is_idempotent_reentry else "depth_cap"
    )

    message = (
        f"{EXPANSION_RECURSION_ERROR_CODE}: Expansion delegation blocked at depth "
        f"{delegated_context.expansion_depth} ({reason}; requestId={rid}; "
        f"origin={delegated_context.origin_session_key}). "
        + _build_expansion_recursion_recovery_guidance(delegated_context.origin_session_key)
    )

    return ExpansionRecursionGuardDecision(
        blocked=True,
        code=EXPANSION_RECURSION_ERROR_CODE,
        reason=reason,
        message=message,
        request_id=rid,
        expansion_depth=delegated_context.expansion_depth,
        origin_session_key=delegated_context.origin_session_key,
    )


def record_expansion_delegation_telemetry(
    *,
    component: str,
    event: TelemetryEvent,
    request_id: str,
    session_key: Optional[str] = None,
    expansion_depth: int,
    origin_session_key: str,
    reason: Optional[str] = None,
    run_id: Optional[str] = None,
) -> None:
    """Emit structured delegated expansion telemetry with monotonic counters."""
    import json

    _telemetry_counters[event] += 1

    payload = {
        "component": component,
        "event": event,
        "requestId": request_id,
        "sessionKey": _normalize_session_key(session_key) or None,
        "expansionDepth": expansion_depth,
        "originSessionKey": origin_session_key,
        "reason": reason,
        "runId": run_id,
        "counters": dict(_telemetry_counters),
    }

    line = f"[lcm][expansion_delegation] {json.dumps(payload)}"
    if event in ("start", "success"):
        logger.info(line)
    else:
        logger.warning(line)


# ═══════════════════════════════════════════════════════════════════════════
# Test helpers
# ═══════════════════════════════════════════════════════════════════════════


def get_delegated_context_for_tests(
    session_key: str,
) -> Optional[DelegatedExpansionContext]:
    """Return the currently stamped delegated expansion context for test assertions."""
    return _delegated_context_by_session_key.get(_normalize_session_key(session_key))


def get_telemetry_snapshot_for_tests() -> Dict[TelemetryEvent, int]:
    """Return the delegated expansion telemetry counters for tests."""
    return dict(_telemetry_counters)


def reset_for_tests() -> None:
    """Reset delegated expansion context and telemetry state between tests."""
    _delegated_context_by_session_key.clear()
    _blocked_request_ids_by_session_key.clear()
    _telemetry_counters["start"] = 0
    _telemetry_counters["block"] = 0
    _telemetry_counters["timeout"] = 0
    _telemetry_counters["success"] = 0
