"""
Depth Router for Edge Inference (WO-EDGE-DEPTH-ROUTER-01)

Shell-based routing using Symbolic Entropy (Se) to choose minimum depth tier
that passes fidelity gate. Based on nested Fibonacci sphere topology.

Shells (onion layers):
    CORE:    blocks 0-4   (~53 MB) - innermost, simple queries
    SHALLOW: blocks 0-8   (~106 MB) - factual Q&A
    MEDIUM:  blocks 0-16  (~212 MB) - multi-step reasoning
    DEEP:    blocks 0-32  (~424 MB) - full model
    CLOUD:   external API - when local won't pass fidelity

Calibration-based routing:
    Se = H × C × D  →  calibration table  →  shell tier  →  fire once
"""

from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Optional, List, Dict, Any, Tuple, NamedTuple
import hashlib
import json
import math
import time


# =============================================================================
# Constants
# =============================================================================

RAM_PER_BLOCK_MB = 13.25  # Conservative estimate with headroom

# Default calibration (can be overridden by loading from file)
DEFAULT_CALIBRATION = {
    "schema": "depth_router_calibration_v1",
    "tiers": [
        {"name": "CORE",    "max_depth": 4,  "se_max": 0.13, "shell_id": 0},
        {"name": "SHALLOW", "max_depth": 8,  "se_max": 0.31, "shell_id": 1},
        {"name": "MEDIUM",  "max_depth": 16, "se_max": 0.62, "shell_id": 2},
        {"name": "DEEP",    "max_depth": 32, "se_max": 1.00, "shell_id": 3},
    ],
    "safety_margin": {"tier_bump_if_close": 0.03},
    "ram_per_block_mb": RAM_PER_BLOCK_MB,
}


# =============================================================================
# DepthTier Enum (Shell-based)
# =============================================================================

class DepthTier(Enum):
    """
    Depth tiers mapped to nested Fibonacci sphere shells.

    Each shell corresponds to a depth range in the transformer:
    - CORE: innermost shell, early layers only
    - SHALLOW/MEDIUM/DEEP: progressively outer shells
    - CLOUD: external fallback when local can't pass fidelity
    """
    CORE = 4       # Shell 0: blocks 0-4
    SHALLOW = 8    # Shell 1: blocks 0-8
    MEDIUM = 16    # Shell 2: blocks 0-16
    DEEP = 32      # Shell 3: blocks 0-32 (full)
    CLOUD = -1     # External API fallback

    @property
    def block_count(self) -> int:
        return self.value if self.value > 0 else 0

    @property
    def shell_id(self) -> int:
        """Shell index in nested Fibonacci sphere (0=core, 3=outer)."""
        mapping = {
            DepthTier.CORE: 0,
            DepthTier.SHALLOW: 1,
            DepthTier.MEDIUM: 2,
            DepthTier.DEEP: 3,
            DepthTier.CLOUD: -1,
        }
        return mapping.get(self, -1)

    @property
    def estimated_ram_mb(self) -> float:
        if self.value <= 0:
            return 0.0
        return self.value * RAM_PER_BLOCK_MB

    @property
    def use_case(self) -> str:
        cases = {
            DepthTier.CORE: "Greetings, simple lookups, echo",
            DepthTier.SHALLOW: "Factual Q&A, simple reasoning",
            DepthTier.MEDIUM: "Multi-step reasoning, analysis",
            DepthTier.DEEP: "Complex analysis (full model)",
            DepthTier.CLOUD: "High-Se queries via cloud API",
        }
        return cases.get(self, "Unknown")


# Backward compatibility alias
InferenceDepth = DepthTier


# =============================================================================
# Se Components
# =============================================================================

class SeComponents(NamedTuple):
    """
    Symbolic Entropy components: Se = H × C × D

    H: Entropy/uncertainty (token diversity, byte entropy)
    C: Coherence (context alignment, retrieval similarity)
    D: Depth requirement (composition complexity, late-layer dependence)
    """
    Se: float
    H: float
    C: float
    D: float


