#!/usr/bin/env python3
"""
FlowTorch Guidance Provider Interface

Provides additional bias guidance for letter selection beyond base φ₀ steering:
- Morphtable: Deterministic lookup table (offline-computed, auditable)
- Satellite: Remote guidance service (ML-driven, adaptive)
- Hybrid: Morphtable with satellite fallback

Environment:
    FLOWTORCH_GUIDANCE_MODE=off|morphtable|satellite|hybrid
    MORPHTABLE_PATH=/path/to/table.json
    SAT_GUIDANCE_ENDPOINT=http://…/v1/guidance
    SAT_GUIDANCE_API_KEY=...
    SAT_GUIDANCE_TIMEOUT_S=1.5

Observability:
    All guidance decisions are logged in receipts:
    - guidance_mode, morphtable_hash, guidance_latency_ms
    - satellite: guidance_cache_hit, guidance_signature
"""
import json
import os
import time
import hashlib
import hmac
import math
from functools import lru_cache
from typing import List, Dict, Any, Tuple, Optional
from collections import deque

# ============================================================================
# Circuit Breaker for Satellite Failures
# ============================================================================

class CircuitBreaker:
    """
    Track satellite failures and auto-disable after threshold.

    Fail-fast pattern: After N consecutive failures, open the circuit
    for cooldown period before retrying.
    """
    def __init__(self, threshold=5, cooldown_sec=60):
        self.threshold = int(os.getenv("FLOWTORCH_GUIDANCE_BREAKER_THRESHOLD", str(threshold)))
        self.cooldown_sec = float(os.getenv("FLOWTORCH_GUIDANCE_BREAKER_COOLDOWN_SEC", str(cooldown_sec)))
        self.failures = 0
        self.opened_at = None

    def is_open(self) -> bool:
        """Check if circuit is open (blocking requests)."""
        if self.opened_at:
            elapsed = time.time() - self.opened_at
            if elapsed > self.cooldown_sec:
                # Cooldown complete - reset and allow retry
                self.failures = 0
                self.opened_at = None
                return False
            return True
        return False

    def record_failure(self):
        """Increment failure count and open circuit if threshold reached."""
        self.failures += 1
        if self.failures >= self.threshold:
            self.opened_at = time.time()

    def record_success(self):
        """Reset failure count on successful request."""
        self.failures = 0
        self.opened_at = None

    def status(self) -> Dict:
        """Return circuit breaker status for health endpoint."""
        return {
            "failures": self.failures,
            "threshold": self.threshold,
            "is_open": self.is_open(),
            "cooldown_sec": self.cooldown_sec,
            "opened_at": self.opened_at
        }

# Global circuit breaker instance
_satellite_breaker = CircuitBreaker()

# ============================================================================
# Vector Clamping (Prevent Guidance from Swamping φ₀)
# ============================================================================

def _clamp_bias(bias: List[float]) -> List[float]:
    """
    Clamp guidance vector to prevent swamping φ₀ steering.

    Two strategies (both applied):
    1. Per-element clamp: |bias[i]| ≤ max_delta
    2. L2 norm clamp: ||bias|| ≤ max_l2

    Environment:
        FLOWTORCH_GUIDANCE_MAX_DELTA=2.0  (per-element max)
        FLOWTORCH_GUIDANCE_MAX_L2=4.0     (L2 norm max)
    """
    max_delta = float(os.getenv("FLOWTORCH_GUIDANCE_MAX_DELTA", "2.0"))
    max_l2 = float(os.getenv("FLOWTORCH_GUIDANCE_MAX_L2", "4.0"))

    # Per-element clamp
    clamped = [max(-max_delta, min(max_delta, b)) for b in bias]

    # L2 norm clamp (preserve direction)
    norm = math.sqrt(sum(b**2 for b in clamped))
    if norm > max_l2:
        scale = max_l2 / norm
        clamped = [b * scale for b in clamped]

    return clamped

# ============================================================================
# Hash and Morphtable Functions
# ============================================================================

def _hash_file(path: str) -> str:
    """Compute SHA-256 hash of file for receipt verification."""
    with open(path, 'rb') as f:
        return hashlib.sha256(f.read()).hexdigest()

