#!/usr/bin/env python3
"""
Basin Control Decode (Deterministic PHI0-based inference)

Minimal wrapper for basin KAT generation. Honors phase control (phi0_deg) with:
- temperature=0.0 (deterministic)
- top_p=1.0 (no nucleus sampling)
- fixed seed=0 (reproducible)

For production: wire to actual FlowTorch/CDC inference engine.
For KAT: can use mock but must be deterministic per phi0.
"""

import math
import hashlib
import os
import re
from typing import Dict, Any, Optional

# Mock mode flag (set to False when real engine is wired)
USE_MOCK = os.environ.get("ECHO_BASIN_USE_MOCK", "1") != "0"

# CDR checkpointing (optional - graceful fallback if unavailable)
CDR_ENABLED = os.environ.get("BASIN_CDR_CHECKPOINTS", "1") != "0"
_cdr_manager = None

def get_cdr_manager():
    """Get or create CDR state manager singleton."""
    global _cdr_manager
    if _cdr_manager is None:
        try:
            from helix_cdc.cdr_state_manager import CDRStateManager
            _cdr_manager = CDRStateManager(tape_size=65536, max_checkpoints=1024)
            # Initialize CDR with deterministic seed for basin tracking
            _cdr_manager.cdr_init(seed=b"basin_phase_tracker_seed_v1" + b"\x00" * 37)
        except ImportError:
            pass  # CDR not available, graceful fallback
    return _cdr_manager

# Canonical prompt version (for KAT reproducibility)
CANON_VER = "v1.0"


def canonical_prompt(prompt: str) -> str:
    """
    Canonicalize a prompt for deterministic basin control.

    Transformations:
    1. Remove zero-width characters (U+200B, U+FEFF, etc.)
    2. Remove bidi control characters (LRM, RLM, LRO, RLO, PDF, etc.)
    3. Remove other Unicode control/format characters (Cf category)
    4. Remove ASCII control characters except space
    5. Replace newlines/tabs with spaces
    6. Collapse multiple spaces to single space
    7. Remove tag prefixes: [TEST], [DEMO], [DEBUG], etc.
    8. Remove tag suffixes: [TEST], etc.
    9. Strip leading/trailing whitespace

    This ensures perturbed prompts map to the same canonical form,
    enabling deterministic φ₀ → output mapping.
    """
    import unicodedata

    s = prompt

    # Remove zero-width and invisible characters
    # U+200B Zero Width Space
    # U+200C Zero Width Non-Joiner
    # U+200D Zero Width Joiner
    # U+FEFF Byte Order Mark / Zero Width No-Break Space
    zero_width = '\u200b\u200c\u200d\ufeff'
    for c in zero_width:
        s = s.replace(c, '')

    # Remove bidi control characters
    # U+200E Left-to-Right Mark, U+200F Right-to-Left Mark
    # U+202A-U+202E Directional formatting
    # U+2066-U+2069 Directional isolates
    bidi_chars = '\u200e\u200f\u202a\u202b\u202c\u202d\u202e\u2066\u2067\u2068\u2069'
    for c in bidi_chars:
        s = s.replace(c, '')

    # Remove other format characters (Unicode category Cf) and control chars (Cc)
    # But preserve standard whitespace (space, tab, newline)
    cleaned = []
    for c in s:
        cat = unicodedata.category(c)
        if cat == 'Cf':  # Format characters
            continue
        if cat == 'Cc' and c not in ' \t\n\r':  # Control chars except whitespace
            continue
        cleaned.append(c)
    s = ''.join(cleaned)

    # Normalize all Unicode whitespace (Zs category) to ASCII space
    # This includes NBSP (U+00A0), EM SPACE (U+2003), THIN SPACE (U+2009), etc.
    normalized = []
    for c in s:
        cat = unicodedata.category(c)
        if cat == 'Zs':  # Space_Separator category
            normalized.append(' ')
        elif c in '\t\n\r':  # Tab, newline, carriage return → space
            normalized.append(' ')
        else:
            normalized.append(c)
    s = ''.join(normalized)

    # Collapse multiple spaces to single space
    s = re.sub(r' +', ' ', s)

    # Strip leading/trailing whitespace
    s = s.strip()

    # Remove tag prefixes like [TEST], [DEMO], [DEBUG]
    s = re.sub(r'^\[(TEST|DEMO|DEBUG|INFO|WARN|ERROR)\]\s*', '', s, flags=re.IGNORECASE)

    # Remove tag suffixes like [TEST]
    s = re.sub(r'\s*\[(TEST|DEMO|DEBUG|INFO|WARN|ERROR)\]$', '', s, flags=re.IGNORECASE)

    # Final strip in case tags left whitespace
    s = s.strip()

    return s