# =============================================================================
# Calibration
# =============================================================================

@dataclass
class Calibration:
    """Loaded calibration table for Se → tier routing."""
    tiers: List[Dict[str, Any]]
    safety_margin: float
    ram_per_block_mb: float
    policy_hash: str  # SHA256 of calibration for audit

    @classmethod
    def load(cls, path: Optional[Path] = None) -> "Calibration":
        """Load calibration from file or use defaults."""
        if path and Path(path).exists():
            with open(path) as f:
                data = json.load(f)
        else:
            data = DEFAULT_CALIBRATION

        # Compute policy hash for audit trail
        policy_bytes = json.dumps(data, sort_keys=True).encode()
        policy_hash = hashlib.sha256(policy_bytes).hexdigest()[:16]

        return cls(
            tiers=data["tiers"],
            safety_margin=data.get("safety_margin", {}).get("tier_bump_if_close", 0.03),
            ram_per_block_mb=data.get("ram_per_block_mb", RAM_PER_BLOCK_MB),
            policy_hash=policy_hash,
        )

    def choose_tier(self, se: float) -> Tuple[DepthTier, int]:
        """
        Choose minimum tier that handles this Se value.

        Returns:
            (tier, shell_id)
        """
        # Apply safety margin near boundaries
        adjusted_se = se
        for tier_info in self.tiers:
            boundary = tier_info["se_max"]
            if abs(se - boundary) < self.safety_margin and se < boundary:
                # Close to boundary, bump up
                adjusted_se = boundary + 0.001
                break

        # Find minimum tier
        for tier_info in self.tiers:
            if adjusted_se <= tier_info["se_max"]:
                tier_name = tier_info["name"]
                shell_id = tier_info["shell_id"]
                return DepthTier[tier_name], shell_id

        # Default to CLOUD if nothing matches
        return DepthTier.CLOUD, -1


# Global calibration (lazy loaded)
_calibration: Optional[Calibration] = None


def get_calibration(path: Optional[Path] = None) -> Calibration:
    """Get or load calibration."""
    global _calibration
    if _calibration is None:
        _calibration = Calibration.load(path)
    return _calibration


def reload_calibration(path: Optional[Path] = None) -> Calibration:
    """Force reload calibration from file."""
    global _calibration
    _calibration = Calibration.load(path)
    return _calibration


# =============================================================================
# Se Computation
# =============================================================================

def compute_se(
    prompt: str,
    context: Optional[str] = None,
    probes: Optional[Dict[str, Any]] = None,
) -> SeComponents:
    """
    Compute Symbolic Entropy components for routing.

    Se = H × C × D where:
    - H: Entropy proxy (token/byte uncertainty)
    - C: Coherence (context alignment) - placeholder for now
    - D: Depth requirement (complexity indicators)

    Args:
        prompt: Input prompt text
        context: Optional context/history for coherence
        probes: Optional early activation probes for better estimates

    Returns:
        SeComponents with (Se, H, C, D)
    """
    H = _estimate_H(prompt)
    C = _estimate_C(prompt, context)
    D = _estimate_D(prompt, probes)

    Se = H * C * D
    return SeComponents(Se=Se, H=H, C=C, D=D)


def _estimate_H(prompt: str) -> float:
    """
    Estimate entropy/uncertainty from prompt.

    Uses byte-level entropy as cheap proxy for token diversity.
    """
    if not prompt:
        return 0.1

    # Byte entropy
    byte_counts = {}
    for b in prompt.encode('utf-8'):
        byte_counts[b] = byte_counts.get(b, 0) + 1

    total = len(prompt.encode('utf-8'))
    if total == 0:
        return 0.1

    entropy = 0.0
    for count in byte_counts.values():
        p = count / total
        if p > 0:
            entropy -= p * math.log2(p)

    # Normalize to [0, 1] (max entropy for uniform 256 = 8 bits)
    normalized = min(1.0, entropy / 8.0)

    # Scale to useful range [0.1, 1.0]
    return 0.1 + normalized * 0.9


