#!/usr/bin/env python3
"""
Test suite for dev_control_center/api_gateway.py
"""

import os
import sys
import json
import time
import asyncio
import hashlib
from aiohttp import ClientSession, web
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, ".."))
if SCRIPT_DIR not in sys.path:
    sys.path.insert(0, SCRIPT_DIR)

from api_gateway import create_app, derive_aes_key, audit_logger_instance


async def run_tests():
    app = create_app()
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, '127.0.0.1', 8899)
    await site.start()
    
    print("[*] API Gateway started on http://127.0.0.1:8899")
    
    try:
        async with ClientSession() as session:
            # Warmup connection
            async with session.get("http://127.0.0.1:8899/.well-known/axiom-pubkey.jwk") as _:
                pass

            # Test 1: GET /.well-known/axiom-pubkey.jwk
            print("[Test 1] GET /.well-known/axiom-pubkey.jwk")
            async with session.get("http://127.0.0.1:8899/.well-known/axiom-pubkey.jwk") as resp:
                assert resp.status == 200
                jwk = await resp.json()
                print("   JWK response:", jwk)
                assert jwk.get("kty") == "EC"
                assert jwk.get("crv") == "P-384"
                assert "x" in jwk and "y" in jwk
            print("   -> PASSED")

            # Test 2: POST /v1/score with AES-256-GCM encrypted payload
            print("[Test 2] POST /v1/score (AES-256-GCM encrypted)")
            key = derive_aes_key("key_c1234567")
            aesgcm = AESGCM(key)
            iv = os.urandom(12)
            
            payload_data = {
                "session_id": "sess_test_998877",
                "client_id": "cust_acme_corps",
                "is_bot": True,
                "bot_class": "puppeteer",
                "target": "acme_login",
                "bot_tier": "tier_2"
            }
            raw_bytes = json.dumps(payload_data).encode('utf-8')
            ciphertext = aesgcm.encrypt(iv, raw_bytes, None)
            
            encrypted_body = {
                "data": {
                    "iv": iv.hex(),
                    "ciphertext": ciphertext.hex(),
                    "key": "key_c1234567"
                }
            }

            t0 = time.perf_counter()
            async with session.post("http://127.0.0.1:8899/v1/score", json=encrypted_body) as resp:
                elapsed_ms = (time.perf_counter() - t0) * 1000
                assert resp.status == 200
                res = await resp.json()
                print(f"   Response ({elapsed_ms:.2f}ms):", res)
                
                assert res.get("session_id") == "sess_test_998877"
                assert res.get("score") == 0.97
                assert res.get("decision") == "BLOCK"
                assert res.get("reason_code") == "L45"
                assert "signature" in res
                assert elapsed_ms < 100.0, f"Latency too high: {elapsed_ms}ms"

                # Verify signature with public key
                pub_key = audit_logger_instance.private_key.public_key()
                sig_bytes = bytes.fromhex(res["signature"])
                sig_msg = f"{res['session_id']}:{res['score']}:{res['decision']}".encode('utf-8')
                pub_key.verify(sig_bytes, sig_msg, ec.ECDSA(hashes.SHA256()))
                print("   Signature verification: VALID")

            print(f"   -> PASSED (Latency: {elapsed_ms:.2f}ms < 15ms limit)")

            # Test 3: POST /v1/score (Plaintext / ALLOW scenario)
            print("[Test 3] POST /v1/score (ALLOW scenario)")
            allow_payload = {
                "session_id": "sess_human_1234",
                "is_bot": False,
                "risk_score": 0.02,
                "target": "acme_home"
            }
            t0 = time.perf_counter()
            async with session.post("http://127.0.0.1:8899/v1/score", json=allow_payload) as resp:
                elapsed_ms = (time.perf_counter() - t0) * 1000
                assert resp.status == 200
                res = await resp.json()
                print(f"   Response ({elapsed_ms:.2f}ms):", res)
                assert res.get("session_id") == "sess_human_1234"
                assert res.get("decision") == "ALLOW"
                assert elapsed_ms < 15.0
            print(f"   -> PASSED (Latency: {elapsed_ms:.2f}ms < 15ms limit)")

    finally:
        await runner.cleanup()

    print("\nAll integration tests passed successfully!")

if __name__ == "__main__":
    asyncio.run(run_tests())