HF_MODEL = os.environ.get("HF_MODEL", "mistralai/Mistral-7B-v0.1")
HF_ATTN = os.environ.get("HF_ATTN", "eager")
BASIN_SEED_SPACING = int(os.environ.get("BASIN_SEED_SPACING", "7919"))
BASIN_WARMUP_TOKENS = int(os.environ.get("BASIN_WARMUP_TOKENS", "0"))

def _phi0_to_seed(phi0_deg: float, cardinality: int, base: int = 42, spacing: int = None) -> int:
    """Deterministic seed from φ₀ and basin count."""
    if spacing is None:
        spacing = BASIN_SEED_SPACING
    basin = int((phi0_deg % 360.0) / (360.0 / cardinality)) % cardinality
    # Use prime spacing to spread basins in seed space
    return base + basin * spacing


def decode_with_phi0_deg(prompt: str, phi0_deg: float, cardinality: int = 8, max_new_tokens: int = 128) -> str:
    """
    Deterministic decode honoring phase control (phi0_deg).

    Args:
        prompt: Input prompt
        phi0_deg: Phase angle in degrees (0-360)
        cardinality: Number of basins (K=2,4,8,16)
        max_new_tokens: Token budget

    Returns:
        Generated text (deterministic per phi0_deg + prompt)
    """
    # Convert degrees to radians (for engines that expect radians)
    phi0_rad = (phi0_deg / 180.0) * math.pi

    # Deterministic config
    cfg = dict(
        temperature=0.0,
        top_p=1.0,
        max_new_tokens=max_new_tokens,
        seed=0,  # Fixed seed for reproducibility
        control=dict(
            phi0_deg=float(phi0_deg),
            phi0_rad=float(phi0_rad),
            cardinality=cardinality,
        ),
    )

    if USE_MOCK:
        # Deterministic mock for KAT validation
        # Hash (prompt + phi0) → stable output
        result = _mock_basin_decode(prompt, phi0_deg, cardinality, max_new_tokens)

        # CDR checkpoint for basin phase tracking (enables reversible exploration)
        if CDR_ENABLED:
            mgr = get_cdr_manager()
            if mgr is not None:
                # Checkpoint using CDR absorb (O(1) memory mode)
                state_msg = f"{phi0_deg:.1f}|{cardinality}|{len(result)}".encode()[:64]
                mgr.cdr_absorb(state_msg, store_message=True)

                # Also create named checkpoint for random access
                mgr.checkpoint_basin_phase(
                    phi0_deg=phi0_deg,
                    output_tokens=[ord(c) for c in result[:64]],  # First 64 chars as "tokens"
                    coherence=1.0,  # Mock has perfect coherence
                )

        return result
    else:
        # Real inference using HELIX CDC decoder (vault-anchored)
        try:
            from helix_cdc.cdc_decode import cdc_decode_batch
        except ImportError:
            raise RuntimeError("helix_cdc not available. Install or set ECHO_BASIN_USE_MOCK=1")

        # Set deterministic seed based on φ₀ and K
        seed = _phi0_to_seed(phi0_deg, cardinality)

        # Optional warmup tokens for basin stability
        warm = ""
        if BASIN_WARMUP_TOKENS > 0:
            warm = " " + ("·" * BASIN_WARMUP_TOKENS)

        # Prepare prompt with warmup
        full_prompt = (prompt + warm).strip()

        # Call CDC decoder with basin-controlled seed
        extras = {
            "phi0_deg": phi0_deg,
            "cardinality": cardinality,
            "basin_seed": seed,
            "temperature": 0.0,  # Greedy for determinism
        }

        # Use CDC batch decode (single prompt)
        texts, trace = cdc_decode_batch(
            prompts=[full_prompt],
            max_new_tokens=max_new_tokens,
            temperature=0.0,
            seed_path=None,  # Could point to .hxz seed if needed
            device="cuda" if os.environ.get("CUDA_VISIBLE_DEVICES", "0") != "" else "cpu",
            extras=extras
        )

        result = texts[0]

        # CDR checkpoint for real inference (enables reversible exploration)
        if CDR_ENABLED:
            mgr = get_cdr_manager()
            if mgr is not None:
                # Get coherence from trace if available
                coherence = trace.get("coherence") if isinstance(trace, dict) else None

                # Checkpoint using CDR absorb
                state_msg = f"{phi0_deg:.1f}|{cardinality}|{len(result)}".encode()[:64]
                mgr.cdr_absorb(state_msg, store_message=True)

                # Named checkpoint with full state
                mgr.checkpoint_basin_phase(
                    phi0_deg=phi0_deg,
                    output_tokens=trace.get("tokens", [])[:64] if isinstance(trace, dict) else [],
                    coherence=coherence,
                )

        return result


