#!/usr/bin/env python3
"""
Axiom Zero — API Gateway Server
Built with aiohttp for sub-15ms telemetry scoring, AES-256-GCM payload decryption,
Unix Domain Socket IPC integration (/tmp/axiom_zero_ipc.sock), and cryptographic audit logging.
"""

import os
import sys
import time
import json
import uuid
import struct
import base64
import asyncio
import hashlib
import logging
from aiohttp import web
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import ec

# Setup directory paths
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, ".."))
AUDIT_LOGGER_DIR = os.path.join(PROJECT_ROOT, "gauntlet_lab", "shared", "audit_logger")

if AUDIT_LOGGER_DIR not in sys.path:
    sys.path.insert(0, AUDIT_LOGGER_DIR)

from audit_logger import AuditLogger

# Global constants & defaults
SOCKET_PATH = "/tmp/axiom_zero_ipc.sock"
DEFAULT_MASTER_KEY = "key_c1234567"
SECRET_KEY_ENV = os.environ.get("AES_KEY") or os.environ.get("option_key") or DEFAULT_MASTER_KEY

# Setup logging
logging.basicConfig(level=logging.INFO, format="[API Gateway] %(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("api_gateway")

# Initialize AuditLogger safely with absolute db_path & key location
abs_db_path = os.path.join(AUDIT_LOGGER_DIR, "audit_log.db")
orig_cwd = os.getcwd()
try:
    os.chdir(AUDIT_LOGGER_DIR)
    audit_logger_instance = AuditLogger(db_path=abs_db_path)
finally:
    os.chdir(orig_cwd)


def get_jwk() -> dict:
    """Return JWK representation of AuditLogger's public key."""
    public_key = audit_logger_instance.private_key.public_key()
    pub_num = public_key.public_numbers()
    
    key_size = public_key.key_size
    coord_len = (key_size + 7) // 8
    
    x_bytes = pub_num.x.to_bytes(coord_len, byteorder='big')
    y_bytes = pub_num.y.to_bytes(coord_len, byteorder='big')
    
    x_b64 = base64.urlsafe_b64encode(x_bytes).rstrip(b'=').decode('utf-8')
    y_b64 = base64.urlsafe_b64encode(y_bytes).rstrip(b'=').decode('utf-8')
    
    crv_name = "P-384" if key_size == 384 else ("P-256" if key_size == 256 else "P-521")
    alg_name = "ES384" if key_size == 384 else ("ES256" if key_size == 256 else "ES512")
    
    jwk_dict = {
        "kty": "EC",
        "crv": crv_name,
        "x": x_b64,
        "y": y_b64,
        "use": "sig",
        "alg": alg_name,
        "kid": "axiom-pubkey-1"
    }
    
    return {
        **jwk_dict,
        "keys": [jwk_dict]
    }


def parse_bytes_field(val) -> bytes:
    """Parse hex or base64 or bytes representation into bytes."""
    if isinstance(val, bytes):
        return val
    if not isinstance(val, str):
        raise ValueError(f"Expected string or bytes, got {type(val)}")
    
    val_clean = val.strip()
    try:
        if len(val_clean) % 2 == 0 and all(c in '0123456789abcdefABCDEF' for c in val_clean):
            return bytes.fromhex(val_clean)
    except Exception:
        pass
    
    try:
        padded = val_clean + "=" * ((4 - len(val_clean) % 4) % 4)
        return base64.b64decode(padded)
    except Exception:
        pass

    try:
        return base64.urlsafe_b64decode(val_clean + "=" * ((4 - len(val_clean) % 4) % 4))
    except Exception:
        pass

    return val.encode('utf-8')


def derive_aes_key(key_input) -> bytes:
    """Derive 32-byte key for AES-256-GCM."""
    if isinstance(key_input, bytes):
        if len(key_input) == 32:
            return key_input
        return hashlib.sha256(key_input).digest()
    
    if isinstance(key_input, str):
        if len(key_input) == 64 and all(c in '0123456789abcdefABCDEF' for c in key_input):
            return bytes.fromhex(key_input)
        
        try:
            raw = base64.b64decode(key_input)
            if len(raw) == 32:
                return raw
        except Exception:
            pass
        
        return hashlib.sha256(key_input.encode('utf-8')).digest()
    
    return hashlib.sha256(str(key_input).encode('utf-8')).digest()


def decrypt_aes_gcm_payload(req_data: dict) -> dict:
    """
    Decrypt incoming AES-256-GCM payload.
    Supports data.iv, data.ciphertext (and optional data.tag, data.key).
    """
    data_container = req_data.get("data") if isinstance(req_data.get("data"), dict) else req_data
    
    iv_val = data_container.get("iv") or req_data.get("iv")
    ciphertext_val = data_container.get("ciphertext") or req_data.get("ciphertext")
    
    if not iv_val or not ciphertext_val:
        return data_container
    
    iv_bytes = parse_bytes_field(iv_val)
    ciphertext_bytes = parse_bytes_field(ciphertext_val)
    
    tag_val = data_container.get("tag") or req_data.get("tag")
    if tag_val:
        tag_bytes = parse_bytes_field(tag_val)
        if not ciphertext_bytes.endswith(tag_bytes):
            ciphertext_bytes = ciphertext_bytes + tag_bytes

    key_val = (data_container.get("key") or data_container.get("aes_key") or 
               req_data.get("key") or SECRET_KEY_ENV)
    
    key_bytes = derive_aes_key(key_val)
    
    aesgcm = AESGCM(key_bytes)
    decrypted_raw = aesgcm.decrypt(iv_bytes, ciphertext_bytes, None)
    
    return json.loads(decrypted_raw.decode('utf-8'))


async def query_daemon_ipc(payload: dict, session_id: str, socket_path: str = SOCKET_PATH, timeout: float = 0.005) -> dict:
    """Connect to /tmp/axiom_zero_ipc.sock using 4-byte length-prefixed framing."""
    if not os.path.exists(socket_path):
        return {}

    ipc_frame = {
        "type": payload.get("type", "EVENT_SCORE"),
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "client_id": session_id,
        "session_id": session_id,
        "payload": payload
    }
    
    payload_bytes = json.dumps(ipc_frame).encode('utf-8')
    header = struct.pack('>I', len(payload_bytes))
    
    try:
        reader, writer = await asyncio.wait_for(
            asyncio.open_unix_connection(socket_path),
            timeout=timeout
        )
        writer.write(header + payload_bytes)
        await writer.drain()
        
        length_bytes = await asyncio.wait_for(reader.readexactly(4), timeout=timeout)
        if length_bytes and len(length_bytes) == 4:
            msg_len = struct.unpack('>I', length_bytes)[0]
            resp_bytes = await asyncio.wait_for(reader.readexactly(msg_len), timeout=timeout)
            resp_data = json.loads(resp_bytes.decode('utf-8'))
            writer.close()
            await writer.wait_closed()
            return resp_data
            
        writer.close()
        await writer.wait_closed()
    except Exception:
        pass
        
    return {}


def create_signature(session_id: str, score: float, decision: str) -> str:
    """Create ECDSA P-384 signature over session_id:score:decision."""
    sig_payload = f"{session_id}:{score}:{decision}".encode('utf-8')
    signature_bytes = audit_logger_instance.private_key.sign(
        sig_payload,
        ec.ECDSA(hashes.SHA256())
    )
    return signature_bytes.hex()


async def handle_jwk(request: web.Request) -> web.Response:
    """GET /.well-known/axiom-pubkey.jwk endpoint."""
    jwk_data = get_jwk()
    return web.json_response(jwk_data)


async def handle_score(request: web.Request) -> web.Response:
    """POST /v1/score endpoint."""
    t0 = time.perf_counter()
    
    try:
        req_json = await request.json()
    except Exception:
        req_json = {}

    # 1. Decrypt incoming AES-256-GCM payload if encrypted
    try:
        decrypted_payload = decrypt_aes_gcm_payload(req_json)
    except Exception as e:
        logger.warning(f"Payload decryption failed: {e}")
        decrypted_payload = req_json

    # 2. Extract or generate session_id
    session_id = (decrypted_payload.get("session_id") or 
                  decrypted_payload.get("client_id") or 
                  req_json.get("session_id") or 
                  f"sess_{uuid.uuid4().hex[:16]}")
    
    # 3. Query daemon IPC socket
    daemon_resp = await query_daemon_ipc(decrypted_payload, session_id)
    
    # 4. Compute score, decision, and reason_code
    if daemon_resp and "score" in daemon_resp:
        score = float(daemon_resp["score"])
        decision = daemon_resp.get("decision", "BLOCK" if score >= 0.70 else "ALLOW")
        reason_code = daemon_resp.get("reason_code", "L45" if decision == "BLOCK" else "L00")
    elif "score" in decrypted_payload:
        score = float(decrypted_payload["score"])
        decision = decrypted_payload.get("decision", "BLOCK" if score >= 0.70 else "ALLOW")
        reason_code = decrypted_payload.get("reason_code", "L45" if decision == "BLOCK" else "L00")
    else:
        bot_tier_name = str(decrypted_payload.get("bot_tier") or req_json.get("bot_tier") or "")
        is_bot = (decrypted_payload.get("is_bot") or req_json.get("is_bot") or
                  decrypted_payload.get("cdp_detected") or 
                  decrypted_payload.get("ja4_mismatch") or 
                  decrypted_payload.get("bot_class") or
                  bot_tier_name.startswith("Tier") or
                  decrypted_payload.get("signals", {}).get("automation", {}).get("webdriver", False))
        if is_bot:
            score = 0.97
            decision = "BLOCK"
            reason_code = decrypted_payload.get("reason_code", "L45")
        else:
            score = float(decrypted_payload.get("risk_score", 0.02))
            decision = "BLOCK" if score >= 0.70 else "ALLOW"
            reason_code = decrypted_payload.get("reason_code", "L45" if decision == "BLOCK" else "L00")

    # 5. Create digital signature asynchronously in thread pool
    signature = await asyncio.to_thread(create_signature, session_id, score, decision)
    
    latency_ms = round((time.perf_counter() - t0) * 1000, 2)
    timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    target = decrypted_payload.get("target") or decrypted_payload.get("customer_id") or "api_gateway"
    bot_tier = decrypted_payload.get("bot_tier") or ("tier_2" if decision == "BLOCK" else "tier_1")
    
    # 6. Log transaction asynchronously to AuditLogger
    asyncio.create_task(
        asyncio.to_thread(
            audit_logger_instance.log_event,
            timestamp,
            session_id,
            target,
            bot_tier,
            decision,
            latency_ms
        )
    )

    # 7. Construct JSON response
    response_body = {
        "session_id": session_id,
        "score": score,
        "decision": decision,
        "reason_code": reason_code,
        "signature": signature
    }

    return web.json_response(response_body)


async def handle_pi5_telemetry(request: web.Request) -> web.Response:
    try:
        import pi5_resource_monitor as mon
        import datetime
        # Use cached Monitor for accurate net/cpu delta tracking between polls
        if '_pi5_monitor' not in request.app:
            request.app['_pi5_monitor'] = mon.Monitor()
        monitor = request.app['_pi5_monitor']
        monitor.update()
        s = monitor.current_stats
        # Memory normalization (raw bytes → MB + %)
        mem_raw = s.get("memory", {})
        total_b = mem_raw.get("total", 1) or 1
        used_b  = mem_raw.get("used", 0)
        free_b  = mem_raw.get("free", 0)
        # Network: pick best physical interface (eth0 > wlan0), returns {rx_mbps, tx_mbps}
        net_all = s.get("network", {})
        eth = net_all.get("eth0", net_all.get("wlan0", {}))
        iface = "eth0" if "eth0" in net_all else ("wlan0" if "wlan0" in net_all else "—")
        metrics = {
            "timestamp": datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
            "cpu_usage_percent": round(s.get("cpu", {}).get("cpu", 0.0), 1),
            "core_temperature_c": s.get("temp_c", 0.0),
            "memory": {
                "total_mb":    round(total_b / 1024 / 1024, 1),
                "used_mb":     round(used_b  / 1024 / 1024, 1),
                "free_mb":     round(free_b  / 1024 / 1024, 1),
                "used_percent": round(used_b / total_b * 100, 1),
            },
            "network": {
                "interface":    iface,
                "download_mbps": round(eth.get("rx_mbps", 0.0), 2),
                "upload_mbps":   round(eth.get("tx_mbps", 0.0), 2),
            },
            "load_avg": [round(x, 2) for x in s.get("load_avg", [0, 0, 0])],
            "throttled": s.get("throttled", "0x0"),
        }
        return web.json_response(metrics)
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)