@lru_cache(maxsize=4)
def _load_morphtable(path: str) -> Tuple[Dict, str]:
    """
    Load and cache morphtable with hash.

    Returns:
        (table_data, sha256_hash)
    """
    with open(path) as f:
        table = json.load(f)

    # Validate schema
    if "degrees" not in table or "bias" not in table:
        raise ValueError(f"Invalid morphtable: missing 'degrees' or 'bias' keys")

    if len(table["degrees"]) != len(table["bias"]):
        raise ValueError(f"Morphtable length mismatch: {len(table['degrees'])} degrees vs {len(table['bias'])} biases")

    return table, _hash_file(path)

def _interp_morphtable_bias(table: Dict, phi_deg: float) -> List[float]:
    """
    Interpolate bias from morphtable at given φ₀.

    Args:
        table: {"degrees": [0, 30, 60, ...], "bias": [[bA, bB, bC, bD], ...]}
        phi_deg: φ₀ in degrees

    Returns:
        [bA, bB, bC, bD] interpolated bias vector
    """
    degrees = table["degrees"]
    biases = table["bias"]

    # Normalize φ to [0, 360)
    phi = phi_deg % 360.0

    # Find segment containing phi
    for i in range(len(degrees) - 1):
        if degrees[i] <= phi < degrees[i+1]:
            # Linear interpolation
            t = (phi - degrees[i]) / (degrees[i+1] - degrees[i])
            a, b = biases[i], biases[i+1]
            return [a[k] + t * (b[k] - a[k]) for k in range(4)]

    # Wrap-around case (between last and first point)
    a, b = biases[-1], biases[0]
    span = (360 - degrees[-1]) + degrees[0]

    if phi >= degrees[-1]:
        t = (phi - degrees[-1]) / span
    else:
        t = (phi + 360 - degrees[-1]) / span

    return [a[k] + t * (b[k] - a[k]) for k in range(4)]

def _satellite_guidance_bias(messages: List[Dict], controls: Dict, phi0: float) -> Tuple[List[float], Dict]:
    """
    Query remote satellite guidance service.

    Args:
        messages: Chat messages
        controls: Steering controls
        phi0: φ₀ value

    Returns:
        (bias_vector, metadata)
    """
    # Check circuit breaker first
    if _satellite_breaker.is_open():
        return [0.0, 0.0, 0.0, 0.0], {"error": "circuit_breaker_open"}

    try:
        import requests
    except ImportError:
        return [0.0, 0.0, 0.0, 0.0], {"error": "requests module unavailable"}

    endpoint = os.getenv("SAT_GUIDANCE_ENDPOINT")
    if not endpoint:
        return [0.0, 0.0, 0.0, 0.0], {"error": "SAT_GUIDANCE_ENDPOINT not set"}

    api_key = os.getenv("SAT_GUIDANCE_API_KEY", "")
    timeout = float(os.getenv("SAT_GUIDANCE_TIMEOUT_S", "1.5"))

    # Build request
    payload = {
        "phi0_deg": phi0,
        "controls": controls,
        "messages": messages
    }

    headers = {}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    t0 = time.time()
    meta = {"guidance_latency_ms": None, "guidance_cache_hit": None, "guidance_signature": None}

    try:
        r = requests.post(endpoint, json=payload, headers=headers, timeout=timeout)
        r.raise_for_status()
        response = r.json()

        meta["guidance_latency_ms"] = int(1000 * (time.time() - t0))
        meta["guidance_cache_hit"] = response.get("cache_hit")
        meta["guidance_signature"] = response.get("signature")  # Raw signature from server
        meta["model_version"] = response.get("model_version")  # For signature verification

        # Extract bias vector
        bias = response.get("bias", [0.0, 0.0, 0.0, 0.0])
        if len(bias) != 4:
            _satellite_breaker.record_failure()
            return [0.0, 0.0, 0.0, 0.0], {**meta, "error": f"Invalid bias length: {len(bias)}"}

        # Success - reset circuit breaker
        _satellite_breaker.record_success()
        return [float(x) for x in bias], meta

    except requests.Timeout:
        meta["guidance_latency_ms"] = int(1000 * (time.time() - t0))
        _satellite_breaker.record_failure()
        return [0.0, 0.0, 0.0, 0.0], {**meta, "error": "timeout"}

    except Exception as e:
        meta["guidance_latency_ms"] = int(1000 * (time.time() - t0))
        _satellite_breaker.record_failure()
        return [0.0, 0.0, 0.0, 0.0], {**meta, "error": str(e)[:100]}