def _estimate_C(prompt: str, context: Optional[str] = None) -> float:
    """
    Estimate coherence from context alignment.

    For now, returns 1.0 (placeholder).
    Future: similarity to cached contexts, retrieval hits.
    """
    # TODO: Implement actual coherence estimation
    # - Compare prompt to recent context embeddings
    # - Check retrieval similarity scores
    # - Detect contradictions/instruction density
    return 1.0


def _estimate_D(prompt: str, probes: Optional[Dict[str, Any]] = None) -> float:
    """
    Estimate depth requirement (late-layer dependence).

    Heuristics for composition complexity:
    - Multi-hop reasoning indicators
    - Constraint density
    - Question complexity
    """
    if not prompt:
        return 0.2

    prompt_lower = prompt.lower()

    # Base from length
    length = len(prompt)
    if length < 20:
        base = 0.3
    elif length < 100:
        base = 0.5
    else:
        base = 0.7

    modifiers = 0.0

    # Multi-hop indicators
    multi_hop_words = [
        "then", "after", "before", "because", "therefore",
        "first", "second", "finally", "step", "process"
    ]
    for word in multi_hop_words:
        if word in prompt_lower:
            modifiers += 0.03

    # Complexity indicators (need deeper layers)
    complex_indicators = [
        "explain", "analyze", "compare", "contrast", "synthesize",
        "evaluate", "derive", "prove", "implement", "design"
    ]
    for word in complex_indicators:
        if word in prompt_lower:
            modifiers += 0.05

    # Simple query indicators (early layers sufficient)
    simple_indicators = ["hello", "hi", "hey", "what is", "who is", "when"]
    for word in simple_indicators:
        if prompt_lower.startswith(word):
            modifiers -= 0.2
            break

    # If probes available, use them
    if probes and "early_logits_entropy" in probes:
        # Higher early entropy = need more depth
        probe_entropy = probes["early_logits_entropy"]
        modifiers += probe_entropy * 0.2

    return max(0.1, min(1.0, base + modifiers))


# =============================================================================
# Decision Receipt
# =============================================================================

@dataclass
class DepthRouteDecision:
    """
    Full routing decision with audit trail.

    This is the "sat guidance" sidecar - captures everything needed
    to reproduce and audit the routing decision.
    """
    # Se components
    se: float
    h: float
    c: float
    d: float

    # Decision
    chosen_tier: str
    max_depth: int
    shell_id: int

    # Calibration audit
    policy_hash: str

    # RAM constraint
    available_ram_mb: float
    was_downgraded: bool
    downgrade_reason: Optional[str]

    # Timing
    route_time_ms: float

    # Schema
    schema: str = "depth_route_decision_v1"
    work_order: str = "WO-EDGE-DEPTH-ROUTER-01"

    def to_dict(self) -> Dict[str, Any]:
        return {
            "schema": self.schema,
            "work_order": self.work_order,
            "se": round(self.se, 6),
            "h": round(self.h, 6),
            "c": round(self.c, 6),
            "d": round(self.d, 6),
            "chosen_tier": self.chosen_tier,
            "max_depth": self.max_depth,
            "shell_id": self.shell_id,
            "policy_hash": self.policy_hash,
            "available_ram_mb": self.available_ram_mb,
            "was_downgraded": self.was_downgraded,
            "downgrade_reason": self.downgrade_reason,
            "route_time_ms": round(self.route_time_ms, 3),
        }


# =============================================================================
# Main Routing Functions
# =============================================================================

