#!/usr/bin/env python3
"""
Axiom Zero - Detection Engine Self-Validation Suite
Validates classification accuracy against the known fingerprint database.
No external network calls or browser automation required.
"""

import sys
import json
import time
from pathlib import Path

# Add dev_control_center to path
sys.path.insert(0, str(Path(__file__).parent))

from ja4_fingerprinter import compute_ja4, check_mismatch, KNOWN_FINGERPRINTS, JA4Cache
from customer_db import CustomerDatabase
from secure_swarm_protocol import SecureSwarmProtocol


def test_ja4_classification_accuracy() -> dict:
    """
    Test that our JA4 fingerprinter correctly classifies every entry
    in the KNOWN_FINGERPRINTS database.
    Returns accuracy metrics.
    """
    fps = list(KNOWN_FINGERPRINTS.values()) if isinstance(KNOWN_FINGERPRINTS, dict) else KNOWN_FINGERPRINTS
    total = len(fps)
    correct = 0
    results = []
    
    for fp in fps:
        name = fp['name']
        ja4 = fp['ja4']
        expected_type = fp['type']  # 'BROWSER', 'AUTOMATION', 'SCRIPTED'
        
        # Test: does check_mismatch() correctly identify mismatches
        # when a bot UA claims to be Chrome?
        if expected_type in ('AUTOMATION', 'SCRIPTED'):
            # These should be flagged when paired with a Chrome UA
            result = check_mismatch('Mozilla/5.0 Chrome/131.0.0.0 Safari/537.36', ja4)
            classified_correctly = result['risk_level'] in ('HIGH', 'CRITICAL')
        else:
            # BROWSER type should not be flagged when using matching browser UA
            ua = 'Mozilla/5.0 (X11; Linux x86_64; rv:121.0) Gecko/20100101 Firefox/121.0' if 'firefox' in name.lower() else 'Mozilla/5.0 Chrome/131.0.0.0 Safari/537.36'
            result = check_mismatch(ua, ja4)
            classified_correctly = result['risk_level'] in ('NONE', 'LOW')
        
        if classified_correctly:
            correct += 1
        
        results.append({
            'name': name,
            'type': expected_type,
            'risk_level': result['risk_level'],
            'correct': classified_correctly
        })
    
    accuracy = correct / total if total > 0 else 0
    return {
        'test': 'JA4 Classification Accuracy',
        'total': total,
        'correct': correct,
        'accuracy': accuracy,
        'results': results
    }


def test_hmac_credential_verification() -> dict:
    """
    Test that the HMAC-SHA256 credential verification correctly:
    1. Accepts valid credentials
    2. Rejects invalid credentials 
    3. Rejects credentials that are almost correct (timing-safe)
    """
    from customer_db import _db
    import os
    
    assert hasattr(_db, 'verify_credential_hash'), 'API method missing'

    secret = getattr(_db, 'secret', None)
    if secret is None:
        from customer_db import AXIOM_SECRET
        secret = AXIOM_SECRET
    
    tests = []
    
    # Test 1: Valid credential round-trip
    test_hash = _db.hash_credential('test_password_123', secret)
    is_valid = _db.verify_credential_hash(test_hash, 'test_password_123', secret)
    tests.append({'name': 'Valid credential accepted', 'passed': is_valid})
    
    # Test 2: Wrong password rejected
    is_invalid = not _db.verify_credential_hash(test_hash, 'wrong_password', secret)
    tests.append({'name': 'Wrong password rejected', 'passed': is_invalid})
    
    # Test 3: Empty string rejected
    is_empty_rejected = not _db.verify_credential_hash(test_hash, '', secret)
    tests.append({'name': 'Empty password rejected', 'passed': is_empty_rejected})
    
    passed = sum(1 for t in tests if t['passed'])
    return {
        'test': 'HMAC Credential Verification',
        'total': len(tests),
        'passed': passed,
        'accuracy': passed / len(tests),
        'results': tests
    }


