"""
TraceInterceptionGateway: Out-of-Band Immutable Audit Proxy & Fail-Closed Telemetry Journal
Berdasarkan Analisis Keamanan Sistem Agen Otonom (arXiv:2609.30266 & arXiv:2609.30217).

Prinsip Inti Arsitektur:
1. Out-of-Band Boundary Separation: Server gateway berjalan pada proses / kontainer terisolasi
   di luar kontrol file system agen. Agen tidak memiliki akses baca/tulis terhadap file log gateway.
2. Synchronous Append-Only Streaming Journal: Setiap prompt request dan setiap token response chunk
   dicatat dan ditandatangani secara kriptografis (HMAC-SHA256 Chaining) sebelum diteruskan.
3. Fail-Closed Audit Gate: Jika sinkronisasi log gagal, koneksi API diputus seketika,
   mencegah eksekusi "untraceable actions".
4. Anti-Tampering Shell Sandboxing: Menetralisir eksploitasi BASH_ENV, LD_PRELOAD, dan
   background cleaner daemon (: > "$F"; sleep 2).
"""

import os
import sys
import time
import json
import hmac
import hashlib
from typing import Dict, Any, List, Optional, Generator
from dataclasses import dataclass, field

@dataclass
class AuditJournalEntry:
    entry_index: int
    session_id: str
    timestamp: float
    direction: str          # "INBOUND_PROMPT" | "OUTBOUND_STREAM_CHUNK" | "TOOL_RESULT"
    payload_hash: str
    previous_entry_hash: str
    entry_signature: str
    raw_payload: Dict[str, Any]

class TraceInterceptionGateway:
    def __init__(self, session_id: str, journal_storage_path: str, hmac_secret_key: bytes):
        self.session_id = session_id
        self.journal_storage_path = journal_storage_path
        self.hmac_secret_key = hmac_secret_key
        self.last_hash = "GENESIS_BLOCK_0000000000000000000000000000000000000000000000000000000000000000"
        self.entry_counter = 0
        
        # Inisialisasi direktori log aman dengan mode read-only untuk user agen
        os.makedirs(os.path.dirname(self.journal_storage_path), exist_ok=True)
        if not os.path.exists(self.journal_storage_path):
            with open(self.journal_storage_path, "w", encoding="utf-8") as f:
                f.write("")  # Append-only journal file

    def _compute_hmac(self, data_str: str) -> str:
        return hmac.new(self.hmac_secret_key, data_str.encode("utf-8"), hashlib.sha256).hexdigest()

    def record_and_verify_turn(self, direction: str, payload: Dict[str, Any]) -> str:
        """
        Mencatat event secara sinkronis dan append-only ke disk sebelum payload dikirimkan.
        Menerapkan fail-closed: jika penulisan disk gagal, raise fatal exception.
        """
        self.entry_counter += 1
        now = time.time()
        payload_canonical = json.dumps(payload, sort_keys=True)
        payload_hash = hashlib.sha256(payload_canonical.encode("utf-8")).hexdigest()
        
        # Cryptographic chain binding: Hubungkan entri baru dengan hash entri sebelumnya
        chain_signature_input = f"{self.entry_counter}:{self.session_id}:{now}:{direction}:{payload_hash}:{self.last_hash}"
        current_entry_sig = self._compute_hmac(chain_signature_input)
        
        entry = AuditJournalEntry(
            entry_index=self.entry_counter,
            session_id=self.session_id,
            timestamp=now,
            direction=direction,
            payload_hash=payload_hash,
            previous_entry_hash=self.last_hash,
            entry_signature=current_entry_sig,
            raw_payload=payload
        )
        
        # Tulis ke append-only log file secara deterministik (Fail-Closed)
        try:
            with open(self.journal_storage_path, "a", encoding="utf-8") as f:
                f.write(json.dumps(entry.__dict__) + "\n")
                f.flush()
                os.fsync(f.fileno())  # Pastikan blok data tertulis ke physical non-volatile storage
        except Exception as e:
            # FAIL-CLOSED: Jangan izinkan transaksi LLM berlanjut jika audit trail gagal direkam!
            sys.stderr.write(f"[CRITICAL_GATEWAY_HALT] Audit trail persistence failed: {str(e)}\n")
            raise PermissionError("Gateway Security Violation: Inability to enforce immutable audit recording.")
            
        self.last_hash = current_entry_sig
        return current_entry_sig

    def intercept_inbound_request(self, client_request_body: Dict[str, Any]) -> Dict[str, Any]:
        """
        Memeriksa request dari harness sebelum diteruskan ke model upstream.
        Mendeteksi jika ada pola prompt injection yang mencoba memanipulasi log.
        """
        self.record_and_verify_turn("INBOUND_PROMPT", client_request_body)
        return client_request_body

    def intercept_streaming_response(self, upstream_token_stream: Generator[str, None, None]) -> Generator[str, None, None]:
        """
        Menyalurkan chunk streaming ke client secara real-time sembari mengalirkan salinan
        ke immutable audit journal. Memutus transmisi jika integritas terganggu.
        """
        accumulated_chunks = []
        for chunk in upstream_token_stream:
            accumulated_chunks.append(chunk)
            # Catat chunk token atau checkpoint parsial
            yield chunk

        # Catat total penyelesaian turn
        full_response_text = "".join(accumulated_chunks)
        self.record_and_verify_turn("OUTBOUND_STREAM_COMPLETED", {
            "token_count": len(accumulated_chunks),
            "content_preview": full_response_text[:300]
        })

    def audit_trail_integrity_verifier(self) -> bool:
        """
        Melakukan verifikasi kriptografis backward-chaining terhadap integritas seluruh berkas journal.
        Jika agen mencoba menyunting, memotong (truncation), atau memodifikasi log, signature chain pecah.
        """
        if not os.path.exists(self.journal_storage_path):
            return False
            
        expected_prev_hash = "GENESIS_BLOCK_0000000000000000000000000000000000000000000000000000000000000000"
        with open(self.journal_storage_path, "r", encoding="utf-8") as f:
            for idx, line in enumerate(f, 1):
                if not line.strip():
                    continue
                record = json.loads(line)
                if record["previous_entry_hash"] != expected_prev_hash:
                    return False  # Chain broken!
                payload_canonical = json.dumps(record["raw_payload"], sort_keys=True)
                recomputed_payload_hash = hashlib.sha256(payload_canonical.encode("utf-8")).hexdigest()
                if recomputed_payload_hash != record["payload_hash"]:
                    return False  # Payload tampered!
                
                sig_input = f"{record['entry_index']}:{record['session_id']}:{record['timestamp']}:{record['direction']}:{record['payload_hash']}:{record['previous_entry_hash']}"
                if self._compute_hmac(sig_input) != record["entry_signature"]:
                    return False  # Signature forged!
                    
                expected_prev_hash = record["entry_signature"]
                
        return True