#!/usr/bin/env python3
"""
FlowTorch Production API Server
Integrated safety rails, security, observability, and multi-state basin control.
"""

import os
import time
import logging
from pathlib import Path
from typing import Dict, Any, Optional
from collections import defaultdict, deque
from fastapi import FastAPI, Request, HTTPException, Response
from fastapi.responses import JSONResponse
import uvicorn

# Optional torch for determinism probes
try:
    import torch
except Exception:
    torch = None

# Import production components
from edge_app.schemas import (
    ChatRequest, HealthResponse, MetricsResponse, CDSResponse, SignedReceipt
)
from edge_app.auto_delta import auto_select_phi0, get_optimal_delta, calculate_phi0
from edge_app.safety import get_monitor, graceful_degrade
from edge_app.security import check_ga_nudge_forbid, check_control_consent, create_signed_receipt
from edge_app.observability import get_metrics, health_check, cds_status
from edge_app.performance import get_concurrency_guard, get_gpu_prewarmer, prewarm_on_startup
from edge_app.basin_decode import decode_with_phi0_deg, backend_decode, get_cdr_manager, CDR_ENABLED

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# GA-OFF tripwire on startup (fail-fast)
try:
    check_ga_nudge_forbid()
    logger.info("✓ GA-OFF mode active (GA_NUDGE_FORBID=1)")
except RuntimeError as e:
    logger.error(f"✗ GA-OFF tripwire failed: {e}")
    raise

# App initialization
app = FastAPI(
    title="FlowTorch Production API",
    version="v1.0-basin-control",
    description="Multi-state basin steering with safety rails and signed receipts"
)

# Global state
ROOT = Path(__file__).resolve().parents[1]
startup_time = time.time()

# Rate limiting (lightweight token bucket)
# PRODUCTION DEFAULTS: 2.0 QPS, burst=10 (conservative for abuse prevention)
# TESTING OVERRIDE: Set RATE_LIMIT_QPS=1000 for grid tests (disable rate limiting)
RATE_LIMIT_QPS = float(os.environ.get("RATE_LIMIT_QPS", "2.0"))         # tokens per second
RATE_LIMIT_BURST = int(os.environ.get("RATE_LIMIT_BURST", "10"))        # bucket size
_ip_bucket = {}
_audit = deque(maxlen=5000)

def _allow(ip: str) -> bool:
    now = time.time()
    tokens, last = _ip_bucket.get(ip, (RATE_LIMIT_BURST, now))
    tokens = min(RATE_LIMIT_BURST, tokens + (now - last) * RATE_LIMIT_QPS)
    if tokens < 1.0:
        _ip_bucket[ip] = (tokens, now)
        return False
    _ip_bucket[ip] = (tokens - 1.0, now)
    return True


@app.on_event("startup")
async def startup_event():
    """Run GPU prewarm on startup."""
    device_hint = os.environ.get("FLOWTORCH_DEVICE", "cpu")
    logger.info(f"Prewarming {device_hint}...")
    prewarm_on_startup(device_hint)

    logger.info("FlowTorch production server ready")
    logger.info(f"- Safety monitoring: ACTIVE")
    logger.info(f"- Signed receipts: ENABLED")
    logger.info(f"- GA-OFF mode: ENFORCED")


@app.get("/health")
async def health():
    """Enhanced health check with detailed status and actual determinism flags."""
    det = {
        "PYTHONHASHSEED": os.environ.get("PYTHONHASHSEED"),
        "CUBLAS_WORKSPACE_CONFIG": os.environ.get("CUBLAS_WORKSPACE_CONFIG"),
        "torch_present": bool(torch),
        "torch_deterministic_algorithms": (bool(torch) and torch.are_deterministic_algorithms_enabled()) if torch else None,
        "cudnn_deterministic": (bool(torch) and getattr(torch.backends.cudnn, "deterministic", None)) if torch else None,
    }
    base_health = health_check()
    base_health["determinism"] = det
    return base_health


@app.get("/metrics", response_model=MetricsResponse)
async def metrics():
    """Prometheus-compatible metrics."""
    metrics_data = get_metrics().get_metrics()
    return metrics_data


@app.get("/receipts")
async def receipts(limit: int = 10):
    """Get recent signed receipts."""
    recent = get_metrics().get_receipts(limit)
    return {"receipts": recent, "count": len(recent)}