# --- Active stress process handle -------------------------------------------
_stress_proc = None
_stress_active = False
_stress_start_time = 0.0


async def handle_pi5_diagnose(request: web.Request) -> web.Response:
    """Full Pi 5 diagnostic: service status, socket, memory, thermal, last logs."""
    import subprocess, shutil
    result = {}
    # systemd service
    try:
        r = subprocess.run(["systemctl", "is-active", "axiom-zero.service"],
                           capture_output=True, text=True, timeout=5)
        result["service_status"] = r.stdout.strip()
    except Exception as e:
        result["service_status"] = f"error: {e}"
    # IPC socket
    result["ipc_socket"] = "OK" if os.path.exists("/tmp/axiom_zero_ipc.sock") else "MISSING"
    # Memory
    try:
        with open("/proc/meminfo") as f:
            lines = {l.split(":")[0]: l.split(":")[1].strip() for l in f}
        total = int(lines.get("MemTotal","0 kB").split()[0])
        avail = int(lines.get("MemAvailable","0 kB").split()[0])
        used_pct = round((1 - avail/total)*100, 1) if total else 0
        result["memory_used_pct"] = used_pct
        result["memory_total_mb"] = round(total/1024, 1)
    except Exception as e:
        result["memory_used_pct"] = f"error: {e}"
    # Thermal
    try:
        r = subprocess.run(["vcgencmd", "measure_temp"], capture_output=True, text=True, timeout=3)
        result["cpu_temp"] = r.stdout.strip()
        r2 = subprocess.run(["vcgencmd", "get_throttled"], capture_output=True, text=True, timeout=3)
        result["throttled"] = r2.stdout.strip()
    except Exception:
        try:
            with open("/sys/class/thermal/thermal_zone0/temp") as f:
                result["cpu_temp"] = f"{int(f.read().strip())/1000:.1f}°C"
        except Exception as e:
            result["cpu_temp"] = f"error: {e}"
    # Last 30 log lines
    try:
        r = subprocess.run(
            ["journalctl", "-u", "axiom-zero.service", "-n", "30", "--no-pager", "--output=short-iso"],
            capture_output=True, text=True, timeout=5
        )
        result["last_logs"] = r.stdout[-3000:]  # cap at 3000 chars
    except Exception as e:
        # fallback: try reading log file
        try:
            log_path = os.path.expanduser("~/api_gateway.log")
            with open(log_path) as f:
                lines = f.readlines()
            result["last_logs"] = "".join(lines[-30:])
        except Exception:
            result["last_logs"] = f"unavailable: {e}"
    return web.json_response(result)


