import grpc
import os
from protos import agent_pb2_grpc
from agent_node.config import SERVER_HOST_PORT, TLS_ENABLED, CERT_CA, CERT_CLIENT_CRT, CERT_CLIENT_KEY

def get_secure_stub():
    """Initializes a gRPC channel (Secure or Insecure) and returns the orchestrator stub."""
    
    options = [
        ('grpc.keepalive_time_ms', 30000),      # Send keepalive ping every 30s
        ('grpc.keepalive_timeout_ms', 10000),    # Wait 10s for pong
        ('grpc.keepalive_permit_without_calls', True),
        ('grpc.http2.max_pings_without_data', 0) # Allow infinite pings
    ]

    if not TLS_ENABLED:
        print(f"[!] TLS is disabled. Connecting via insecure channel to {SERVER_HOST_PORT}")
        channel = grpc.insecure_channel(SERVER_HOST_PORT, options=options)
        return agent_pb2_grpc.AgentOrchestratorStub(channel)

    print(f"[*] Connecting via secure (mTLS) channel to {SERVER_HOST_PORT}")
    try:
        with open(CERT_CLIENT_KEY, 'rb') as f: pkey = f.read()
        with open(CERT_CLIENT_CRT, 'rb') as f: cert = f.read()
        with open(CERT_CA, 'rb') as f: ca = f.read()
        
        creds = grpc.ssl_channel_credentials(ca, pkey, cert)
        channel = grpc.secure_channel(SERVER_HOST_PORT, creds, options=options)
        return agent_pb2_grpc.AgentOrchestratorStub(channel)
    except FileNotFoundError as e:
        print(f"[!] mTLS Certificate files not found: {e}. Falling back to standard TLS (Server Verify)...")
        # Fallback to standard TLS (uses system CA roots by default)
        creds = grpc.ssl_channel_credentials()
        channel = grpc.secure_channel(SERVER_HOST_PORT, creds, options=options)
        return agent_pb2_grpc.AgentOrchestratorStub(channel)