@app.get("/cds", response_model=CDSResponse)
async def cds():
    """Control Diagnostics Status - real-time steering health."""
    return cds_status()


@app.get("/prewarm")
async def prewarm_status():
    """Get GPU prewarm status."""
    prewarmer = get_gpu_prewarmer()
    return prewarmer.get_status()


@app.post("/v1/chat/completions")
async def chat_completions(request: ChatRequest, http_request: Request):
    """
    OpenAI-compatible chat completions with basin control.

    Integrated features:
    - Auto-Δ selector (optimal DELTA_BASE per cardinality)
    - Safety monitoring (MI degradation, jitter budget, rate limiting)
    - Authorization (control consent + allowlist)
    - Signed receipts (HMAC-SHA256)
    - Concurrency guard (prevent overload)
    - Graceful degradation (fallback to lower cardinality)
    """
    start_time = time.time()
    metrics = get_metrics()
    safety = get_monitor()

    # Record request
    metrics.record_request()

    try:
        # 1. RATE LIMITING (IP-based token bucket)
        client_ip = http_request.client.host
        if not _allow(client_ip):
            metrics.record_error()
            raise HTTPException(status_code=429, detail="Rate limit exceeded (IP-based)")

        # 2. CONCURRENCY GUARD
        with get_concurrency_guard():

            # 3. SAFETY RATE LIMITING (RPM)
            if not safety.check_rate_limit():
                metrics.record_error()
                raise HTTPException(
                    status_code=429,
                    detail=f"Rate limit exceeded ({safety.rate_limit_rpm} rpm)"
                )

            # 3. CONTROL AUTHORIZATION
            if request.quant:
                # Check if basin control requested
                control_requested = (
                    request.quant.cardinality is not None or
                    request.quant.state_letter is not None or
                    request.quant.phi0_deg is not None
                )

                if control_requested:
                    # CRITICAL: Block phi0_deg without explicit consent
                    if request.quant.phi0_deg is not None and not request.quant.control_consent:
                        metrics.record_error()
                        raise HTTPException(
                            status_code=403,
                            detail="phi0_deg requires control_consent=true"
                        )

                    # Require consent
                    client_ip = http_request.client.host
                    authorized, error_msg = check_control_consent(
                        client_ip,
                        request.quant.control_consent
                    )

                    if not authorized:
                        metrics.record_error()
                        raise HTTPException(status_code=403, detail=error_msg)

                    # CRITICAL: Enforce deterministic settings in basin mode
                    request_temp = request.temperature if request.temperature is not None else 0.7
                    request_top_p = request.top_p if request.top_p is not None else 0.9

                    if request_temp != 0.0:
                        metrics.record_error()
                        raise HTTPException(
                            status_code=400,
                            detail=f"Basin control requires temperature=0.0, got {request_temp}"
                        )
                    if request_top_p != 1.0:
                        metrics.record_error()
                        raise HTTPException(
                            status_code=400,
                            detail=f"Basin control requires top_p=1.0, got {request_top_p}"
                        )

            # 4. AUTO-Δ SELECTION
            phi0_applied = None
            delta_base = None
            cardinality = None
            state_letter = None
            effective_K = 8  # Default cardinality

            if request.quant and request.quant.cardinality:
                cardinality = request.quant.cardinality

                # CRITICAL: Enforce valid cardinality
                valid_K = {2, 4, 8, 16}
                if cardinality not in valid_K:
                    metrics.record_error()
                    raise HTTPException(
                        status_code=400,
                        detail=f"Invalid cardinality {cardinality}; allowed {sorted(valid_K)}"
                    )

                effective_K = cardinality
                state_letter = request.quant.state_letter

                # Audit log (before expensive operations)
                _audit.append({
                    "ts": time.time(),
                    "ip": client_ip,
                    "phi0_deg": request.quant.phi0_deg if request.quant else None,
                    "cardinality": cardinality,
                    "control_consent": (request.quant.control_consent if request.quant else False),
                    "ua": http_request.headers.get("user-agent", "?"),
                })

                if state_letter:
                    # Auto-select PHI0 from state letter
                    phi0_applied, delta_base = auto_select_phi0(state_letter, cardinality)
                    logger.info(f"Auto-Δ: {state_letter} (n={cardinality}) → PHI0={phi0_applied}° (Δ={delta_base}°)")
                elif request.quant.phi0_deg is not None:
                    # Manual PHI0 override
                    phi0_applied = request.quant.phi0_deg
                    delta_base = get_optimal_delta(cardinality)
                    logger.info(f"Manual PHI0: {phi0_applied}° (n={cardinality}, Δ={delta_base}°)")

            # 5. RUN INFERENCE (placeholder - integrate with actual engine)
            # This would call the actual FlowTorch engine with phi0_applied
            # For now, return mock response showing integration

            # Simulate basin measurement
            if phi0_applied is not None:
                # Mock PLL result (in production, this comes from actual inference)
                phi_measured = phi0_applied  # Perfect lock
                R_spectral = 0.28  # Typical locking strength

                # Calculate MI
                theoretical_mi = {2: 1.0, 3: 1.585, 4: 2.0, 8: 3.0}
                mi_bits = theoretical_mi.get(cardinality, 0.0)

                # Mock energy (in production, measured during inference)
                energy_uJ = 1250.0  # Typical CPU inference

                # 6. SAFETY MONITORING
                safety_ok = safety.record_mi(mi_bits, theoretical_mi[cardinality])
                jitter_ok, jitter = safety.record_phi(phi_measured)

                if not safety_ok or not jitter_ok:
                    # Graceful degradation
                    new_card, degrade_msg = graceful_degrade(cardinality, safety.degradation_reason)
                    logger.warning(degrade_msg)

                    if new_card != cardinality:
                        # Recalculate with degraded cardinality
                        cardinality = new_card
                        if state_letter:
                            # Remap state letter to new cardinality
                            # (in production, you'd have a mapping strategy)
                            state_letter = 'a'  # Default to first state
                            phi0_applied, delta_base = auto_select_phi0(state_letter, cardinality)

                # 7. CREATE SIGNED RECEIPT
                receipt = create_signed_receipt(
                    phi=phi_measured,
                    R=R_spectral,
                    MI=mi_bits,
                    energy_uJ=energy_uJ,
                    cardinality=cardinality,
                    state=state_letter or "manual",
                    delta_base=delta_base,
                    phi0_applied=phi0_applied,
                )

                # Record metrics
                latency = time.time() - start_time
                metrics.record_basin_lock(mi_bits, latency, energy_uJ, cardinality, receipt)

            else:
                # No basin control requested, standard inference
                receipt = None

            # 8. RUN INFERENCE (deterministic decode with PHI0)
            prompt = request.messages[-1].content if request.messages else ""
            max_tokens = request.max_tokens or 128

            if phi0_applied is not None:
                # Basin control: deterministic decode with PHI0
                content = decode_with_phi0_deg(prompt, phi0_applied, effective_K, max_new_tokens=max_tokens)
                logger.info(f"Basin decode: φ={phi0_applied}° (K={effective_K}) → {len(content)} chars")
            else:
                # Standard inference path
                temperature = request.temperature or 0.7
                top_p = request.top_p or 0.9
                content = backend_decode(prompt, temperature=temperature, top_p=top_p, max_new_tokens=max_tokens)
                logger.info(f"Standard decode: T={temperature} → {len(content)} chars")

            # 9. BUILD RESPONSE
            response = {
                "id": f"flowtorch-{int(time.time())}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": request.model or "flowtorch/basin-control",
                "choices": [{
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": content
                    },
                    "finish_reason": "stop"
                }],
                "usage": {
                    "prompt_tokens": len(request.messages[0].content.split()) if request.messages else 0,
                    "completion_tokens": len(content.split()),
                    "total_tokens": 0
                },
            }

            # Add receipt to response
            if receipt:
                response["receipt"] = receipt
                response["basin_control"] = {
                    "cardinality": effective_K,
                    "state": state_letter,
                    "phi0_applied": phi0_applied,
                    "delta_base": delta_base,
                    "safety_status": safety.get_status()["status"],
                }

            # Echo effective basin config (self-describing responses)
            if phi0_applied is not None:
                response["echo"] = {
                    "basin": {
                        "phi0_deg": phi0_applied,
                        "cardinality": effective_K
                    }
                }

            # CDR state tracking (for reversible inference exploration)
            if CDR_ENABLED:
                mgr = get_cdr_manager()
                if mgr is not None:
                    response["cdr"] = {
                        "enabled": True,
                        "step_count": mgr.cdr_state.step_count if mgr.cdr_state.initialized else 0,
                        "state_hash": mgr.cdr_state_hash()[:16],
                        "checkpoints": len(mgr.checkpoints),
                    }

            return response

    except HTTPException:
        raise
    except Exception as e:
        metrics.record_error()
        logger.error(f"Request failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/basin/audit_tail")