def _mock_basin_decode(prompt: str, phi0_deg: float, cardinality: int, max_new_tokens: int) -> str:
    """
    Mock decode for KAT validation (deterministic per phi0).

    Uses hash(prompt + phi0) to generate stable but distinct outputs per basin.
    Supports variable cardinality (K=2,4,8,16).
    """
    # Hash prompt + phi0 to create stable seed per basin
    seed_str = f"{prompt}|{phi0_deg:.1f}"
    h = hashlib.sha256(seed_str.encode()).digest()
    seed_int = int.from_bytes(h[:4], 'big')

    # Compute basin_id from phi0 and cardinality
    delta_deg = 360.0 / cardinality
    basin_id = int(phi0_deg / delta_deg) % cardinality

    # Template pool (16 total to support up to K=16)
    # ⚠️ FROZEN FOR KAT REPRODUCIBILITY
    # Golden SHA256: 2cc4d22f76528d929a8039166d5c0cbdc03f4004065097a8cb42ab7586129d85
    # Changing this list alters KAT results. If modified, you MUST:
    # 1. Update golden hash in tools/validate_basin_receipt.py
    # 2. Re-run K-sweep: bash basin/k_sweep_test.sh
    # 3. Regenerate receipt: python3 tools/basin_emit_receipt.py
    templates = [
        "Quantum computing leverages superposition and entanglement to process information in parallel states.",
        "Unit tests validate individual components in isolation, ensuring code reliability and catching regressions early.",
        "Information flows through\nentangled bits of light,\nentropy defined.",
        "Deterministic systems produce identical outputs given the same inputs and initial conditions.",
        "Phase-locked loops synchronize oscillator frequencies by comparing phase differences.",
        "Computational models map abstract problems to concrete algorithmic solutions.",
        "Error correction encodes redundancy to detect and repair corrupted data bits.",
        "Basin control steers outputs toward specific attractors in the solution space.",
        "Distributed consensus protocols coordinate state across replicated nodes.",
        "Symbolic regression searches function space to fit observed data patterns.",
        "Temporal logic specifies system behavior over discrete time steps.",
        "Gradient descent optimizes parameters by following the steepest descent direction.",
        "Binary search trees maintain sorted order for efficient lookup operations.",
        "Hash functions map arbitrary data to fixed-size digests.",
        "Cryptographic signatures provide authenticity and non-repudiation guarantees.",
        "State machines transition between discrete configurations based on input events.",
    ]

    # Select template ONLY by basin_id (not prompt hash)
    # This ensures basin separation: same phi0 → same basin → same base template
    base_text = templates[basin_id % len(templates)]

    # Add basin-specific suffix for differentiation
    suffix = f" [Basin {basin_id}, φ={phi0_deg:.1f}°, K={cardinality}]"

    return base_text + suffix


def backend_decode(prompt: str, temperature: float = 0.7, top_p: float = 0.9,
                   max_new_tokens: int = 128, **kwargs) -> str:
    """
    Fallback decode without basin control (standard inference).

    For production: route to your normal inference path.
    For testing: returns mock output.
    """
    if USE_MOCK:
        return f"Standard inference response to: '{prompt[:50]}...'"
    else:
        # TODO: Wire to actual inference engine
        raise NotImplementedError("Standard inference not yet wired.")