def get_guidance_bias(messages: List[Dict], controls: Optional[Dict] = None) -> Tuple[List[float], Dict]:
    """
    Get guidance bias vector based on configured mode.

    Args:
        messages: Chat messages (list of {role, content} dicts)
        controls: Steering controls dict (may contain phi0_deg or mod.value)

    Returns:
        (bias_vector, metadata_dict)

        bias_vector: [bA, bB, bC, bD] logit deltas
        metadata: {
            "guidance_mode": str,
            "morphtable_hash": Optional[str],
            "guidance_latency_ms": Optional[int],
            "guidance_cache_hit": Optional[bool],
            "guidance_signature": Optional[str],
            "error": Optional[str]
        }
    """
    mode = os.getenv("FLOWTORCH_GUIDANCE_MODE", "off").lower()

    meta = {
        "guidance_mode": mode,
        "morphtable_hash": None,
        "guidance_latency_ms": None,
        "guidance_cache_hit": None,
        "guidance_signature": None,
        "error": None
    }

    # Early exit for off mode
    if mode == "off":
        return [0.0, 0.0, 0.0, 0.0], meta

    # Extract φ₀ from controls
    phi0 = None
    if controls:
        # Try direct phi0_deg first
        phi0 = controls.get("phi0_deg")

        # Try nested mod.value format
        if phi0 is None and "mod" in controls:
            mod = controls["mod"]
            if isinstance(mod, dict):
                phi0 = mod.get("value")

    if phi0 is None:
        # No φ₀ specified - return zero bias
        meta["error"] = "no_phi0_in_controls"
        return [0.0, 0.0, 0.0, 0.0], meta

    phi0 = float(phi0)

    # Morphtable mode
    if mode in ("morphtable", "hybrid"):
        path = os.getenv("MORPHTABLE_PATH")
        if path and os.path.exists(path):
            try:
                table, table_hash = _load_morphtable(path)
                bias = _interp_morphtable_bias(table, phi0)
                meta["morphtable_hash"] = table_hash
                # Clamp bias to prevent swamping φ₀
                bias = _clamp_bias(bias)
                return bias, meta
            except Exception as e:
                meta["error"] = f"morphtable_error: {str(e)[:50]}"
                # Fall through to satellite if hybrid

    # Satellite mode
    if mode in ("satellite", "hybrid"):
        bias, sat_meta = _satellite_guidance_bias(messages, controls or {}, phi0)
        meta.update(sat_meta)

        # Verify signature if present
        if sat_meta.get("guidance_signature") and not sat_meta.get("error"):
            sec = os.getenv("SAT_SIGNATURE_SECRET")
            if sec:
                # Verify HMAC-SHA256 signature
                sig_from_server = sat_meta["guidance_signature"]
                # Reconstruct payload (without signature field)
                payload = {"bias": bias, "cache_hit": sat_meta.get("guidance_cache_hit")}
                if "model_version" in sat_meta:
                    payload["model_version"] = sat_meta["model_version"]

                # Compute expected signature
                msg = json.dumps(payload, separators=(',', ':'), sort_keys=True).encode()
                expected_sig = hmac.new(sec.encode(), msg, hashlib.sha256).hexdigest()

                if sig_from_server != expected_sig:
                    # Signature invalid - fail-open to zero bias
                    meta["guidance_signature"] = "invalid"
                    meta["error"] = "signature_verification_failed"
                    return [0.0, 0.0, 0.0, 0.0], meta
                else:
                    meta["guidance_signature"] = "valid"

        if bias != [0.0, 0.0, 0.0, 0.0] or mode == "satellite":
            # Clamp bias to prevent swamping φ₀
            bias = _clamp_bias(bias)
            return bias, meta

    # Fallback: no guidance available
    if not meta["error"]:
        meta["error"] = "no_guidance_available"

    return [0.0, 0.0, 0.0, 0.0], meta