async def audit_tail(n: int = 50):
    """Get last N basin control audit log entries (redacted, no PII)."""
    return {"audit": list(list(_audit)[-min(max(n, 1), 500):]), "count": len(_audit)}


@app.get("/cdr/status")
async def cdr_status():
    """
    CDR (Classical Reversible Computing) state tracking status.

    Shows:
    - CDR step count (absorptions performed)
    - State hash (current sponge state fingerprint)
    - Checkpoint count (named checkpoints for random access)
    - Memory usage (O(1) VM memory, external tape for messages)
    """
    if not CDR_ENABLED:
        return {"enabled": False, "message": "CDR disabled (BASIN_CDR_CHECKPOINTS=0)"}

    mgr = get_cdr_manager()
    if mgr is None:
        return {"enabled": False, "message": "CDR manager unavailable (helix_cdc not installed)"}

    stats = mgr.get_stats()
    return {
        "enabled": True,
        "cdr_state": {
            "initialized": mgr.cdr_state.initialized,
            "step_count": mgr.cdr_state.step_count if mgr.cdr_state.initialized else 0,
            "state_hash": mgr.cdr_state_hash()[:16] + "...",
            "vm_memory_bytes": 200,  # O(1) constant
        },
        "checkpoints": {
            "count": len(mgr.checkpoints),
            "max": mgr.max_checkpoints,
            "ids": list(mgr.checkpoints.keys())[-10:],  # Last 10 checkpoint IDs
        },
        "stats": stats,
        "memory_model": "O(1) VM + external tape (caller tracks messages for LIFO rollback)",
    }