def choose_tier(
    prompt: str,
    context: Optional[str] = None,
    probes: Optional[Dict[str, Any]] = None,
    available_ram_mb: float = 512.0,
    force_local: bool = False,
    calibration_path: Optional[Path] = None,
) -> DepthRouteDecision:
    """
    Choose depth tier for a query using calibration-based routing.

    This is the main entry point for shell-based routing:
    1. Compute Se = H × C × D
    2. Look up tier from calibration table
    3. Apply RAM constraints
    4. Return decision with full audit trail

    Args:
        prompt: Input query text
        context: Optional context for coherence estimation
        probes: Optional early activation probes
        available_ram_mb: Available RAM for constraint checking
        force_local: Never use cloud fallback
        calibration_path: Path to calibration JSON (uses default if None)

    Returns:
        DepthRouteDecision with tier and audit info
    """
    start_time = time.perf_counter()

    # Load calibration
    calibration = get_calibration(calibration_path)

    # Compute Se
    se_components = compute_se(prompt, context, probes)

    # Choose tier from calibration
    ideal_tier, shell_id = calibration.choose_tier(se_components.Se)

    # Apply RAM constraints
    final_tier = ideal_tier
    was_downgraded = False
    downgrade_reason = None

    if ideal_tier == DepthTier.CLOUD and force_local:
        final_tier = DepthTier.DEEP
        was_downgraded = True
        downgrade_reason = "force_local=True, cloud disabled"

    if final_tier != DepthTier.CLOUD:
        required_ram = final_tier.estimated_ram_mb
        if required_ram > available_ram_mb:
            was_downgraded = True
            original = final_tier

            # Cascade down
            for tier in [DepthTier.MEDIUM, DepthTier.SHALLOW, DepthTier.CORE]:
                if tier.estimated_ram_mb <= available_ram_mb:
                    final_tier = tier
                    break
            else:
                if force_local:
                    final_tier = DepthTier.CORE
                    downgrade_reason = f"RAM insufficient, force_local=True, using CORE with pressure"
                else:
                    final_tier = DepthTier.CLOUD
                    downgrade_reason = f"RAM insufficient for any local tier"

            if downgrade_reason is None:
                downgrade_reason = f"RAM ({available_ram_mb:.0f}MB) < {original.name} ({original.estimated_ram_mb:.0f}MB)"

    route_time_ms = (time.perf_counter() - start_time) * 1000

    return DepthRouteDecision(
        se=se_components.Se,
        h=se_components.H,
        c=se_components.C,
        d=se_components.D,
        chosen_tier=final_tier.name,
        max_depth=final_tier.block_count,
        shell_id=final_tier.shell_id,
        policy_hash=calibration.policy_hash,
        available_ram_mb=available_ram_mb,
        was_downgraded=was_downgraded,
        downgrade_reason=downgrade_reason,
        route_time_ms=route_time_ms,
    )


# =============================================================================
# Backward Compatibility Layer
# =============================================================================

# Old thresholds (for backward compat)
SE_THRESHOLD_SHALLOW = 0.20
SE_THRESHOLD_MEDIUM = 0.45
SE_THRESHOLD_DEEP = 0.70

# Old block counts
BLOCKS_SHALLOW = 2
BLOCKS_MEDIUM = 8
BLOCKS_DEEP = 16
BLOCKS_FULL = 32

MIN_RAM_FOR_LOCAL_MB = RAM_PER_BLOCK_MB * 4  # CORE minimum


@dataclass
class DepthRouteResult:
    """Backward-compatible result (wraps new decision)."""
    depth: DepthTier
    block_count: int
    se_score: float
    available_ram_mb: float
    force_local: bool
    ideal_depth: DepthTier
    was_downgraded: bool
    downgrade_reason: Optional[str] = None
    estimated_ram_mb: float = 0.0
    ram_headroom_mb: float = 0.0
    route_time_ms: float = 0.0
    schema: str = "depth_route_result_v1"
    work_order: str = "WO-EDGE-DEPTH-ROUTER-01"

    def to_dict(self) -> Dict[str, Any]:
        return {
            "schema": self.schema,
            "work_order": self.work_order,
            "depth": self.depth.name,
            "block_count": self.block_count,
            "se_score": self.se_score,
            "available_ram_mb": self.available_ram_mb,
            "force_local": self.force_local,
            "ideal_depth": self.ideal_depth.name,
            "was_downgraded": self.was_downgraded,
            "downgrade_reason": self.downgrade_reason,
            "estimated_ram_mb": self.estimated_ram_mb,
            "ram_headroom_mb": self.ram_headroom_mb,
            "route_time_ms": self.route_time_ms,
        }


