#!/usr/bin/env python3
"""
BEASTBULLET Mobile API Server
Real-time hallucination detection with logging
"""

from flask import Flask, request, jsonify
import numpy as np
from pathlib import Path
from datetime import datetime
import json


class ValidatorExpert:
    """Validator expert using only NumPy"""
    
    def __init__(self, weights_path: str = "models/validator_weights.npz"):
        weights_path = Path(weights_path)
        if not weights_path.exists():
            raise FileNotFoundError(f"Weights not found: {weights_path}")
        
        data = np.load(weights_path)
        self.weights = {key: data[key] for key in data.files}
    
    def sigmoid(self, x: np.ndarray) -> np.ndarray:
        return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))
    
    def layer_norm(self, x: np.ndarray, weight: np.ndarray, bias: np.ndarray, eps: float = 1e-5) -> np.ndarray:
        mean = np.mean(x, axis=-1, keepdims=True)
        var = np.var(x, axis=-1, keepdims=True)
        x_norm = (x - mean) / np.sqrt(var + eps)
        return x_norm * weight + bias
    
    def linear(self, x: np.ndarray, weight: np.ndarray, bias: np.ndarray) -> np.ndarray:
        return np.dot(x, weight.T) + bias
    
    def forward(self, embedding: np.ndarray):
        if embedding.ndim == 1:
            embedding = embedding.reshape(1, -1)
        
        x = embedding
        
        # Auto-detect weight naming scheme
        # Original model uses: encoder.0, encoder.1, encoder.4, encoder.5
        # Patched model uses: encoder.0, encoder.2, encoder.4, encoder.6
        if 'encoder.1.weight' in self.weights:
            # Original model
            norm1_w, norm1_b = 'encoder.1.weight', 'encoder.1.bias'
            norm2_w, norm2_b = 'encoder.5.weight', 'encoder.5.bias'
        else:
            # Patched model
            norm1_w, norm1_b = 'encoder.2.weight', 'encoder.2.bias'
            norm2_w, norm2_b = 'encoder.6.weight', 'encoder.6.bias'
        
        # Encoder
        x = self.linear(x, self.weights['encoder.0.weight'], self.weights['encoder.0.bias'])
        x = np.maximum(0, x)
        x = self.layer_norm(x, self.weights[norm1_w], self.weights[norm1_b])
        x = self.linear(x, self.weights['encoder.4.weight'], self.weights['encoder.4.bias'])
        x = np.maximum(0, x)
        x = self.layer_norm(x, self.weights[norm2_w], self.weights[norm2_b])
        
        # Validator head
        val = self.linear(x, self.weights['validator.0.weight'], self.weights['validator.0.bias'])
        val = np.maximum(0, val)
        verdict = self.linear(val, self.weights['validator.2.weight'], self.weights['validator.2.bias'])
        
        # Confidence
        conf = self.linear(x, self.weights['confidence_adjuster.0.weight'], self.weights['confidence_adjuster.0.bias'])
        conf = np.maximum(0, conf)
        confidence = self.linear(conf, self.weights['confidence_adjuster.2.weight'], self.weights['confidence_adjuster.2.bias'])
        confidence = self.sigmoid(confidence)
        
        return verdict, confidence
    
    def predict(self, embedding: np.ndarray):
        logits, conf = self.forward(embedding)
        pred_class = np.argmax(logits, axis=-1)[0]
        classes = ["VALID", "INVALID", "UNCERTAIN"]
        return classes[pred_class], float(conf[0, 0])


# Initialize Flask app
app = Flask(__name__)

# Load validator
print("🚀 Loading BEASTBULLET Validator...")
validator = ValidatorExpert()
print("✓ Validator loaded!")

# Log file
log_file = Path("validator_log.jsonl")


def log_validation(request_data, verdict, confidence, is_hallucination):
    """Log validation to file"""
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "specialist": request_data.get("specialist", "unknown"),
        "statement": request_data.get("statement", ""),
        "verdict": verdict,
        "confidence": confidence,
        "is_hallucination": is_hallucination,
        "embedding_shape": request_data.get("embedding_shape", None)
    }
    
    with open(log_file, 'a') as f:
        f.write(json.dumps(log_entry) + '\n')


@app.route('/health', methods=['GET'])
def health():
    """Health check endpoint"""
    return jsonify({
        'status': 'healthy',
        'service': 'beastbullet-validator',
        'version': '1.0.0'
    })


@app.route('/validate', methods=['POST'])
def validate():
    """
    Validate a statement embedding
    
    Request:
    {
        "embedding": [384-dim array],
        "specialist": "code_expert",
        "statement": "original statement text (optional)"
    }
    
    Response:
    {
        "verdict": "VALID|INVALID|UNCERTAIN",
        "confidence": 0.0-1.0,
        "is_hallucination": true/false,
        "latency_ms": 0.27
    }
    """
    import time
    start = time.time()
    
    try:
        data = request.json
        
        # Get embedding
        embedding = np.array(data['embedding'], dtype=np.float32)
        
        # Validate
        verdict, confidence = validator.predict(embedding)
        
        # Determine if hallucination (INVALID with high confidence)
        is_hallucination = verdict == "INVALID" and confidence < 0.5
        
        latency_ms = (time.time() - start) * 1000
        
        # Log the validation
        log_validation(
            request_data={
                "specialist": data.get("specialist", "unknown"),
                "statement": data.get("statement", ""),
                "embedding_shape": embedding.shape
            },
            verdict=verdict,
            confidence=confidence,
            is_hallucination=is_hallucination
        )
        
        # Print to console
        emoji = "🚨" if is_hallucination else "✅"
        print(f"{emoji} [{data.get('specialist', 'unknown')}] {verdict} (conf: {confidence:.3f})")
        
        return jsonify({
            'verdict': verdict,
            'confidence': confidence,
            'is_hallucination': is_hallucination,
            'latency_ms': round(latency_ms, 2)
        })
    
    except Exception as e:
        return jsonify({
            'error': str(e)
        }), 400


@app.route('/stats', methods=['GET'])
def stats():
    """Get validation statistics"""
    if not log_file.exists():
        return jsonify({'total': 0, 'hallucinations': 0})
    
    total = 0
    hallucinations = 0
    
    with open(log_file, 'r') as f:
        for line in f:
            entry = json.loads(line)
            total += 1
            if entry['is_hallucination']:
                hallucinations += 1
    
    return jsonify({
        'total_validations': total,
        'hallucinations_caught': hallucinations,
        'catch_rate': f"{100*hallucinations/total:.1f}%" if total > 0 else "0%"
    })


if __name__ == '__main__':
    print("\n" + "=" * 60)
    print("🚀 BEASTBULLET Validator API Server")
    print("=" * 60)
    print("\nEndpoints:")
    print("  GET  /health     - Health check")
    print("  POST /validate   - Validate embedding")
    print("  GET  /stats      - Get statistics")
    print("\nLog file:", log_file.absolute())
    print("=" * 60)
    print("\nStarting server on 0.0.0.0:8080...")
    print("Press Ctrl+C to stop\n")
    
    app.run(host='0.0.0.0', port=8080, debug=False)