@app.post("/cdr/rollback")
async def cdr_rollback(checkpoint_id: str):
    """
    Rollback CDR state to a previous checkpoint.

    NOTE: Requires LIFO order for pure CDR mode (step_back).
    Named checkpoints allow random access but discard later checkpoints.
    """
    if not CDR_ENABLED:
        raise HTTPException(status_code=400, detail="CDR disabled")

    mgr = get_cdr_manager()
    if mgr is None:
        raise HTTPException(status_code=500, detail="CDR manager unavailable")

    state = mgr.rollback(checkpoint_id)
    if state is None:
        raise HTTPException(status_code=404, detail=f"Checkpoint '{checkpoint_id}' not found")

    return {
        "rolled_back_to": checkpoint_id,
        "restored_state": state,
        "current_cdr_hash": mgr.cdr_state_hash()[:16],
    }


@app.get("/")
async def root():
    """Root endpoint with API info."""
    return {
        "service": "FlowTorch Production API",
        "version": "v1.0-basin-control",
        "features": [
            "Multi-state basin control (2/3/4/8 states)",
            "Auto-Δ selector (optimal DELTA_BASE)",
            "Safety monitoring (MI/jitter/rate)",
            "Signed receipts (HMAC-SHA256)",
            "Graceful degradation",
            "GA-OFF enforcement",
            "CDR reversible state tracking (O(1) memory)",
        ],
        "endpoints": {
            "chat": "/v1/chat/completions",
            "health": "/health",
            "metrics": "/metrics",
            "receipts": "/receipts",
            "cds": "/cds (Control Diagnostics Status)",
            "cdr_status": "/cdr/status (CDR state tracking)",
            "cdr_rollback": "/cdr/rollback (Rollback to checkpoint)",
        },
        "docs": "/docs",
    }


if __name__ == "__main__":
    # Get config from environment
    host = os.getenv("FLOWTORCH_HOST", "127.0.0.1")
    port = int(os.getenv("FLOWTORCH_PORT", "8089"))

    logger.info(f"Starting FlowTorch production server on {host}:{port}")

    uvicorn.run(
        app,
        host=host,
        port=port,
        log_level="info",
    )