async def handle_pi5_diagnose_fix(request: web.Request) -> web.Response:
    """Automated fix: restart axiom-zero.service and clean stale IPC socket."""
    import subprocess
    actions = []
    # Remove stale socket
    if os.path.exists("/tmp/axiom_zero_ipc.sock"):
        try:
            os.remove("/tmp/axiom_zero_ipc.sock")
            actions.append("removed stale IPC socket")
        except Exception as e:
            actions.append(f"socket remove failed: {e}")
    # Restart service
    try:
        r = subprocess.run(["systemctl", "restart", "axiom-zero.service"],
                           capture_output=True, text=True, timeout=15)
        actions.append(f"service restart: {'ok' if r.returncode == 0 else r.stderr.strip()}")
    except Exception as e:
        actions.append(f"service restart error: {e}")
    return web.json_response({"actions": actions, "ok": True})


async def handle_stress_start(request: web.Request) -> web.Response:
    """Launch stress test as background subprocess on Pi 5."""
    global _stress_proc, _stress_active, _stress_start_time
    import subprocess, sys
    if _stress_active and _stress_proc and _stress_proc.poll() is None:
        return web.json_response({"error": "stress test already running", "pid": _stress_proc.pid}, status=409)
    body = await request.json()
    duration    = int(body.get("duration", 30))
    concurrency = int(body.get("concurrency", 100))
    mode        = body.get("mode", "sustained")
    target_url  = body.get("target_url", "http://127.0.0.1:8085/v1/score")

    script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "live_pi5_stress_test.py")
    cmd = [sys.executable, script,
           f"--url={target_url}",
           f"--concurrency={concurrency}",
           f"--duration={duration}"]
    try:
        _stress_proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        _stress_active = True
        _stress_start_time = time.time()
        return web.json_response({"ok": True, "pid": _stress_proc.pid, "duration": duration, "concurrency": concurrency})
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


