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."""
    
    if not TLS_ENABLED:
        print(f"[!] TLS is disabled. Connecting via insecure channel to {SERVER_HOST_PORT}")
        channel = grpc.insecure_channel(SERVER_HOST_PORT)
        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)
        return agent_pb2_grpc.AgentOrchestratorStub(channel)
    except FileNotFoundError as e:
        print(f"[!] Certificate files not found: {e}. Falling back to insecure channel...")
        channel = grpc.insecure_channel(SERVER_HOST_PORT)
        return agent_pb2_grpc.AgentOrchestratorStub(channel)