def route_se_to_depth(
    se_score: float,
    available_ram_mb: float = 512.0,
    force_local: bool = False,
) -> DepthRouteResult:
    """
    Backward-compatible routing function.

    Maps Se score to depth tier using calibration.
    """
    start_time = time.perf_counter()

    calibration = get_calibration()
    se_score = max(0.0, min(1.0, se_score))

    ideal_tier, _ = calibration.choose_tier(se_score)
    final_tier = ideal_tier
    was_downgraded = False
    downgrade_reason = None

    if ideal_tier == DepthTier.CLOUD and force_local:
        ideal_tier = DepthTier.DEEP
        final_tier = DepthTier.DEEP
        was_downgraded = True
        downgrade_reason = "force_local=True, cloud disabled"

    if final_tier != DepthTier.CLOUD:
        required_ram = final_tier.estimated_ram_mb
        if required_ram > available_ram_mb:
            was_downgraded = True
            original = final_tier

            for tier in [DepthTier.MEDIUM, DepthTier.SHALLOW, DepthTier.CORE]:
                if tier.estimated_ram_mb <= available_ram_mb:
                    final_tier = tier
                    break
            else:
                if force_local:
                    final_tier = DepthTier.CORE
                else:
                    final_tier = DepthTier.CLOUD

            if final_tier != DepthTier.CLOUD:
                downgrade_reason = f"RAM ({available_ram_mb:.1f}MB) insufficient for {original.name}"
            else:
                downgrade_reason = "RAM insufficient for any local depth"

    estimated_ram = final_tier.estimated_ram_mb
    ram_headroom = available_ram_mb - estimated_ram if estimated_ram > 0 else available_ram_mb
    route_time_ms = (time.perf_counter() - start_time) * 1000

    return DepthRouteResult(
        depth=final_tier,
        block_count=final_tier.block_count,
        se_score=se_score,
        available_ram_mb=available_ram_mb,
        force_local=force_local,
        ideal_depth=ideal_tier,
        was_downgraded=was_downgraded,
        downgrade_reason=downgrade_reason,
        estimated_ram_mb=estimated_ram,
        ram_headroom_mb=ram_headroom,
        route_time_ms=route_time_ms,
    )


def get_block_indices_for_depth(depth: DepthTier) -> List[int]:
    """Get block indices for a depth tier (0 to max_depth-1)."""
    if depth == DepthTier.CLOUD or depth.block_count <= 0:
        return []

    # For shell-based routing, use contiguous blocks 0 to N
    return list(range(depth.block_count))


def estimate_se_from_prompt(prompt: str) -> float:
    """Backward-compatible Se estimation (returns just Se value)."""
    components = compute_se(prompt)
    return components.Se


def get_depth_summary() -> Dict[str, Dict[str, Any]]:
    """Get summary of all depth tiers."""
    calibration = get_calibration()

    result = {}
    for tier_info in calibration.tiers:
        name = tier_info["name"]
        tier = DepthTier[name]
        result[name] = {
            "blocks": tier.block_count,
            "ram_mb": tier.estimated_ram_mb,
            "use_case": tier.use_case,
            "se_max": tier_info["se_max"],
            "shell_id": tier_info["shell_id"],
        }

    result["CLOUD"] = {
        "blocks": 0,
        "ram_mb": 0,
        "use_case": DepthTier.CLOUD.use_case,
        "se_max": None,
        "shell_id": -1,
    }

    return result