async def handle_stress_stop(request: web.Request) -> web.Response:
    global _stress_proc, _stress_active
    import signal
    if _stress_proc:
        try:
            _stress_proc.terminate()
            try:
                _stress_proc.wait(timeout=3)
            except Exception:
                _stress_proc.kill()
        except Exception:
            pass
    _stress_active = False
    _stress_proc = None
    return web.json_response({"ok": True, "stopped": True})


async def handle_stress_status(request: web.Request) -> web.Response:
    global _stress_proc, _stress_active, _stress_start_time
    if _stress_proc and _stress_proc.poll() is not None:
        _stress_active = False
    elapsed = round(time.time() - _stress_start_time, 1) if _stress_active else 0
    return web.json_response({
        "active": _stress_active,
        "pid": _stress_proc.pid if _stress_proc and _stress_active else None,
        "elapsed": elapsed,
    })


def create_app() -> web.Application:
    app = web.Application()
    app.router.add_post("/v1/score", handle_score)
    app.router.add_post("/api/telemetry/submit", handle_score)
    app.router.add_get("/.well-known/axiom-pubkey.jwk", handle_jwk)
    # Pi 5 Telemetry & Diagnostic API
    app.router.add_get("/api/pi5/telemetry",         handle_pi5_telemetry)
    app.router.add_get("/api/pi5/diagnose",           handle_pi5_diagnose)
    app.router.add_post("/api/pi5/diagnose/fix",      handle_pi5_diagnose_fix)
    app.router.add_post("/api/pi5/stress/start",      handle_stress_start)
    app.router.add_post("/api/pi5/stress/stop",       handle_stress_stop)
    app.router.add_get("/api/pi5/stress/status",      handle_stress_status)
    # Serve static UI dashboard files from dev_control_center
    static_dir = os.path.dirname(os.path.abspath(__file__))
    app.router.add_static("/", static_dir, show_index=True)
    return app


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Axiom Zero API Gateway")
    parser.add_argument("--host", default="0.0.0.0", help="Host address to bind")
    parser.add_argument("--port", type=int, default=8080, help="Port to bind")
    args = parser.parse_args()

    app = create_app()
    logger.info(f"Starting Axiom Zero API Gateway on {args.host}:{args.port}...")
    web.run_app(app, host=args.host, port=args.port)