def test_customer_db_persistence() -> dict:
    """
    Test that customer DB correctly stores and retrieves records.
    """
    from customer_db import _db
    import json
    
    tests = []
    
    # Test 1: Can retrieve all customers
    customers = _db.get_all_customers()
    tests.append({'name': 'get_all_customers returns list', 'passed': isinstance(customers, list)})
    tests.append({'name': 'At least 2 customers seeded', 'passed': len(customers) >= 2})
    
    # Test 2: Can retrieve specific customer
    if customers:
        cust_id = customers[0]['id']
        cust = _db.get_customer(cust_id)
        tests.append({'name': 'get_customer returns dict', 'passed': isinstance(cust, dict)})
        tests.append({'name': 'Customer has company_name', 'passed': 'company_name' in cust})
    
    # Test 3: Log a detection event
    _db.log_detection_event(
        customer_id='test_validation',
        event_type='EVENT_BLOCK',
        ip_address='10.0.0.1',
        ja4_hash='t13d190900_abcdef123456_xyz789012345',
        bot_class='TIER_1_SCRIPTED',
        outcome='BLOCKED',
        latency_ms=1.2
    )
    events = []
    for _ in range(3):
        time.sleep(0.5)  # Let the queue writer process
        events = _db.get_recent_events(limit=5)
        if len(events) >= 1:
            break
    tests.append({'name': 'Detection event logged and retrievable', 'passed': len(events) >= 1})
    
    passed = sum(1 for t in tests if t['passed'])
    return {
        'test': 'Customer DB Persistence',
        'total': len(tests),
        'passed': passed,
        'accuracy': passed / len(tests),
        'results': tests
    }


def test_socket_protocol_framing() -> dict:
    """
    Test the 4-byte length-prefix message framing in SocketProtocol
    using a loopback socket pair.
    """
    import socket
    import struct
    import threading
    from secure_swarm_protocol import SocketProtocol
    
    tests = []
    received_frames = []
    
    # Create a socketpair for testing
    server_sock, client_sock = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM)
    
    # Server thread: receive one frame
    def receive_one():
        try:
            frame = SocketProtocol.recv_frame(server_sock)
            received_frames.append(frame)
        except Exception as e:
            received_frames.append({'error': str(e)})
        finally:
            server_sock.close()
    
    t = threading.Thread(target=receive_one, daemon=True)
    t.start()
    
    # Client: send a test frame
    test_event = {
        'type': 'EVENT_BLOCK',
        'timestamp': '2026-08-04T19:00:00Z',
        'client_id': 'test_node_01',
        'payload': {'ip': '1.2.3.4', 'bot_class': 'TIER_1', 'latency_ms': 1.5}
    }
    
    try:
        SocketProtocol.send_frame(client_sock, test_event)
        client_sock.close()
        t.join(timeout=3.0)
        
        if t.is_alive():
            tests.append({'name': 'Frame roundtrip timed out', 'passed': False})
        elif received_frames and 'error' not in received_frames[0]:
            received = received_frames[0]
            tests.append({'name': 'Frame sent and received', 'passed': True})
            tests.append({'name': 'Event type preserved', 'passed': received.get('type') == 'EVENT_BLOCK'})
            tests.append({'name': 'Payload preserved', 'passed': received.get('payload', {}).get('ip') == '1.2.3.4'})
            tests.append({'name': 'client_id preserved', 'passed': received.get('client_id') == 'test_node_01'})
        else:
            tests.append({'name': 'Frame roundtrip failed', 'passed': False})
    except Exception as e:
        tests.append({'name': f'Socket test error: {e}', 'passed': False})
    
    passed = sum(1 for t in tests if t['passed'])
    return {
        'test': 'Socket Protocol Framing',
        'total': len(tests),
        'passed': passed,
        'accuracy': passed / len(tests),
        'results': tests
    }


def run_all_validations() -> dict:
    print('=' * 60)
    print('  AXIOM ZERO — DETECTION ENGINE SELF-VALIDATION')
    print('=' * 60)
    
    all_results = []
    
    tests_to_run = [
        ('JA4 Classification', test_ja4_classification_accuracy),
        ('HMAC Credentials', test_hmac_credential_verification),
        ('Customer DB', test_customer_db_persistence),
        ('Socket Framing', test_socket_protocol_framing),
    ]
    
    for test_name, test_fn in tests_to_run:
        print(f'\n[*] Running: {test_name}...')
        try:
            result = test_fn()
            all_results.append(result)
            accuracy = result.get('accuracy', result.get('passed', 0) / max(result.get('total', 1), 1))
            status = '✅ PASS' if accuracy >= 0.8 else '❌ FAIL'
            print(f'    {status} — Accuracy: {accuracy:.1%}')
        except Exception as e:
            print(f'    ❌ ERROR: {e}')
            all_results.append({'test': test_name, 'error': str(e), 'accuracy': 0})
    
    overall = sum(r.get('accuracy', 0) for r in all_results) / len(all_results)
    
    print(f'\n{"=" * 60}')
    print(f'  OVERALL VALIDATION SCORE: {overall:.1%}')
    print(f'{"=" * 60}\n')
    
    report = {
        'timestamp': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
        'overall_score': overall,
        'tests': all_results
    }
    
    report_path = Path(__file__).parent / 'validation_report.json'
    with open(report_path, 'w') as f:
        json.dump(report, f, indent=2)
    print(f'[*] Validation report saved: {report_path}')
    
    return report


if __name__ == '__main__':
    run_all_validations()